Résumé Technique : Modèles de Langage à Diffusion de Contexte Résiduel (RCD)

Problématique

Les modèles de langage à diffusion (dLLM) offrent une alternative prometteuse aux modèles autorégressifs (AR) en permettant un décodage de jetons (tokens) en parallèle, déplaçant potentiellement l'inférence des régimes limités par la bande passante mémoire vers des régimes de l'utilisation du calcul. Cependant, les dLLM par blocs à l'état de l'art souffrent d'un écart de précision significatif par rapport aux modèles AR. Cet écart est exacerbé par le mécanisme de « re-masquage » utilisé lors de l'inférence : à chaque itération de débruitage, le modèle ne valide que les jetons les plus confiants et rejette les autres en les réinitialisant à un jeton de masque statique.

Les auteurs observent que ce processus gaspille efficacement le calcul intermédiaire effectué sur les jetons à faible confiance (rejetés). Comme le démontre l'analyse du rappel de jetons, ces distributions intermédiaires contiennent des signaux structurés et sémantiquement informatifs concernant le contexte global, qui sont cruciaux pour les étapes de décodage ultérieures. Les dLLM standards rejettent cette information, ce qui conduit à une précision sous-optimale malgré des coûts de calcul par jeton plus élevés.

Méthodologie : Diffusion de Contexte Résiduel (RCD)

Le papier propose la Diffusion de Contexte Résiduel (RCD), un cadre qui recycle le calcul des jetons rejetés en traitant leurs représentations latentes comme des « résidus contextuels » injectés à nouveau dans le modèle pour l'étape de débruitage suivante.

1. Injection de Résidus Pondérée par l'Entropie

Au lieu de rejeter les jetons à faible confiance, le RCD convertit leurs distributions de probabilité prédites en vecteurs d'embedding continus (soft tokens) et les injecte dans l'entrée de l'étape suivante.

• Construction du Résidu : Pour un jeton à la position $i$ et à l'étape $t$, l'information résiduelle $\Delta_i^{(t)}$ est calculée comme une somme pondérée des embeddings du vocabulaire basée sur la distribution de probabilité prédite $p_i^{(t)}$ :
  $$\Delta_i^{(t)} = \sum_{j=1}^{V} p_{i,j}^{(t)} E_{j,:}$$
• Connexion de type Highway : Pour éviter un désalignement de magnitude entre l'entrée originale et le résidu, le RCD utilise une connexion de type highway pour interpoler l'embedding du jeton masqué actuel avec le résidu de l'étape précédente :
  $$\tilde{e}_i^{(t)} = (1 - \alpha_i^{(t-1)}) E(x_i^{(t)}) + \alpha_i^{(t-1)} \Delta_i^{(t-1)}$$
  (où $x_i^{(t)} = [M]$ pour les jetons masqués).
• Pondération par l'Entropie : Le poids de contribution $\alpha_i^{(t)}$ est déterminé dynamiquement par l'entropie de Shannon normalisée de la distribution de probabilité du jeton prédit. Les jetons à haute entropie (faible confiance) se voient assigner des poids plus élevés, car ils sont supposés porter des informations structurelles plus critiques pour l'affinement de la séquence.

2. Pipeline d'Entraînement Découplé en Deux Étapes

L'entraînement de bout en bout du RCD est d'un coût de calcul prohibitif en raison de la dépendance récursive (l'entrée de l'étape $t$ dépend de la sortie de l'étape $t-1$), ce qui crée un graphe de calcul déroulé long, similaire aux RNN. Pour contourner les goulots d'étranglement de la mémoire, les auteurs proposent une stratégie d'entraînement en deux étapes :

• Étape 1 (Modèle de Référence) : Un dLLM léger et pré-entraîné est affiné pour servir de « Modèle de Référence » ($M_{ref}$) gelé. Il génère des distributions de probabilité et des poids d'entropie de haute qualité pour les entrées masquées.
• Étape 2 (Entraînement du Modèle Cible) : Un « Modèle Cible » ($M_{target}$) est entraîné en utilisant les signaux provenant du $M_{ref}$ gelé. Pendant l'entraînement, $M_{ref}$ fournit les distributions de probabilité et les poids d'entropie pour construire les vecteurs résiduels, qui sont ensuite injectés dans les embeddings d'entrée de $M_{target}$. Cela découple la génération des résidus de leur utilisation, permettant un apprentissage supervisé standard par étape sans rétropropagation à travers le temps.

3. Stratégie d'Inférence

Lors de l'inférence, le modèle passe à une boucle auto-référentielle où il génère ses propres résidus. Pour combler l'écart de distribution entre les phases d'entraînement (proxy) et d'inférence (auto-générée), les auteurs introduisent :

• Entropie à Température Échelonnée : Un scalaire $T_{res}$ ajuste la « douceur » de la distribution de probabilité utilisée pour le calcul de l'entropie, calibrant la confiance du modèle pour qu'elle corresponde à la distribution d'entraînement.
• Initialisation : Le processus peut être « amorcé à chaud » (warm-started) en utilisant la prédiction initiale du Modèle de Référence ou « amorcé à froid » (cold-started) avec un vecteur nul.

Principales Contributions

1. Mécanisme de Contexte Résiduel : Un module novateur qui recycle les représentations de jetons rejetées dans les dLLM, transformant le calcul gaspillé en un signal contextuel directeur.
2. Agrégation Pilotée par l'Entropie : Une méthode rigoureuse de pondération de l'information résiduelle basée sur l'entropie de Shannon normalisée, garantissant que les jetons incertains contribuent davantage à l'affinement du contexte.
3. Pipeline d'Entraînement Efficace : Une approche d'entraînement découplée en deux étapes qui évite les goulots d'étranglement de la mémoire liés à la rétropropagation à travers le temps, permettant la conversion de dLLM standards vers le paradigme RCD avec un minimum de données (environ 300 millions de jetons).
4. Nouveau Compromis Pareto : La méthode introduit un nouveau levier de compromis entre les étapes de débruitage et la transmission des résidus, permettant une précision plus élevée pour une latence équivalente ou un nombre d'étapes nettement inférieur pour une même précision.

Résultats Expérimentaux

Les auteurs ont validé le RCD sur deux familles distinctes de dLLM : LLaDA (contexte global bidirectionnel) et SDAR (semi-autorégressif par blocs).

• Gains de Précision : Le RCD améliore systématiquement les dLLM de pointe de 4 à 11 points de pourcentage sur des benchmarks incluant GSM8K, MATH500 et MinervaMath.
• Raisonnement Complexe : Sur les benchmarks exigeants AIME24/25, le RCD fait presque doubler la précision de base (par exemple, SDAR-8B-b64 est passé de 7,08 % à 18,75 % sur AIME24).
• Efficacité : Le RCD atteint ces gains avec un surcoût de calcul minimal. Il atteint la précision maximale de la ligne de base avec 4 à 5 fois moins d'étapes de débruitage.
• Efficacité des Données : Le RCD atteint une haute précision de raisonnement avec seulement ~300M de jetons d'entraînement, là où les méthodes basées sur les latents concurrentes (ex: Loopholing) échouent à générer des séquences cohérentes sous le même budget.
• Scalabilité : La méthode passe de manière robuste de modèles de 4B à 8B paramètres et à travers différents tailles de blocs (32 et 64 jetons).

Signification et Revendications

Le papier affirme que le principal goulot d'étranglement des dLLM actuels n'est pas un manque d'étapes d'entraînement ou de capacité du modèle, mais plutôt la perte d'information causée par la stratégie de re-masquage. En recyclant le calcul « gaspillé » des jetons rejetés, le RCD libère le potentiel des modèles de diffusion pour égaler ou dépasser les capacités de raisonnement des modèles autorégressifs tout en conservant les avantages du décodage parallèle.

Les auteurs positionnent le RCD comme une solution pratique et évolutive pour la génération de texte haute fidélité en parallèle. Ils soulignent que la méthode ne nécessite pas de changements architecturaux au modèle de base, mais plutôt un changement de paradigme de décodage et d'entraînement, ce qui en fait une voie viable pour améliorer la précision et l'efficacité des prochains modèles de langage à diffusion.