Orthrus: Memory-Efficient Parallel Token Generation via Dual-View Diffusion

Orthrus est un cadre à double architecture économe en mémoire qui intègre un module de diffusion léger avec un LLM autorégressif figé pour permettre la génération parallèle de tokens tout en garantissant une fidélité d'inférence sans perte grâce à un cache KV partagé et un mécanisme de consensus exact.

L'essentiel

Imaginez que vous essayiez d'écrire une histoire longue et complexe. Vous avez deux façons de le faire, mais les deux présentent un défaut majeur :

1. L'écrivain « un mot à la fois » (modèles autoregressifs) : Cet écrivain est incroyablement intelligent et précis. Il réfléchit soigneusement à chaque mot avant de l'écrire, s'assurant que l'histoire a un sens parfait. Cependant, il est lent. Il doit terminer un mot, consulter ses notes, réfléchir au suivant, puis l'écrire. Il ne peut pas accélérer car il craint de faire une erreur.
2. L'écrivain « par lots » (modèles de diffusion) : Cet écrivain tente d'écrire un paragraphe entier d'un coup. Il est très rapide ! Mais parce qu'il devine plusieurs mots simultanément sans vérifier chacun d'eux soigneusement, il fait souvent des erreurs logiques, perd le fil de l'intrigue ou écrit des absurdités.

Orthrus est un nouveau cadre qui combine le meilleur des deux mondes. Il crée un système à « double voix » qui vous permet d'écrire un paragraphe entier d'un coup sans perdre la précision de l'écrivain soigneux.

Voici comment cela fonctionne, en utilisant une analogie simple :

L'analogie de « l'Architecte et le Bâtisseur »

Imaginez le modèle d'IA comme un chantier de construction avec deux ouvriers : l'Architecte et le Bâtisseur.

• L'Architecte (le LLM figé) : C'est le modèle original, hautement entraîné et super-intelligent. C'est l'expert qui sait exactement à quoi le bâtiment devrait ressembler. Il est « figé », ce qui signifie qu'il ne change pas d'avis ni n'apprend de nouvelles choses pendant ce processus ; il fournit simplement le plan parfait.
• Le Bâtisseur (le module de diffusion) : C'est un nouvel ouvrier, léger, ajouté à l'équipe. Sa tâche est de poser des briques (des jetons) rapidement.

Comment ils travaillent ensemble :

1. Mise en place (Pré-remplissage) : D'abord, l'Architecte lit l'ensemble de l'invite (les instructions) et construit une « carte mémoire » parfaite et haute fidélité (appelée cache KV). Cette carte contient tout le contexte nécessaire pour construire le reste de l'histoire.
2. Le sprint parallèle (Génération) : Au lieu que l'Architecte pose une brique à la fois, le Bâtisseur regarde la carte de l'Architecte et tente de poser toute une rangée de briques (disons 32 briques) d'un coup.
3. Le contrôle de sécurité (Consensus) : C'est la partie magique. Avant que le travail du Bâtisseur ne soit accepté, l'Architecte vérifie instantanément le lot du Bâtisseur.
   • Si le Bâtisseur a deviné le mot suivant correctement selon la logique parfaite de l'Architecte, l'Architecte dit : « Super ! Gardez-le ! »
   • Si le Bâtisseur a deviné faux, l'Architecte dit : « Non, ce n'est pas ça », et corrige ce mot spécifique immédiatement.
   • Le processus se répète pour le lot suivant.

Pourquoi est-ce une grande avancée ?

• Pas de gaspillage de mémoire : Habituellement, si vous avez deux modèles qui travaillent, vous avez besoin de deux jeux de notes en mémoire. Orthrus est malin car le Bâtisseur et l'Architecte partagent la même carte mémoire exacte. Le Bâtisseur n'a pas besoin de faire ses propres notes ; il regarde simplement celles de l'Architecte. Cela économise une énorme quantité de mémoire informatique.
• Pas de perte de qualité : Parce que l'Architecte (le modèle intelligent original) a le dernier mot sur chaque mot, l'histoire est tout aussi bonne que si l'Architecte l'avait écrite mot par mot. Il n'y a pas de « dérive » ni de perte de qualité.
• Vitesse massive : En permettant au Bâtisseur de poser 32 briques à la fois et en ne les vérifiant que de manière instantanée, Orthrus est jusqu'à 7,8 fois plus rapide que la méthode lente, un mot à la fois.

Les résultats

L'article a testé cela sur des tâches difficiles comme la résolution de problèmes mathématiques (MATH-500), l'écriture de code et la réponse à des énigmes logiques.

• Vitesse : Il était nettement plus rapide que les modèles standards.
• Précision : Il était tout aussi précis que le modèle lent original.
• Efficacité : Il n'a nécessité l'entraînement que d'une toute petite fraction (environ 16 %) des paramètres du modèle, le rendant peu coûteux et facile à ajouter aux systèmes d'IA existants.

En bref, Orthrus est comme embaucher un lecteur rapide capable de deviner les 30 prochains mots d'une histoire instantanément, mais qui a un éditeur strict debout juste à côté de lui pour corriger toute erreur immédiatement. Le résultat est une histoire écrite à la vitesse de l'éclair qui reste parfaitement précise.

Résumé technique

Résumé Technique : Orthrus – Génération Parallèle de Tokens Économe en Mémoire via une Diffusion à Double Vue

1. Énoncé du Problème

Les modèles de langage de grande taille (LLM) autorégressifs (AR) dominent actuellement le traitement du langage naturel grâce à leur génération haute fidélité et à leurs capacités de raisonnement robustes. Cependant, ils souffrent d'une inefficacité fondamentale lors de la phase de décodage : la génération de tokens est strictement séquentielle. Alors que l'étape de pré-remplissage traite les invites en parallèle, la phase de génération nécessite $N$ passes avant distinctes pour produire $N$ tokens. Cette dépendance séquentielle crée un goulot d'étranglement de bande passante mémoire, entraînant une sous-utilisation du matériel et une latence d'inférence élevée.

À l'inverse, les modèles de langage par diffusion (DLM) offrent une génération native parallèle en débruitant des blocs de tokens simultanément. Cependant, les DLM existants font face à des obstacles majeurs :

• Dégradation des performances : Ils sous-performent souvent les modèles AR de taille similaire, en particulier dans les tâches de raisonnement complexe, en raison d'une « dérive conditionnelle » où l'hypothèse d'indépendance conditionnelle viole les dépendances causales strictes.
• Coûts d'entraînement : Atteindre une cohérence de base nécessite souvent des ensembles de données d'entraînement massifs (par exemple, des centaines de milliards de tokens) ou un pré-entraînement continu.
• Divergence architecturale : L'adaptation de modèles AR pré-entraînés vers des cadres de diffusion modifie souvent les poids de base, détruisant la distribution prédictive exacte du modèle original et échouant à égaler ses capacités de raisonnement.

Le défi central consiste à unifier la conditionnement causal haute fidélité des modèles AR avec la vitesse de décodage parallèle des modèles de diffusion sans sacrifier l'un ni l'autre.

2. Méthodologie : L'Architecture Orthrus

Orthrus propose un cadre à double architecture qui unifie ces paradigmes au sein d'un seul Transformer. Au lieu de remplacer le socle AR, Orthrus enrichit un modèle AR pré-entraîné et figé avec un module de diffusion léger et entraînable.

2.1 Attention Unifiée à Double Vue

L'architecture introduit deux chemins d'attention distincts opérant sur un cache Key-Value (KV) partagé :

1. La Tête AR Figée (Chemin Bleu) : Ce chemin reste strictement figé. Sa seule fonction est de traiter le contexte lors de l'étape de pré-remplissage pour construire des représentations KV causales haute fidélité ($K_{AR}, V_{AR}$). Il agit comme le « professeur » pour la distribution prédictive exacte.
2. La Tête de Diffusion Entraînable (Chemin Rouge) : Un module léger (initialisé à partir de contreparties AR) est injecté à côté des têtes d'attention AR. Il est conçu spécifiquement pour une génération parallèle haute vitesse.

2.2 Entraînement : Masquage de Blocs en Double Passage

L'entraînement vise à aligner les prédictions parallèles de la vue de diffusion avec la distribution cible exacte du modèle AR figé.

• Construction des Données : Pour une séquence, des blocs aléatoires de longueur $K$ sont sélectionnés. Le premier token du bloc est conservé comme « ancre » visible, tandis que les $K-1$ tokens suivants sont remplacés par des tokens <mask>.
• Mécanisme d'Attention : La tête de diffusion traite ces blocs corrompus en utilisant un masque de bloc spécialisé ($M_{diff}$). Ce masque impose deux règles :
  1. Contexte Causal : Les positions du bloc s'attendent de manière causale au contexte AR propre précédant l'ancre du bloc.
  2. Bloc Bidirectionnel : Les positions au sein du même bloc masqué s'attendent bidirectionnellement les unes aux autres, permettant une agrégation de contexte parallèle.
• Objectif : La tête de diffusion minimise la divergence KL avant par rapport à la distribution prédictive complète de la tête AR figée. Les gradients ne circulent que via le module de diffusion, laissant le socle AR intact.

2.3 Inférence : Correspondance Exacte de Distribution via le Consensus Intra-Modèle

Orthrus réalise une génération parallèle sans dérive distributionnelle grâce à un mécanisme de consensus :

1. Projection Parallèle : La tête de diffusion prend le token ancre actuel et les $K-1$ masques, les traitant en une seule passe avant pour projeter simultanément $K$ tokens candidats.
2. Validation Structurelle : Le bloc projeté est immédiatement acheminé vers la tête AR figée. Parce que la tête AR voit le bloc entièrement rempli, elle calcule les probabilités cibles exactes pour les $K$ positions en une seule passe.
3. Consensus et Engagement : L'architecture effectue une évaluation stricte de gauche à droite. Un token projeté est accepté si et seulement s'il correspond à la prédiction gloutonne de la tête AR figée. Si une divergence se produit à l'indice $j$, le système engage le préfixe synchronisé jusqu'à $j-1$, ajoute le token de correction AR exact, et tronque le cache. Cela garantit une inférence sans perte, assurant que la sortie correspond strictement à la distribution prédictive du modèle de base.

3. Contributions Clés

• Cadre Nouvelle Architecture Double : Orthrus intègre un module de diffusion parallèle au sein d'un Transformer AR standard, permettant aux deux vues d'opérer sur un cache KV partagé avec zéro stockage redondant de cache KV historique.
• Garantie d'Inférence Sans Perte : En employant un mécanisme de consensus intra-modèle, Orthrus préserve la distribution prédictive exacte du LLM de base, assurant une génération strictement sans perte qui surpasse les adaptations de diffusion antérieures.
• Accélération d'Inférence Significative : En exploitant nativement la tête de diffusion pour la génération parallèle de tokens, Orthrus brise le goulot d'étranglement séquentiel, offrant jusqu'à un accélération de 7,8×.
• Efficacité Extrême en Paramètres et Mémoire : L'intégration est légère. Les capacités parallèles sont injectées en affinant uniquement ~16% des paramètres totaux du modèle en utilisant moins de 1 milliard de tokens (nécessitant moins de 24 heures sur un seul nœud 8xH200).

4. Résultats Expérimentaux

Les auteurs ont évalué Orthrus sur la famille de modèles Qwen3 (1,7B, 4B et 8B paramètres) à travers des benchmarks de raisonnement mathématique (GSM8K, MATH-500, AIME) et de génération de code (HumanEval, MBPP).

• Efficacité : Orthrus a atteint une moyenne de Tokens par Passe Avant (TPF) de 5,39 sur le modèle 8B, se traduisant par des accélérations allant de 3,07× à 7,83× selon la tâche et les paramètres de température.
• Précision : Contrairement aux méthodes d'adaptation qui subissent des baisses de performance, Orthrus a atteint la précision zero-shot exacte du modèle de base Qwen3-8B. Par exemple, sur MATH-500, Orthrus a atteint 86,2 % de précision, alors que les adaptations de diffusion de pointe comme Fast-dLLM-v2 ont subi une chute de 11,1 points (75,1 % contre 86,2 % pour la référence).
• Comparaison avec le Décodage Spéculatif : Comparé aux méthodes externes de décodage spéculatif (EAGLE-3, DFlash), Orthrus a atteint une Longueur Moyenne d'Acceptation significativement plus élevée (11,7 sur MATH-500 contre 7,9 pour DFlash et 3,5 pour EAGLE-3) car il ne nécessite pas le maintien de caches KV séparés et redondants pour un modèle brouillon.

5. Importance et Revendications

L'article revendique qu'Orthrus réconcilie fondamentalement le compromis entre la fidélité de génération autorégressive et le parallélisme basé sur la diffusion.

• Unification Structurelle : En découplant la génération parallèle des contraintes séquentielles tout en l'ancrant dans des représentations AR figées et haute fidélité, Orthrus élimine la « dérive distributionnelle » qui affecte les autres approches de diffusion.
• Évolutivité et Plug-and-Play : Le cadre est présenté comme une solution hautement évolutive qui peut être adaptée de manière transparente à n'importe quel modèle AR open-source existant de haute qualité pour débloquer un débit parallèle sans sacrifier les capacités de raisonnement d'élite.
• Viabilité de Production : Avec une surcharge de cache mémoire de $O(1)$ et des ajouts de paramètres minimes, Orthrus offre une voie pratique et économe en mémoire au déploiement de LLM à haut débit, évitant les coûts de calcul liés à la rééducation de modèles de diffusion massifs à partir de zéro.

Les auteurs concluent qu'Orthrus offre une accélération d'inférence strictement sans perte, proposant un nouvel état de l'art pour la fidélité de la génération parallèle.