💬 NLP

Improving Diffusion Language Model Decoding through Joint Search in Generation Order and Token Space

Cet article introduit l'Order-Token Search, une nouvelle méthode de décodage pour les modèles de langage par diffusion qui explore conjointement l'ordre de génération et l'espace des jetons afin de surpasser les bases de référence existantes sur les tests de raisonnement mathématique et de codage.

Auteurs originaux : Yangyi Shen, Tianjian Feng, Jiaqi Han, Wen Wang, Tianlang Chen, Chunhua Shen, Jure Leskovec, Stefano Ermon

Publié 2026-02-06

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Yangyi Shen, Tianjian Feng, Jiaqi Han, Wen Wang, Tianlang Chen, Chunhua Shen, Jure Leskovec, Stefano Ermon

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez de résoudre un puzzle complexe, comme un problème de mathématiques ou un défi de programmation, mais que vous avez un assistant magique (le Modèle de Langage de Diffusion) qui ne se contente pas d'écrire la réponse de gauche à droite comme un humain qui tape une lettre. Au lieu de cela, cet assistant commence avec une page blanche remplie de marqueurs « MASK » et essaie de les remplir tous d'un coup, en devinant quels mots appartiennent à quels emplacements.

Le problème est que l'assistant est un peu chaotique. Il peut remplir les blancs dans n'importe quel ordre qu'il veut. Parfois, il trouve le bon mot pour le premier blanc, mais se retrouve bloqué parce qu'il a deviné le mauvais mot pour le second. D'autres fois, il devine les bons mots mais dans le mauvais ordre, ce qui mène à une impasse.

L'ancienne méthode : Deviner et Élaguer

Auparavant, les gens essayaient de corriger cela de deux manières principales, toutes deux présentant des défauts :

L'approche « Confiante » : L'assistant regarde les blancs pour lesquels il se sent le plus sûr de lui et les remplit en premier. C'est comme un randonneur qui ne marche que sur le sentier qui semble le plus solide.
- Le bon côté : Il obtient généralement une bonne réponse rapidement.
- Le mauvais côté : Si le « sentier solide » mène à une falaise (une mauvaise réponse), le randonneur est coincé. Il n'explore jamais d'autres chemins qui auraient pu être meilleurs.
L'approche « Aléatoire » : L'assistant choisit les blancs à remplir de manière totalement aléatoire.
- Le bon côté : Il explore une immense variété de chemins, il est donc très probable qu'il finisse par trouver la solution correcte un jour ou l'autre.
- Le mauvais côté : Il est tellement dispersé qu'il choisit rarement le meilleur chemin du premier coup. C'est comme un randonneur qui erre en cercles ; il finira par trouver le trésor, mais il creusera aussi beaucoup de trous vides avant.

La nouvelle solution : La Recherche Ordre-Token (Order-Token Search)

Les auteurs de ce papier ont introduit une nouvelle méthode appelée Order-Token Search. Voyez cela comme une Équipe d'Explorateurs travaillant ensemble.

Au lieu d'envoyer un seul randonneur ou une foule chaotique, la méthode envoie une petite équipe (un « faisceau » ou beam) d'explorateurs. Voici comment ils fonctionnent :

Des Chemins Divergents (La Recherche) : À intervalles réguliers, l'équipe se sépare. Chaque explorateur essaie une stratégie différente :
- L'explorateur A décide de remplir le premier mot manquant.
- L'explorateur B décide de remplir le dernier mot manquant.
- L'explorateur C essaie un mot différent pour le milieu.
- Analogie : Ils explorent à la fois où écrire ensuite (l'ordre) et ce qu'il faut écrire (le token).
La Fiche de Notation (L'Estimateur de Vraisemblance) : C'est la partie magique. L'équipe dispose d'un juge spécial (l'estimateur de vraisemblance) qui ne se contente pas de regarder la réponse finale. Au lieu de cela, le juge examine chaque étape franchie par les explorateurs.
- L'explorateur a-t-il fait un mouvement logique ?
- Cette phrase partielle a-t-elle du sens compte tenu de ce qui a été écrit auparavant ?
- Analogie : Imaginez un entraîneur regardant une course de relais. Si un coureur trébuche tôt, l'entraîneur n'attend pas qu'il termine la course pour l'arrêter ; il l'arrête immédiatement parce que l'étape qu'il a franchie était mauvaise.
Couper les Impasses (L'Élagage) : Le juge évalue les progrès de chaque explorateur. Si un explorateur suit un chemin qui semble peu probable de réussir (même s'il n'a pas encore terminé), l'équipe coupe ce chemin et concentre ses ressources sur les explorateurs qui sont sur les meilleures pistes.

Pourquoi cela importe

Le papier a testé cela sur des problèmes mathématiques difficiles (comme les ensembles de données GSM8K et MATH500) et des tâches de codage (HumanEval).

Le Résultat : La « Équipe d'Explorateurs » (Order-Token Search) a systématiquement trouvé les bonnes réponses plus souvent que les anciennes méthodes « Confiante » ou « Aléatoire ».
La Comparaison : Elle a si bien performé qu'elle a égalé ou même battu des méthodes nécessitant un réentraînement coûteux de plusieurs mois (comme diffu-GRPO). Cela signifie que vous pouvez rendre une IA beaucoup plus intelligente simplement en changeant sa façon de penser pendant le test, sans avoir besoin de réentraîner son cerveau lui-même.

Une note spéciale sur le Sudoku

Le papier a également testé cela sur des puzzles de Sudoku. Curieusement, cela n'a pas bien fonctionné là. Les auteurs expliquent que le Sudoku exige des règles strictes et globales (comme « aucun chiffre répété dans une ligne ») que le « score interne » de l'IA ne semble pas pouvoir comprendre. C'est comme donner à une équipe de randonneurs une carte qui ne montre pas les falaises ; peu importe la qualité de leur exploration, ils ne pourront pas éviter de tomber dans le vide si la carte est erronée. Cela suggère que pour certaines tâches, l'IA elle-même doit être entraînée différemment, et pas seulement la méthode de recherche.

En résumé

Le papier démontre qu'en laissant une IA explorer plusieurs façons différentes d'écrire (l'ordre) et plusieurs mots différents (les tokens) simultanément, et en utilisant un système de notation intelligent pour couper les mauvaises idées précocement, nous pouvons obtenir de bien meilleurs résultats avec les Modèles de Langue de Diffusion sans avoir besoin de les réentraîner. Cela transforme un jeu de devinettes chaotique en une recherche structurée et efficace de la vérité.

Résumé Technique : Améliorer le décodage des modèles de langage par diffusion grâce à une recherche conjointe dans l'ordre de génération et l'espace des jetons

Énoncé du Problème

Les modèles de langage par diffusion (DLM), particulièrement les modèles de diffusion masqués (MDM), offrent un avantage distinct par rapport aux modèles autorégressifs (AR) : la génération indépendante de l'ordre. Contrairement aux modèles AR qui génèrent les jetons strictement de gauche à droite, les MDM peuvent débruiter plusieurs positions en parallèle, permettant ainsi des ordres de génération adaptatifs. Cependant, les méthodes de décodage actuelles ne parviennent pas à exploiter pleinement cette flexibilité.

Les stratégies de décodage existantes s'engagent généralement dans une trajectoire unique par exécution, régie par des règles de re-masquage heuristiques :

Re-masquage à faible confiance : Fixe de manière gloutonne les jetons à haute confiance tôt dans le processus. Bien que cela produise souvent une forte précision sur un échantillon unique (pass@1), cela restreint la recherche à un ensemble étroit de trajectoires, entraînant une croissance lente de la couverture multi-échantillons (pass@k) et un risque d'enfermement dans des optima locaux incorrects.
Re-masquage aléatoire : Échantillonne les motifs de re-masquage de manière uniforme pour maximiser la diversité des trajectoires. Bien que cela améliore le pass@k en explorant un espace de solutions plus large, cela souffre d'un faible pass@1 car il manque de mécanisme principled pour sélectionner la trajectoire correcte parmi les candidats concurrents.

Cela crée un compromis exploration–exploitation : les méthodes actuelles soit sur-exploitent la confiance locale (manquant des solutions valides), soit explorent aveuglément (échouant à sélectionner la meilleure solution). L'article soutient que le potentiel des DLM à explorer divers ordres de génération et valeurs de jetons reste inexploité car le décodage standard n'effectue pas de recherche explicite dans cet espace conjoint.

Méthodologie : Recherche Ordre-Jetons (Order-Token Search)

Pour remédier à cela, les auteurs proposent l'Order-Token Search, un algorithme de décodage qui effectue une recherche en faisceau (beam search) dans l'espace conjoint des ordres de génération et des valeurs de jetons. L'algorithme alterne entre l'expansion de trajectoires candidates et l'élagage de celles qui sont peu prometteuses.

1. Processus de Recherche (Expansion)

L'Order-Token Search maintient un faisceau de $K$ candidats partiellement débruités. Contra 오는 à l'échantillonnage standard qui suit une séquence de décisions fixes, l'algorithme bifurque explicitement sur deux sources d'incertitude à des intervalles définis par l'utilisateur :

Ordre de Génération : Quelles positions (jetons) démasquer ensuite.
Valeurs de Jetons : Quels jetons assigner à ces positions.

Pour gérer la complexité computationnelle, la recherche utilise la diffusion par blocs. Au lieu de bifurquer à chaque étape de débruitage individuelle (ce qui engendrerait un surcoût prohibitif), l'algorithme étend les candidats aux limites de blocs de jetons contigus. À l'intérieur de chaque bloc, le débruitage est indépendant. Cela réduit les points de décision de la longueur totale de la séquence au nombre de blocs, permettant une diversité substantielle tant dans l'ordre que dans l'assignation des jetons tout en maintenant des coûts d'inférence gérables.

2. Processus d'Élagage (Estimation de Vraisemblance)

L'innovation centrale réside dans l'estimateur de vraisemblance utilisé pour noter et élaguer les trajectoires partielles. Un objectif MDM standard (vraisemblance de reconstruire la séquence complète propre à partir d'un état bruité) est instable pour les états intermédiaires comportant de nombreux masques restants, car il force le modèle à résoudre des problèmes de remplissage difficiles pour lesquels il n'a pas été optimisé lors de l'entraînement.

Les auteurs introduisent une fonction de score incrémentale au niveau de l'action :

Au lieu de noter l'état entier, l'algorithme note les actions de débruitage spécifiques entreprises pour passer d'un état plus corrompu $x_t$ à un état moins corrompu $x_s$ .
Le score $s(x_t; x_s)$ mesure avec quelle confiance le modèle peut repredire les nouveaux jetons révélés dans $x_s$ , conditionné par la prédiction de séquence complète du modèle $x_0$ (qui sert de contexte).
Cette approche évalue si les décisions passées d'une trajectoire (choix de l'ordre et des jetons) sont soutenues par le modèle, fournissant un signal à plus faible variance et plus discriminatif pour l'élagage.

Le score cumulatif d'un candidat est la somme de ces scores incrémentaux sur les intervalles de recherche. L'algorithme conserve les $K$ meilleurs candidats après chaque expansion, filtrant efficacement les chemins incorrects ou sous-optimaux avant de poursuivre le décodage.

Principales Contributions

Identification du compromis : L'article démontre empiriquement que les méthodes de décodage MDM standard occupent les deux extrémités opposées du spectre exploration–exploitation. Le re-masquage à faible confiance sur-exploite (pass@1 élevé, croissance faible du pass@k), tandis que le re-masquage aléatoire explore aveuglément (pass@k élevé, pass@1 faible).
Algorithme Order-Token Search : La proposition d'un nouvel algorithme de décodage qui recherche l'espace conjoint des ordres de génération et des valeurs de jetons. Il utilise un estimateur de vraisemblance incrémental stable pour élaguer les trajectoires partielles, convertissant la diversité des trajectoires en une amélioration de la précision sur échantillon unique.
Validation Empirique : La méthode est démontrée comme surpassant systématiquement les bases de référence sur les benchmarks de raisonnement mathématique et de codage, égalant ou dépassant les gains typiquement obtenus via des méthodes coûteuses de post-entraînement comme diffu-GRPO.

Résultats Expérimentaux

Les auteurs ont évalué l'Order-Token Search (OTS) sur GSM8K, MATH500, Countdown et HumanEval en utilisant LLaDA-8B-Instruct et LLaDA-1.5 (une variante post-entraînée par RL).

Gains de Performance : OTS a systématiquement obtenu la précision moyenne la plus élevée sur les benchmarks.
- Sur LLaDA-8B, OTS a amélioré le pass@1 de +3,1 % sur GSM8K, +3,8 % sur MATH500, +7,9 % sur Countdown et +6,8 % sur HumanEval par rapport à la base de référence de re-masquage à faible confiance.
- Sur LLaDA-1.5, OTS a obtenu des améliorations relatives similaires, atteignant une précision moyenne de 40,4 % contre 37,2 % pour la meilleure base de référence (Low-conf + MV).
Comparaison au Post-Entraînement : OTS a égalé ou surpassé les performances de modèles soumis à un entraînement par renforcement (diffu-GRPO), suggérant que la recherche au moment de l'inférence peut être aussi efficace que le réentraînement du modèle pour améliorer les capacités de raisonnement.
Études d'Ablation :
- Estimateur de Vraisemblance : La fonction de score incrémentale spécifique (Éq. 2) était critique. Les variantes utilisant la vraisemblance jointe de tous les blocs révélés ou des vraisemblances de type AR ont été nettement moins performantes, confirmant que l'estimateur sur mesure est nécessaire pour un élagage stable dans les MDM.
- Espace de Recherche : OTS a surpassé "Order Search" (recherche uniquement de l'ordre) et "Token Search" (recherche uniquement des jetons) qui utilisaient des vraisemblances de type AR, soulignant la nécessité d'une recherche conjointe et l'inadéquation des heuristiques AR avec les MDM.
Efficacité : OTS passe plus efficacement à l'échelle avec le calcul au moment du test (NFE) que le vote majoritaire. À budgets de calcul équivalents, OTS a atteint une précision plus élevée (ex: 29,3 % contre 19,9 % sur Countdown) en exploitant une recherche structurée plutôt qu'un échantillonnage non structuré.

Signification et Revendications

L'article affirme que la recherche conjointe ordre-jeton est un composant clé pour faire progresser le décodage des DLM. En explorant explicitement l'espace des ordres de génération et des valeurs de jetons et en élaguant sur la base d'un estimateur de vraisemblance stable, la méthode convertit le potentiel intrinsèque de multi-trajectoires des modèles de diffusion en une performance fiable sur échantillon unique.

Les auteurs positionnent ce travail comme une solution au goulot d'étranglement "exploration–exploitation" dans le décodage MDM actuel. Ils affirment que leur approche permet aux DLM d'échapper aux optima locaux qui piègent les méthodes glouttes et de sélectionner les solutions correctes que l'échantillonnage aléatoire manque, le tout sans nécess ৰ de réentraînement supplémentaire du modèle. Ce travail établit que les stratégies de recherche au moment de l'inférence, spécifiquement adaptées à la nature parallèle et indépendante de l'ordre des modèles de diffusion, peuvent produire des améliorations systématiques dans les tâches de raisonnement et de codage.