Le bouclage des couches réduit jusqu'à 3,3 fois le coût d'entraînement des LLM de diffusion

Les modèles de langage à diffusion avec bruit uniforme (UDLM) passent depuis longtemps pour une réponse intéressante mais coûteuse à l'autorégression. Un nouveau travail montre comment réduire nettement ce coût. En faisant passer plusieurs des premières couches du transformeur en boucle récursive à l'intérieur de chaque étape de débruitage, le looped UDLM atteint la validation NLL d'un UDLM classique entraîné sur environ 0,5 T de tokens, pour un coût jusqu'à 2,8 fois inférieur sur OpenWebText et jusqu'à 3,3 fois inférieur sur FineWeb-Edu. L'objectif d'entraînement et le sampler, eux, n'ont pas changé.

Pourquoi les UDLM coûtent cher à entraîner, et pourquoi l'accepter

Les modèles de diffusion masqués (MDLM) travaillent sur un contexte propre. Une partie des tokens est remplacée par un masque, les autres restent corrects, et le modèle reconstitue les trous. Les modèles uniformes sont plus exigeants. N'importe quel token de la séquence peut être remplacé par un token aléatoire du vocabulaire ordinaire, puis corrigé aux étapes suivantes. Pendant l'entraînement, le modèle doit débruiter le texte en s'appuyant sur un contexte qui peut lui-même être faux.

La tâche est plus difficile, c'est pourquoi les UDLM exigent plus de FLOPs que les MDLM. En contrepartie, les auteurs relèvent de meilleurs résultats sur les tâches aval liées au raisonnement. Pouvoir revenir sur un token déjà écrit donne au modèle une chance de corriger sa propre erreur, au lieu de la traîner jusqu'à la fin de la génération. Reste la question du coût, et c'est précisément celle que ce travail cherche à réduire.

Une boucle plutôt que de nouveaux paramètres

L'idée du selective looping a déjà fait ses preuves sur les MDLM à budget FLOPs égal. Plusieurs des premières couches du transformeur sont appliquées de façon récursive au sein d'une même étape de débruitage. Les poids des couches restent les mêmes, mais la profondeur effective du calcul augmente. Les auteurs ont transposé l'approche aux UDLM sans toucher à rien d'autre : ni à la fonction de perte, ni à la procédure d'échantillonnage.

Les principaux chiffres du travail :

• Entraînement : le looped UDLM atteint la validation NLL de l'UDLM de base entraîné sur ~0,5 T de tokens, pour un coût jusqu'à 2,8x inférieur sur OpenWebText.
• Sur FineWeb-Edu, le gain est plus élevé : jusqu'à 3,3x de FLOPs en moins.
• Génération : à budget NFE identique (nombre d'évaluations du réseau par échantillon), la generation perplexity est inférieure de jusqu'à 9,8 % à celle du modèle sans boucles.
• L'essentiel de l'amélioration s'obtient avec un nombre modéré de boucles ; au-delà, le gain est nettement plus faible.
• Distribution du nombre de boucles à l'entraînement : un biais vers les petites valeurs réduit de 12 % le nombre de boucles à l'inférence par rapport à un échantillonnage uniforme, sans baisse de précision sur les tâches aval.

Il est important de noter que c'est la combinaison des métriques qui compte. La validation NLL mesure la qualité de l'entraînement, tandis que la generation perplexity à NFE fixé mesure la qualité de la génération elle-même. Le gain apparaît sur les deux, et le test sur les tâches aval confirme qu'il ne se limite pas à des métriques favorables au modèle.

Combien de boucles prévoir, et à qui cela servira

La conclusion pratique est très terre à terre. Le rendement n'est pas linéaire avec le nombre de boucles, il est donc inutile de viser la récursion maximale. Une valeur modérée apporte l'essentiel de l'amélioration ; au-delà, le coût de chaque étape augmente plus vite que le bénéfice. Pour les modèles de diffusion, dont la génération exige déjà de nombreuses itérations de débruitage, c'est particulièrement important : chaque boucle supplémentaire est multipliée par le nombre d'étapes.

Un deuxième point pratique concerne la distribution du nombre de boucles. L'échantillonnage uniforme pendant l'entraînement semble le choix naturel, mais des petites valeurs plus fréquentes rendent le modèle robuste à un régime de fonctionnement allégé. Résultat : à l'inférence, il faut 12 % de boucles en moins pour la même précision. Pour les équipes qui construisent des LLM de diffusion, c'est une recette de réglage prête à l'emploi. Elle n'impose aucun changement dans le sampler et peut donc s'ajouter à un pipeline existant.

L'intérêt principal de ce travail est ailleurs. Les UDLM raisonnaient mieux, mais perdaient sur le coût d'entraînement. Si le bouclage réduit de façon stable les coûts de 2,8–3,3 fois, l'écart de prix avec les MDLM se resserre et l'argument en faveur du bruit uniforme devient plus convaincant. La question du passage à l'échelle au-delà d'OpenWebText et de FineWeb-Edu reste ouverte. Mais la direction est claire : la qualité des modèles de langage à diffusion peut s'acheter par l'architecture, et pas seulement avec de nouveaux tokens et des heures de GPU.