Résumé technique : Ne pas réentraîner, mais aligner : Adaptation des modèles de langage autorégressifs aux modèles de diffusion par alignement des représentations

Énoncé du problème

Les modèles de langage de diffusion (DLM) offrent des avantages distincts par rapport aux modèles autorégressifs (AR) standards, en particulier dans la génération non séquentielle, l'édition bidirectionnelle et le raffinement itératif. Cependant, la mise à l'échelle des DLM est coûteuse en termes de calcul. Théoriquement, les DLM doivent apprendre $L!$ chemins de génération, contre un seul chemin de gauche à droite pour les modèles AR, ce qui nécessite souvent beaucoup plus de puissance de calcul.

Les approches récentes tentent d'atténuer ce coût en convertissant des points de contrôle (checkpoints) AR préentraînés en DLM. Bien que ces méthodes réussissent à initialiser des DLM à partir de poids AR, elles reposent principalement sur un entraînement continu de débruitage avec des modifications des objectifs, des motifs d'attention ou des conventions d'échantillonnage. Crucialement, les recettes existantes ne contraignent pas explicitement le DLM converti à préserver la géométrie des représentations internes apprise lors du préentraînement AR original. Cela soulève une question centrale : est-il nécessaire de traiter les modèles AR et DLM comme des paradigmes disjoints nécessitant la réapprentissage des représentations linguistiques, ou la structure sémantique apprise par le préentraînement AR peut-elle être explicitement préservée lors de la conversion ?

Méthodologie : REPR-ALIGN

Les auteurs proposent REPR-ALIGN, une méthode d'adaptation minimaliste qui traite l'entraînement des DLM comme un problème d'alignement plutôt que comme un problème d'apprentissage de représentations. L'hypothèse centrale est que la « partie difficile » de la génération de langage réside dans l'apprentissage des structures sémantiques et syntaxiques (représentations), ce que le préentraînement AR a déjà maîtrisé. La tâche restante consiste simplement à adapter ces caractéristiques à un mécanisme de décodage itératif et d'ordre quelconque.

La méthode implique deux modèles ayant une architecture identique :

1. Enseignant : Un transformateur autorégressif (AR) préentraîné et figé, avec une attention causale.
2. Élève : Un modèle de langage de diffusion masqué (DLM) avec une attention bidirectionnelle, initialisé à partir des mêmes poids AR.

Procédure d'entraînement

Pendant l'entraînement, le modèle élève consomme une séquence corrompue (avec des tokens masqués) et tente de prédire les tokens masqués en utilisant un objectif standard de débruitage masqué. Simultanément, le modèle enseignant figé traite la séquence propre sous attention causale.

Pour préserver la géométrie interne du modèle AR, les auteurs introduisent une fonction de perte d'alignement des représentations au niveau des couches. Pour chaque couche $\ell$ de la pile de transformateurs, les états cachés de l'élève ($h^{(\ell)}_D$) sont alignés avec les états cachés de l'enseignant ($h^{(\ell)}_{AR}$) en utilisant une fonction de perte de similarité cosinus.

L'objectif d'entraînement total est :
$$\mathcal{L}(\theta) = \mathcal{L}_{diff}(\theta) + \lambda \mathcal{L}_{align}(\theta)$$
Où :

• $\mathcal{L}_{diff}$ est la perte standard de l'entropie croisée de débruitage masqué.
• $\mathcal{L}_{align}$ est la distance cosinus moyenne entre les états cachés de l'élève et de l'enseignant sur toutes les couches et toutes les positions masquées.
• $\lambda$ est un hyperparamètre de pondération (valeur par défaut 10).

Les choix de conception clés incluent :

• Pas d'adaptateurs : La méthode ne nécessite aucun paramètre supplémentaire ni aucun changement d'architecture au-delà du passage du masque d'attention de causal à bidirectionnel.
• Enseignant figé : L'enseignant AR reste figé, fournissant un système de coordonnées représentatif stable.
• Convention de décalage : La méthode respecte la convention de décalage standard utilisée dans l'adaptation AR→DLM pour concilier l'indexation du token suivant avec le débruitage positionnel.

Contributions clés

1. Identification de la préservation des représentations : L'article identifie que la préservation explicite de la géométrie des représentations internes du modèle AR est un ingrédient manquant dans les recettes actuelles de conversion AR→DLM.
2. Une recette de conversion minimaliste : REPR-ALIGN introduit une méthode simple, sans adaptateur, qui ancre l'élève DLM à l'enseignant AR figé via un alignement cosinus au niveau des couches. Cette approche améliore la qualité de la conversion et l'efficacité de l'échantillonnage à travers différentes échelles de modèles.
3. Efficacité en données et en calcul : Les auteurs démontrent que la conversion AR→DLM n'est pas intrinsèquement gourmande en données. Avec l'alignement des représentations, l'entraînement sur un tout petit sous-ensemble de données (0,8 milliard de tokens) peut surpasser l'entraînement sur l'ensemble complet des données (50 milliards de tokens) avec le même budget d'étapes. De plus, le gel de grands blocs (incorporations et MLP) améliore le débit d'entraînement jusqu'à 2$\times$ sans dégrader la qualité.
4. Validation de la mise à l'échelle : Les auteurs ont entraîné un « oDLM » de 4 milliards de paramètres utilisant cette recette. Comparé à Dream-7B (un DLM public performant également construit sur une initialisation AR), oDLM a obtenu de meilleurs scores HumanEval et HumanEval+ pass@10, malgré l'utilisation de moins de paramètres et d'un budget de données et de calculs considérablement plus léger.

Résultats expérimentaux

La méthode a été évaluée sur des benchmarks de génération de code (HumanEval, HumanEval+, MBPP, MBPP+) en utilisant des points de contrôle Qwen3 aux échelles 0,6B, 1,7B et 4B.

• Vitesse d'adaptation : REPR-ALIGN a constamment accéléré l'adaptation. Sur HumanEval pass@10, le modèle aligné a surpassé à la fois le fine-tuning AR et l'entraînement à partir de zéro tout au long de la conversion précoce. À 200k étapes, le modèle 0,6B a enregistré un gain de 6,1 points, et le modèle 1,7B un gain de 9,4 points par rapport à la ligne de base.
• Mise à l'échelle : Les bénéfices de l'alignement ont augmenté avec la taille du modèle. Le modèle 1,7B a montré des améliorations absolues plus importantes que le modèle 0,6B, suggérant que les géométries de préentraînement AR plus grandes sont plus précieuses à préserver.
• Efficacité des données : L'entraînement sur un sous-échantillon aléatoire de 0,8 milliard de tokens avec alignement a produit des scores pass@10 plus élevés que l'entraînement sur le flux complet de 50 milliards de tokens au même nombre d'étapes.
• Gel : Le gel des incorporations de tokens et des blocs MLP a augmenté le débit d'environ 2$\times$ tout en maintenant ou en améliorant légèrement les performances, soutenant l'idée que le signal d'apprentissage principal réside dans l'adaptation du mécanisme de débruitage, et non dans le réapprentissage des représentations.
• Ablations :
  • Métrique : L'alignement cosinus a surpassé la distance L2, suggérant que le signal utile est géométrique (direction) plutôt que métrique (magnitude).
  • Poids : $\lambda=10$ a fourni le compromis optimal entre ancrage et adaptation.
  • Couches : L'alignement de toutes les couches a fourni le meilleur pass@1, tandis que les couches intermédiaires et supérieures étaient particulièrement importantes pour le pass@10.

Importance et revendications

L'article revendique que les représentations linguistiques sont universelles, indépendamment de l'ordre de génération. En redéfinissant l'entraînement des DLM comme un problème d'alignement, les auteurs soutiennent que le coût élevé du préentraînement par diffusion est largement dû à la réapprentissage inutile de représentations que le préentraînement AR a déjà fournies.

L'importance de ce travail réside dans l'offre d'une « technique de référence » pour l'entraînement de modèles de langage de diffusion qui :

• Réduit le calcul : Atteint des performances compétitives avec une fraction du coût de préentraînement standard.
• Simplifie l'entraînement : Élimine le besoin d'adaptateurs complexes ou de fine-tuning lourd, reposant plutôt sur un objectif d'alignement simple.
• Démocratise l'accès : Rend l'entraînement de DLM de haute qualité réalisable dans des régimes à faible volume de données et avec des ressources de calcul contraintes.

Les auteurs positionnent leur travail comme un passage du « réentraînement » à l'« alignement », suggérant que les avancées futures dans la génération discrète pourraient provenir d'un meilleur alignement des mécanismes plutôt que de paradigmes d'entraînement à partir de zéro.