Who Should Lead Decoding Now? Tracking Reliable Trajectories for Ensembling Masked Diffusion Language Models

Résumé Technique : Qui doit diriger le décodage maintenant ? Suivi de trajectoires fiables pour l'assemblage de modèles de langage de diffusion masqués

Énoncé du problème
Les modèles de langage de diffusion masqués (MDLM) sont apparus comme un paradigme distinct pour la génération de séquences, utilisant un re-masquage itératif pour débruiter les séquences en parallèle plutôt que de manière autorégressive. À mesure que la famille des MDLM se diversifie avec des modèles présentant des forces et des dynamiques de décodage différentes, un défi critique surgit : comment orchestrer ou fusionner efficacement les connaissances de MDLM hétérogènes. Les stratégies d'assemblage traditionnelles pour les modèles autorégressifs — telles que la moyenne des probabilités du jeton suivant ou l'aiguillage vers le modèle le plus confiant — ne sont pas directement applicables aux MDLM. Dans les MDLM, les séquences sont générées dans un ordre flexible et non de gauche à droite, ce qui signifie que différents modèles peuvent opérer sur des séquences partiellement débruités différentes à chaque étape. Ce manque d'une position commune de « jeton suivant » rend les définitions d'assemblage standard difficiles à appliquer, laissant l'assemblage des MDLM largement inexploré.

Méthodologie : TIE (Ensemble Itératif basé sur les Trajectoires)
Pour répondre à cela, les auteurs proposent TIE, un cadre de fusion de connaissances conçu spécifiquement pour les MDLM. TIE repose sur deux observations fondamentales dérivées de l'analyse des dynamiques de décodage :

1. Stabilité des trajectoires correctes : Les échantillons correctement répondus présentent des dynamiques plus stables et plus confiantes sur les jetons pertinents pour la réponse (entropie plus faible, probabilité top-1 plus élevée et moins de changements de jetons) par rapport aux échantillons incorrects, même lorsque les jetons sont encore masqués.
2. Correctibilité via l'injection : Les modèles générant des réponses incorrectes peuvent souvent être orientés vers des réponses correctes s'ils reçoivent des états intermédiaires partiellement débruités prometteurs provenant de modèles suivant des trajectoires plus fiables.

Sur la base de ces observations, TIE implémente un cycle récurrent en trois étapes pendant le processus de génération :

1. Génération de trajectoire : Chaque modèle constitutif ($M$) décode indépendamment pendant un nombre fixe d'étapes (l'intervalle d'assemblage, $n$), produisant une trajectoire de décodage partielle.
2. Évaluation de la trajectoire : Le cadre évalue la fiabilité de chaque trajectoire en scorant les dynamiques de confiance sur les positions de jetons de réponse (jetons apparaissant après un suffixe spécifique tel que « Answer: »). Le scoring utilise quatre métriques :
   • Nombre de changements de jetons négatifs : Suit l'historique des changements de jeton top-1 sur les positions de réponse masquées. Un compte plus faible indique une plus grande stabilité.
   • Probabilité Top-1 : La probabilité du jeton le plus probable.
   • Entropie Négative : Mesure la certitude de la distribution.
   • Marge de Probabilité : L'écart entre les probabilités top-1 et top-2.
   • Scoring Cross-Modèle : Pour les métriques basées sur les logits, le cadre fait passer chaque trajectoire par tous les modèles constitutifs et calcule la moyenne des scores. Cela atténue les différences de calibration entre les modèles, garantissant que la trajectoire sélectionnée est systématiquement confiante à travers l'ensemble plutôt que simplement favorisée par son propre modèle source.
3. Relais de trajectoire : La trajectoire ayant le score agrégé le plus élevé est sélectionnée. Cette trajectoire « gagnante » est transmise à tous les modèles constitutifs, remplaçant leur état partiel actuel. Les modèles reprennent ensuite le décodage indépendant à partir de cet état intermédiaire partagé et plus fiable.

Ce cycle se répète jusqu'à ce que la génération se termine. Enfin, le système sélectionne la réponse finale parmi les $M$ candidats en fonction de la trajectoire présentant la dynamique de jetons de réponse la plus stable (compte de changement de jetons normalisé le plus bas).

Contributions Clés

• Nouvelle analyse de décodage : L'article identifie que les jetons pertinents pour la réponse dans les générations correctes de MDLM suivent des trajectoires de confiance stables, fournissant un signal pour identifier les chemins fiables avant la génération de la réponse complète.
• Cadre d'assemblage spécifique aux MDLM : TIE introduit un mécanisme pour fusionner les connaissances en relayant itérativement des trajectoires partielles, permettant aux modèles de se corriger mutuellement dans leurs écarts au cours du processus de débruitage plutôt que d'attendre les sorties finales.
• Mécanisme de scoring Cross-Modèle : Les auteurs démontrent la nécessité d'évaluer les trajectoires à travers tous les modèles constitutifs pour tenir compte des disparités de calibration, empêchant une sélection biaisée vers des modèles trop confiants.

Résultats Expérimentaux
Les auteurs ont évalué TIE sur huit benchmarks dans quatre domaines : raisonnement général (MMLU, ARC-Challenge, Winogrande), raisonnement mathématique (GSM8K, MATH500), codage (HumanEval, MBPP) et planification (Countdown).

• Gains de performance : TIE a systématiquement surpassé les MDLM individuels (LLaDA-1.5, Dream-7B, DreamCoder-7B, DiffuCoder-7B) et les bases d'assemblage post-génération. Par exemple, sur GSM8K, l'assemblage de LLaDA et Dream via TIE a amélioré la précision de 78,39 % (Dream) à 83,47 %.
• Métriques de scoring : Parmi les fonctions de scoring, le Compte de changement de jetons (spécifiquement la version normalisée) s'est avéré le plus robuste, obtenant les meilleurs résultats sur quatre des huit benchmarks.
• Synergie des modèles : Le cadre est plus efficace lorsque les modèles constitutifs ont des capacités individuelles comparables et fortes. Lorsque les écarts de performance sont importants (ex. : >15 %), l'assemblage devient moins efficace car les modèles plus faibles introduisent des signaux bruités.
• Leadership dynamique : L'analyse révèle que le modèle produisant la trajectoire la mieux scorée change fréquemment tout au long du processus de décodage (taux de changement de modèle de 12 à 44 % selon la tâche). Cela confirme que différents modèles apportent des forces complémentaires à différentes étapes de la génération.
• Compatibilité : TIE reste efficace lorsqu'il est combiné avec des techniques d'accélération standard des MDLM, telles que le seuillage et le démasquage top-$k$.

Signification et Revendications
L'article affirme que TIE offre une approche pratique du problème peu exploré de l'assemblage des MDLM sans nécessiter d'entraînement supplémentaire. En exploitant les dynamiques de confiance sur les jetons pertinents pour la réponse, TIE permet à des modèles hétérogènes de collaborer pour affiner les réponses, permettant aux modèles plus faibles de se rétablir de trajectoires sous-optimales. Les auteurs positionnent ce travail comme une étape significative vers une orchestration plus efficace des modèles de langage basés sur la diffusion, suggérant qu'une inférence collaborative guidée par la stabilité des trajectoires peut considérablement améliorer les capacités de raisonnement à travers diverses tâches. L'article note modestement des limites, spécifiquement que l'efficacité diminue avec de grands écarts de performance entre les modèles et que l'étude actuelle est limitée à des assemblages de deux modèles.