DSL-LLaDA: Scaling Continuous Denoising to 8B Masked Diffusion LMs

Résumé technique : DSL-LLaDA : Mise à l'échelle du débruitage continu pour les modèles de langage de diffusion masqués de 8B

Énoncé du problème
Les modèles de langage de diffusion discrète (DLM), tels que LLaDA, génèrent du texte via un décodage parallèle itératif. Bien que ces modèles aient atteint des échelles de milliards de paramètres, leur qualité d'inférence se dégrade fortement sous de faibles budgets d'étapes (génération en peu d'étapes). Le problème central est un compromis entre la longueur de la sortie et la qualité : avec un nombre fixe de passages avant (forward passes), le démasquage itératif standard produit soit des sorties courtes et de haute qualité, soit des textes longs, répétitifs et incohérents. Cet échec provient de la factorisation de la transition inverse, où chaque position de jeton est décodée indépendamment de son contexte par rapport à sa distribution marginale conditionnelle, négligeant ainsi les corrélations inter-jetons. Lorsque les budgets d'étapes sont faibles, chaque étape doit engager de nombreuses positions simultanément, amplifiant ces erreurs et menant à une terminaison prématurée ou à des boucles de répétition.

Le débruitage continu dans l'espace des plongements (embeddings) offre une solution théorique en faisant évoluer toutes les positions conjointement pour contourner la factorisation de l'espace des jetons, mais la construction d'un tel modèle à grande échelle reste un problème de conception ouvert en raison de la complexité de la géométrie des plongements, de la planification du bruit (noise scheduling) et des objectifs d'entraînement.

Méthodologie
Les auteurs proposent DSL-LLaDA, une méthode pour adapter légèrement un DLM préentraîné existant (spécifiquement LLaDA-8B-Instruct) afin de supporter le débruitage continu dans l'espace des plongements, plutôt que d'entraîner un modèle continu de zéro.

1. Adaptation par Localisation Stochastique Discrète (DSL) :

   • La méthode remplace l'entrée binaire « propre ou masquée » des DLM standards par des états continus bruités. Pour un jeton propre $x_i$, un rapport signal sur bruit (SNR) $\gamma_i$ est échantillonné, et un état bruité $z_i$ est construit comme un mélange du plongement du jeton et d'un bruit gaussien.
   • Un convertisseur softmax réinjecte ces états DSL continus dans l'espace d'entrée des plongements de l'architecture de base (backbone). Ce convertisseur produit un « masque doux » où un faible SNR correspond à un état de type masque, un SNR moyen à des jetons aléatoires peu fiables, et un SNR élevé à des plongements de jetons clairs.
   • Le backbone subit un pré-entraînement continu pendant seulement 1 000 étapes sur FineWeb-Edu. L'objectif d'entraînement reste l'entropie croisée standard des jetons, mais les états d'entrée sont désormais des mélanges continus plutôt que des masques binaires. Crucialement, le convertisseur permet au modèle de traiter ces entrées sans changements architecturaux, car l'invariance au SNR du débruiteur optimal élimine le besoin d'entrées explicites de pas de temps (timestep).

2. Inférence Continue (Échantillonneur SDE) :

   • Pendant l'inférence, le modèle maintient un état continu pour toutes les positions. Il utilise un échantillonneur basé sur une équation différentielle stochastique (SDE) (étapes de Heun) pour faire évoluer l'état d'un faible SNR (incertitude élevée) vers un SNR élevé.
   • Le modèle prédit une distribution de probabilité sur les jetons, qui est convertie en une moyenne pondérée des plongements de jetons (en utilisant les $M$ jetons les plus probables). Cette estimation de plongement met à jour l'état continu.
   • L'engagement ferme du jeton (hard token commitment) est différé jusqu'à l'étape finale. Cela permet au modèle d'affiner conjointement toutes les positions, préservant les corrélations inter-jetons tout au long du processus de génération.

3. Génération Conditionnelle :

   • Pour des tâches comme la résumé (summarization), les jetons de contexte sont figés à leurs plongements propres, tandis que la région de réponse subit un débruitage continu. Le backbone utilise une attention bidirectionnelle, conditionnant le débruitage de la réponse sur le contexte propre.

Principales Contributions

1. Conversion Légère : A démontré qu'un DLM masqué préentraîné de 8B peut être converti pour supporter le débruitage continu dans l'espace des plongements via une adaptation minimale de 1 000 étapes utilisant le DSL, évitant ainsi d'entraîner un modèle continu de zéro.
2. Évitement des modes de défaillance à faible NFE : A montré que l'inférence SDE continue évite efficacement la terminaison prématurée et les boucles de répétition caractéristiques du démasquage itératif à de faibles budgets d'étapes (NFE $\le$ 16).
3. Robustesse Sélective aux États Bruités : A identifié que l'adaptation induit une robustesse où le modèle peut corriger les jetons corrompus tout en préservant les jetons propres (correction sélective), un comportement absent des bases discrètes de même calcul.

Résultats

• Résumé Zero-Shot : Sur quatre benchmarks (XSum, CNN/DailyMail, arXiv, PubMed), DSL-LLaDA-SDE a obtenu les meilleurs scores ROUGE-1 à de faibles budgets d'étapes ($\le$ 16 passages avant) sans heuristique de contrôle de longueur. En revanche, les bases discrètes ont soit tronqué les sorties prématurément, soit souffert de taux de répétition élevés.
• Répétition et Longueur : En génération ouverte, DSL-LLaDA-SDE a maintenu des taux de répétition inférieurs à 9 % sur tous les budgets d'étapes, alors que les bases discrètes dépassaient 60 % de répétition à faible NFE ou se tr amountaient en sorties très courtes.
• Correction d'Erreurs : Le modèle adapté a démontré une « correction sélective », récupérant les jetons corrompus tout en préservant plus de 98 % des jetons propres. En comparaison, les bases discrètes avec corruption de jetons aléatoires (XDLM) ont endommagé une partie significative des jetons propres, et le débruitage de diffusion masqué standard (MDM-CPT) présentait une capacité de correction négligeable.
• Expériences de Contrôle : Les expériences de contrôle utilisant le même calcul de diffusion masquée standard avec corruption de jetons aléatoires (MDM-CPT et XDLM) n'ont pas réussi à reproduire les bénéfices du débruitage continu ou la robustesse sélective, confirmant que les améliorations découlent de l'exposition au bruit continu plutôt que d'un calcul supplémentaire seul.

Signification et Revendications
L'article affirme que le DSL-LLaDA fournit une procédure de conversion polyvalente pour les DLM masqués, permettant de tirer parti des avantages du débruitage continu (évolution conjointe des jetons, engagement différé) avec un surcoût d'entraînement minimal. Les résultats suggèrent que l'adaptation au bruit continu peut résoudre le compromis longueur-qualité inhérent au démasquage itératif lors de faibles budgets d'étapes.

Les auteurs restent modestes quant aux implications plus larges :

• L'évaluation est principalement limitée à la génération ancrée (résumé et suivi d'instructions). La création ouverte montre encore des taux de dégénérescence plus élevés (24 % contre $\le$ 4 % pour le résumé).
• La méthode est actuellement validée uniquement à l'échelle LLaDA-8B ; les performances sur d'autres architectures ou à des échelles plus grandes ne sont pas testées.
• Bien que la SDE surpasse les méthodes discrètes à faible NFE, le modèle LLaDA classique égale ou dépasse les performances de la SDE sur les tâches à longue référence à haut NFE (où la terminaison prématurée est moins problématique).
• Les auteurs ne revendiquent pas la parité avec les modèles autorégressifs (AR) de même échelle, notant que les modèles AR dominent encore sur les tâches à longue référence avec un grand nombre d'étapes.
• L'évaluation humaine est identifiée comme un travail futur.