Résumé technique : LESS : Échantillonnage par stabilité mutuelle pour les modèles de langage de diffusion

Énoncé du problème

Les modèles de langage de diffusion (dLLM) offrent une alternative prometteuse au décodage autorégressif (AR) en permettant des mises à jour parallèles de jetons et un conditionnement bidirectionnel via le débruitage itératif de séquences masquées. Cependant, leur efficacité pratique est actuellement limitée par des procédures d'échantillonnage qui reposent sur des programmes de débruitage inverse fixes. Les échantillonneurs standards exécutent un nombre prédéterminé d'étapes, gaspillant souvent de la puissance de calcul sur des positions où les prédictions se sont déjà stabilisées, tout en risquant un engagement prématuré sur des jetons instables. Une fois qu'un jeton est démasqué dans un processus de diffusion, il est généralement traité comme fixe ; si un jeton incorrect est engagé trop tôt en raison d'un programme rigide, les étapes de retour ultérieures ne peuvent pas le réviser, ce qui entraîne une dégradation de la qualité. Les échantillonneurs adaptatifs existants, sans réentraînement, traitent partiellement ce problème en utilisant des critères de confiance ou de marge, mais ils échouent souvent à capturer la pleine stabilité temporelle et distributionnelle de la trajectoire prédictive d'un jeton.

Méthodologie

Les auteurs proposent LESS (Less is More), un échantillonneur adaptatif, sans réentraînement et agnostique au modèle, qui cadre l'engagement des jetons comme un problème d'arrêt en ligne par position. LESS implémente l'échantillonnage par stabilité mutuelle, où une position masquée n'est éligible au démasquage que lorsque trois signaux de stabilité spécifiques indiquent conjointement que la prédiction s'est stabilisée.

La règle de stabilité conjointe

Pour chaque position masquée $i$ à l'étape inverse $t$, LESS évalue trois signaux dérivés de la distribution prédictive du modèle $p_{t,i}$ :

1. Confiance prédictive du Top-1 : La probabilité assignée au jeton top-1 actuel ($w_{t,i}$). Une position doit dépasser un seuil de confiance $c$.
2. Persistance du jeton Top-1 : Le jeton top-1 actuel doit rester inchangé sur une fenêtre récente de $P$ étapes inverses. Cela empêche l'engagement d'un jeton qui vient de changer, même si sa confiance actuelle est élevée.
3. Divergence de Jensen–Shannon (JSD) inter-étapes Top-K : C'est la contribution novatrice de LESS. Elle mesure la stabilité de la distribution prédictive elle-même entre des étapes inverses consécutives.
   • Au lieu d'utiliser la divergence de Kullback-Leibler (KL) asymétrique, LESS utilise la JSD, qui est symétrique et bornée.
   • Pour garantir l'efficacité computationnelle et se concentrer sur les candidats pertinents, la JSD est calculée sur un vocabulaire grossier constitué de l'union des $K$ meilleurs jetons des étapes actuelle et précédente, le reste de la masse étant regroupé dans un symbole résiduel.
   • La JSD doit être inférieure à un seuil $d$.

Une position est acceptée pour le démasquage si et seulement si :
$$\text{accept}_{t,i} = (\text{conf}_{t,i} \ge c) \land \text{persist}_{t,i} \land (\text{JSD}_{t,i} \le d)$$

Politique de démasquage

LESS emploie une politique de priorité à la frontière (frontier-first policy) pour les modèles de diffusion de séquence complète. Il tente de démasquer la position masquée la plus à gauche si elle est éligible. Si ce n'est pas le cas, il balaie une fenêtre de prévision (look-ahead window) pour trouver la première position éligible. Si aucune position n'est éligible, l'étape est sautée. Pour garantir la terminaison, un recours au démasquage forcé est déclenché après $\Pi$ étapes sautées consécutives, bien que ce recours ne soit pas comptabilisé comme un engagement « stable ».

Contributions clés

1. Formulation de l'engagement des jetons : L'article cadre le démasquage adaptatif dans les dLLM comme un problème d'arrêt en ligne par position induit par l'irréversibilité du démasquage.
2. JSD inter-étapes Top-K : Les auteurs introduisent la divergence de Jensen–Shannon inter-étapes Top-K comme un signal de stabilité symétrique et borné. Cela répond aux limites des travaux précédents où la confiance seule ne peut détecter un mouvement distributionnel en cours, et où la divergence KL est directionnelle et non bornée.
3. Échantillonneur LESS : Un échantillonneur sans réentraînement qui combine la confiance, la persistance du jeton et la JSD inter-étapes en une règle de stabilité conjointe. Il ne nécessite aucun réentraînement du modèle et est compatible avec les points de contrôle (checkpoints) de dLLM existants.
4. Validation empirique : Évaluation complète sur trois familles de dLLM (Dream-7B, LLaDA-8B, LLaDA-1.5-8B) et sept benchmarks couvrant les mathématiques, le code et la culture générale.

Résultats

L'évaluation démontre que LESS améliore significativement le compromis précision-efficacité par rapport au décodage à budget fixe et aux autres échantillonneurs adaptatifs sans réentraînement performants (Prophet et KLASS).

• Efficacité : À travers trois familles de modèles, LESS utilise en moyenne 72,1 % de moins d'étapes inverses par rapport au décodage à budget fixe (Base).
• Précision : LESS obtient un gain de précision moyen de +0,91 point par rapport à la référence à budget fixe tout en utilisant moins d'étapes. Il surpasse également Prophet et KLASS en précision moyenne sur toutes les familles de modèles.
• Performance des tâches : Les gains sont plus prononcés dans les tâches de génération de mathématiques et de code. Par exemple, sur Dream-7B, LESS atteint une précision moyenne de 65,18 % (contre 63,30 % pour Base) avec seulement 64,3 étapes (contre 256 pour Base).
• Économies de calcul : Étant donné que chaque étape inverse nécessite un passage direct (forward pass) du Transformer, la réduction des étapes se traduit directement par une réduction du calcul d'inférence. Sur Dream-7B GSM8K, LESS réduit la latence de temps réel de 19,45s à 5,16s et réduit le nombre d'évaluations directes d'environ 75 % lors de budgets d'étapes élevés.
• Études d'ablation :
  • Complémentarité des signaux : Aucun signal unique (confiance, persistance ou JSD) ne performe aussi bien que la règle conjointe. La confiance est le moteur principal, mais la persistance et la JSD fournissent des garde-fous essentiels contre l'engagement prématuré et la dérive distributionnelle.
  • Robustesse : La méthode est cohérente à travers différents modèles et régimes de décodage (séquence complète vs bloc sémi-autorégressif) en utilisant une configuration globale unique $(c=0,75, d=0,040)$.

Signification et affirmations

L'article affirme que LESS fournit un mécanisme pratique et sans réentraînement pour accélérer le décodage des modèles de langage de diffusion en engageant les jetons uniquement lorsqu'ils sont prédictibles avec confiance, persistance temporelle et stabilité distributionnelle.

Les auteurs soulignent que la stabilité est plus que la simple confiance. Une prédiction à haute confiance peut rester peu fiable si le jeton top-1 vient de changer ou si la masse de probabilité est encore en mouvement parmi les alternatives. En traitant la stabilité comme une propriété conjointe de la prédiction actuelle et de sa trajectoire récente, LESS évite les pièges de l'engagement prématuré.

Le travail positionne LESS comme orthogonal et complémentaire d'autres méthodes d'accélération :

• Il complète les méthodes de phase d'entraînement (ex: distillation) qui raccourcissent le processus inverse.
• Il complète les méthodes système (ex: mise en cache, décodage spéculatif) qui réduisent la latence par étape.
• Il agit comme une politique d'inférence prête à l'emploi qui peut être combinée avec ces méthodes sans modifier le modèle sous-jacent.

L'article conclut que la stabilité mutuelle est un critère viable pour un engagement précoce robuste des jetons, permettant aux dLLM d'atteindre une efficacité supérieure sans sacrifier, et souvent en améliorant, la qualité de la génération. Les auteurs reconnaissent que la méthode expose un compromis qualité-efficacité contrôlé par les seuils de stabilité, ce qui peut nécessiter un réglage pour des modèles ou des tâches spécifiques, bien qu'une configuration globale fixe se soit révélée efficace à travers leur suite d'évaluation diversifiée.