NVIDIA résout le compromis majeur des LLM modernes

Lorsqu'on travaille avec n'importe quel grand modèle de langage (LLM), il faut choisir entre la vitesse et la qualité. Les modèles autoregressifs comme GPT offrent une haute qualité mais fonctionnent lentement car ils génèrent les tokens un par un.

Les modèles à diffusion fonctionnent à l'inverse. Ils génèrent des blocs entiers en parallèle, ce qui les rend plus rapides, mais cette vitesse se fait au détriment de la qualité.

C'est précisément ce problème que les chercheurs de NVIDIA ont cherché à résoudre.

Ils ont pris un modèle de 30 milliards de paramètres et l'ont divisé en deux parties pour générer les tokens en parallèle plutôt qu'un par un. Ainsi est né Nemotron-Labs-TwoTower, un modèle de langage à diffusion issu de NVIDIA Research, adapté sur la base du modèle Nemotron-3-Nano-30B-A3B.

Mais avant de comprendre le fonctionnement de TwoTower, il est essentiel de saisir pourquoi les LLM à diffusion sont en retard par rapport aux modèles autoregressifs.

Un LLM à diffusion commence avec un bloc de tokens masqués et les révèle au fil de plusieurs passes. À chaque passe, il lit tout le texte déjà généré tout en prédisant simultanément les tokens masqués du bloc actuel. Cela signifie que le même réseau effectue deux tâches entièrement différentes à la fois : il doit comprendre le contexte existant tout en prédisant des tokens qui n'existent pas encore.

Un seul jeu de poids est dirigé vers deux tâches distinctes, empêchant le réseau de se spécialiser pleinement dans l'une ou l'autre.

TwoTower résout ce problème en ne forçant pas un seul réseau à effectuer les deux tâches. Au lieu de cela, le modèle autoregressif pré-entraîné est scindé en deux "tours".

1. La Tour de Contexte (Context Tower)

Il s'agit de la partie figée du modèle qui lit tout le texte déjà généré, exactement comme le faisait le modèle original. Elle n'est plus entraînée, préservant ainsi pleinement les capacités intellectuelles du modèle.

2. La Tour de Débruitage (Denoiser Tower)

C'est la partie entraînable du modèle qui génère chaque nouveau bloc de tokens en parallèle. Pour maintenir le contexte, elle utilise constamment une attention croisée (cross-attention) vers la Tour de Contexte.

L'aspect le plus intéressant réside dans l'interaction entre ces deux tours. Elles sont connectées couche par couche : la cinquième couche de la Tour de Débruitage interagit avec la cinquième couche de la Tour de Contexte, la sixième avec la sixième, et ainsi de suite. Cela permet au débruiteur d'accéder à toute la représentation multi-niveaux du contexte issue du modèle de base, et non seulement à sa représentation cachée finale.

Résultats clés :

• Débit de génération augmenté de 2,42 fois.
• Rétention de la qualité de 98,7 % par rapport au modèle original.
• Architecture construite sur un modèle hybride Mamba–Transformer MoE de 30 milliards de paramètres.
• Seuls environ 2,1 billions de tokens ont été utilisés pour l'entraînement — une fraction infime des 25 billions de tokens utilisés lors du pré-entraînement du modèle original.

Ce dernier chiffre est particulièrement important. Aucun composant n'a été entraîné depuis zéro. Il s'agit d'une adaptation qui peut être ajoutée à un modèle autoregressif existant. Elle transforme relativement à peu de frais le décodage lent token par token en une génération rapide de blocs en parallèle.

Par conséquent, pour tous les systèmes sensibles à la latence, c'est une approche qui mérite assurément d'être explorée. Il n'est plus nécessaire de choisir entre un modèle fiable et la vitesse requise.

Les poids, le code et la méthodologie d'entraînement sont publiés en accès libre sur Hugging Face.

Article scientifique