💬 NLP

Nemotron-TwoTower: Diffusion Language Modeling with Pretrained Autoregressive Context

Némotron-TwoTower introduit une nouvelle architecture de diffusion autorégressive par blocs qui découple le traitement du contexte et le débruitage en deux tours spécialisés, atteignant un débit de génération 2,42 fois plus rapide tout en conservant 98,7 % de la qualité de sa base de référence autorégressive.

Auteurs originaux : Fitsum Reda, John Kamalu, Roger Waleffe, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro

Publié 2026-06-26

📖 6 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Fitsum Reda, John Kamalu, Roger Waleffe, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez d'écrire une longue histoire, mais que vous avez une règle stricte : vous ne pouvez écrire qu'un seul mot à la fois, et vous devez attendre que le mot précédent soit entièrement terminé avant de pouvoir commencer le suivant. C'est ainsi que fonctionnent la plupart des modèles de langage IA actuels (appelés modèles « autoregressifs »). C'est comme un dactylo très méticuleux, lent, qui ne fait jamais d'erreur mais met un temps infini à terminer une phrase.

Maintenant, imaginez une nouvelle méthode où vous pouvez taper une phrase entière d'un coup, l'examiner, puis corriger les erreurs dans tous les mots simultanément. C'est plus rapide, mais généralement, la qualité de l'histoire en pâtit car l'IA s'embrouille en essayant de tout corriger en même temps.

Nemotron-TwoTower est une nouvelle invention qui tente d'obtenir le meilleur des deux mondes : la vitesse de la dactylographie de tout d'un coup, avec la haute qualité du dactylo méticuleux qui écrit un mot à la fois.

Voici comment cela fonctionne, en utilisant une analogie simple :

Les Deux Tours : Le Bibliothécaire et L'Éditeur

Les chercheurs ont construit un système avec deux « tours » distinctes (parties de l'IA) qui travaillent ensemble, comme une équipe de deux personnes :

La Tour de Contexte Gelée (Le Bibliothécaire) :
- Rôle : C'est l'« expert » qui connaît l'histoire jusqu'ici. C'est une IA pré-entraînée qui est gelée (verrouillée en place). Elle agit comme un bibliothécaire strict qui lit l'histoire mot par mot, du début à la fin.
- Travail : Elle ne change pas. Elle tient simplement un registre parfait et organisé de tout ce qui a été écrit jusqu'à présent. Elle dit à l'autre tour : « Voici le contexte de l'histoire jusqu'à ce point. »
- Pourquoi c'est important : Parce qu'elle est gelée, elle conserve toute l'intelligence et les connaissances du modèle massif de 30 milliards de paramètres dont elle est issue. Elle n'oublie jamais et ne se laisse pas confondre.
La Tour de Débruiteur de Diffusion Entraînable (L'Éditeur) :
- Rôle : C'est le « travailleur rapide ». C'est une IA entraînable qui est autorisée à changer et à apprendre.
- Travail : Au lieu d'écrire un mot à la fois, elle prend un bloc de 16 mots (ou tokens) qui sont actuellement vides (masqués) et essaie de remplir ces 16 blancs d'un seul coup. Elle regarde les notes du « Bibliothécaire » et tente de remplir les 16 blancs.
- Le processus : Elle ne réussit pas parfaitement du premier coup. Elle fait une supposition, puis examine sa propre supposition et les notes du Bibliothécaire, et affine les mots. Elle répète ce processus (itérativement) jusqu'à ce que le bloc de mots soit propre et cohérent.
- La magie : Comme elle peut regarder les 16 mots en même temps et les corriger ensemble, elle est beaucoup plus rapide que l'écriture un par un.

Comment ils travaillent ensemble (La chaîne de montage)

Imaginez une chaîne de montage d'usine :

Le Bibliothécaire lit l'invite (prompt) et les mots déjà écrits, gardant une carte mentale parfaite de l'histoire.
L'Éditeur saisit un bloc de papier vierge (long de 16 mots).
L'Éditeur demande au Bibliothécaire : « Que vient-il ensuite ? » et le Bibliothécaire fournit le contexte.
L'Éditeur remplit les 16 blancs, vérifie son travail et corrige les erreurs.
Une fois que l'Éditeur est convaincu que les 16 mots sont bons, ils sont « engagés » (finalisés).
Le Bibliothécaire met à jour sa carte mentale pour inclure ces nouveaux 16 mots.
L'Éditeur saisit le bloc suivant de 16 mots vierges, et le cycle se répète.

Les Résultats : Vitesse vs Qualité

L'article affirme que ce système atteint un « point d'équilibre » :

Vitesse : Il est 2,42 fois plus rapide que la méthode traditionnelle d'un mot à la fois. C'est comme passer d'un escargot à un sprinter.
Qualité : Il conserve 98,7 % de la qualité du modèle original. Il n'a pas perdu beaucoup d'intelligence pour gagner cette vitesse.
Efficacité : Ils ont entraîné cette nouvelle tour « Éditeur » sur environ 2,1 billions de mots, ce qui est une infime fraction des 25 billions de mots utilisés pour entraîner le « Bibliothécaire » original. Cela signifie qu'ils n'ont pas eu besoin de réenseigner toute l'IA à partir de zéro ; ils ont juste appris à l'Éditeur comment utiliser les notes du Bibliothécaire.

Pourquoi est-ce différent des tentatives précédentes ?

Les tentatives précédentes essayaient d'utiliser un seul cerveau pour faire les deux tâches : mémoriser l'histoire et corriger les mots en même temps. L'article soutient que c'est comme demander à une personne d'être à la fois un gardien de mémoire parfait et un éditeur rapide — c'est trop de travail, et elle finit par ne bien faire aucun des deux.

En divisant les rôles en deux tours (une gelée et une entraînable), ils permettent à chaque partie de faire ce pour quoi elle est la meilleure. Le Bibliothécaire reste parfait, et l'Éditeur devient rapide.

L'astuce de la « Confiance »

Le système possède également une fonction intelligente de « démasquage de confiance » (confidence unmasking).

Si l'Éditeur est très confiant concernant un mot, il le verrouille immédiatement.
S'il est incertain, il laisse ce mot vide (masqué) et essaie de le corriger lors du tour d'édition suivant.
Cela signifie qu'il ne perd pas de temps à revérifier les mots qu'il sait déjà être corrects, mais qu'il consacre du temps supplémentaire pour polir les parties délicates.

Résumé

Nemotron-TwoTower est un moyen de rendre la génération de texte par l'IA beaucoup plus rapide sans la rendre « plus stupide ». Il y parvient en séparant la tâche de « se souvenir de l'histoire » (qui reste lente et parfaite) de celle de « rédiger le bloc suivant » (qui devient rapide et itérative). Le résultat est un système qui écrit des histoires presque aussi bien que l'ancienne méthode lente, mais plus de deux fois plus vite.

Résumé Technique : Nemotron-TwoTower

Énoncé du Problème

Les modèles de langage autorégressifs (AR) dominent la génération de texte mais souffrent d'un goulot d'étranglement de débit dû à leur processus de décodage séquentiel, un jeton à la fois. Les modèles de langage de diffusion discrète offrent une alternative prometteuse en générant des jetons en parallèle et en les affinant de manière itérative, ce qui peut potentiellement offrir un débit plus élevé et un contrôle plus fin. Cependant, les approches de diffusion existantes reposent généralement sur un réseau unique pour remplir simultanément deux rôles distincts : représenter les jetons de contexte propres et débruiter itérativement les jetons corrompus. Cet enchevêtrement force un ensemble unique de poids à servir des objectifs conflictuels, limitant la capacité du modèle à exceller dans l'un ou l'autre de ces rôles. De plus, il reste incertain si un découplage complet de ces rôles est réalisable à grande échelle, particulièrement sur les architectures hybrides modernes combinant Mamba, l'attention et des couches de mélange d'experts (MoE).

Méthodologie

Les auteurs proposent TwoTower, un modèle de diffusion autorégressif par blocs qui sépare les rôles de représentation du contexte et de débruitage en deux tours distincts, construits sur l'architecture hybride Nemotron-3-Nano-30B-A3B (qui entrelace des couches Mamba-2, d'auto-attention et MoE).

Architecture

Le modèle utilise deux copies du réseau préentraîné avec des rôles complémentaires :

Tour de Contexte AR Gelée (Frozen AR Context Tower) : Cette tour reste gelée et agit comme un modèle autorégressif causal sur le prompt et les jetons précédemment engagés. Elle traite les jetons propres pour produire des caches Clé-Valeur (KV) alignés sur les couches et des états Mamba. Elle préserve les capacités autorégressives originales du backbone et maintient des états réutilisables à travers les blocs.
Tour de Débruitage de Diffusion Entraînable (Trainable Diffusion Denoiser Tower) : Cette tour est entraînée pour affiner les blocs bruités. Elle emploie une attention de bloc bidirectionnelle au sein du bloc bruité actuel (permettant aux jetons de s'attendre mutuellement) tout en maintenant une attention croisée causale vers les couches correspondantes de la tour de contexte gelée. Cette attention croisée alignée sur les couches fournit au débruiteur un accès multi-échelle aux représentations du backbone, contrastant avec les méthodes précédentes qui ne diffusaient que l'état caché final.

Processus de Génération

La génération procède bloc par bloc (Algorithme 1) :

La tour de contexte traite le prompt et les jetons engagés de manière causale.
Un nouveau bloc est initialisé avec $S$ jetons masqués.
La tour de débruitage affine itérativement ce bloc sur $T$ étapes. À chaque couche, le débruiteur est attentif au cache KV de la tour de contexte pour les blocs passés et à son propre cache KV bidirectionnel pour le bloc actuel.
Une fois qu'un bloc est suffisamment affiné (basé sur des seuils de confiance), il est engagé. La tour de contexte met ensuite à jour ses états avec les nouveaux jetons propres, et le processus se répète pour le bloc suivant.

Entraînement et Conditionnement

Diffusion Masquée : Le débruiteur est entraîné à l'aide d'un objectif de diffusion masquée. Les blocs propres sont corrompus en remplaçant les jetons par [MASK] selon un calendrier linéaire. Le modèle minimise la log-vraisemblance négative des jetons originaux aux positions masquées.
Conditionnement Temporel : Le débruiteur est conditionné sur l'étape de diffusion $t$ via la normalisation adaptative de couche (adaLN), une technique courante dans la diffusion d'images mais moins utilisée dans les modèles de langage masqués. Cela ajoute un nombre minimal de paramètres (~1,5M) mais améliore la qualité de l'échantillonnage.
Curriculum : L'entraînement suit un curriculum en deux étapes. La Phase 1 utilise un mélange de données diversifiées pour une couverture large, tandis que la Phase 2 surpondère les sources de haute qualité et axées sur les STEM. La taille du bloc d'entraînement est progressivement réduite (de $S=32$ à $S=16$ ) pour équilibrer le parallélisme et la qualité du conditionnement.
Échantillonnage : Les auteurs emploient une stratégie de démasquage par confiance (Algorithme 2). À chaque étape, le débruiteur prédit tous les jetons masqués en parallèle ; les jetons dont la confiance est supérieure à un seuil $\gamma$ sont engagés, tandis que les jetons incertains restent masqués pour un affinement ultérieur. Cela rend le nombre de jetons engagés adaptatif.

Résultats Clés

Évalué sur le backbone Nemotron-3-Nano-30B-A3B (30 milliards de paramètres, MoE hybride), Nemotron-TwoTower démontre les points suivants :

Compromis Qualité-Débit : Le modèle conserve 98,7 % de la qualité globale des benchmarks de la ligne de base autorégressive tout en atteignant une amélioration de 2,42× du débit de génération en temps réel (wall-clock).
Performance des Benchmarks :
- Connaissances Générales : La performance reste à environ 1 point de la ligne de base AR.
- Code et Mathématiques : Montre une légère dégradation par rapport à la ligne de base.
- Sens Commun et Multilinguisme : La performance est récupérée ou légèrement améliorée.
Ablations de Conception :
- Mamba Bidirectionnel : Tester une variante Mamba bidirectionnelle a montré des gains de qualité marginaux mais a doublé le calcul SSM ; la conception finale conserve un Mamba causal.
- Conditionnement Temporel : Le conditionnement temporel explicite a considérablement amélioré les scores de génération, de code et de mathématiques.
- Découplage des Tours : Garder la tour de contexte gelée et entraîner une tour de débruitage séparée a donné les meilleurs résultats. Lier les poids entre les tours ou continuer l'entraînement AR sur la tour de contexte a entraîné des chutes de précision significatives.
Sensibilité de la Taille de Bloc : Des tailles de blocs d'entraînement plus petites ( $S=16$ ) ont amélioré la qualité par rapport à des tailles plus grandes ( $S=128$ ), mais ont réduit le débit. Les auteurs ont choisi $S=16$ par défaut pour équilibrer la rétention de qualité et la vitesse. L'échantillonnage avec des blocs plus grands que le bloc d'entraînement a entraîné une dégradation sévère des performances.

Signification et Revendications

L'article affirme que Nemotron-TwoTower démontre avec succès que les représentations autorégressives préentraînées peuvent supporter un débruiteur de diffusion entraînable sans ré-entraîner l'intégralité du modèle. En découplant les rôles de contexte et de débruitage, l'approche préserve les larges capacités du modèle hybride MoE de 30B original tout en débloquant les avantages de la génération parallèle de la diffusion.

Les auteurs soulignent que cette récupération de qualité est obtenue grâce à une adaptation sur une fraction des données utilisées pour le préentraînement original (~2,1 billions de jetons contre 25 billions de jetons). Ce travail établit une nouvelle frontière de Pareto pour les modèles de langage, offrant une voie viable vers une génération à haut débit sans sacrifier la qualité fondamentale des modèles autorégressifs à grande échelle. Le code et les poids du modèle sont publiés pour faciliter la recherche ultérieure dans cette direction.