Résumé Technique : Modèles de Langage de Diffusion Efficaces via le Décodage Parallèle Temporel-Spatial et l'Extrapolation de Confiance

Problématique

Les modèles de langage basés sur la diffusion (dLLM) offrent une alternative au décodage autorégressif en générant du texte par un débruitage itératif, permettant des mises à jour parallèles et des ordres de génération flexibles. Cependant, l'inférence pratique reste lourde en latence. L'inférence actuelle des dLLM nécessite des passes avant complètes sur toute la séquence sous attention bidirectionnelle pour chaque étape de débruitage. Une partie significative de ces étapes est consacrée à un « raffinement redondant », où les jetons ont déjà convergé vers leurs valeurs finales mais continuent d'être re-masqués et mis à jour en raison de critères d'arrêt conservateurs.

Les méthodes d'accélération existantes reposent principalement sur des heuristiques locales à l'étape (par exemple, des seuils de confiance fixes), des calendriers statiques ou des seuils globaux. Les auteurs soutiennent que ces approches sont fragiles car la confiance de la diffusion n'est pas uniformément calibrée selon les prompts, les domaines ou les positions des jetons. Des observations empiriques révèlent des comportements temporels et spatiaux structurés dans le décodage de diffusion, tels qu'une stabilisation retardée, une confiance oscillatoire et une convergence dépendante de la position (les jetons aux positions ultérieures se stabilisent plus tard). Par conséquent, les règles simples basées sur l'étape mènent soit à un engagement prématuré (perte de qualité), soit à une attente inutile (gaspillage de calcul), échouant à distinguer les trajectoires stables des trajectoires instables qui partagent des scores de confiance instantanés similaires.

Méthodologie

L'article reformule le décodage de diffusion comme un problème de contrôle dynamique plutôt que comme une séquence de tests de seuil indépendants. Le cadre proposé se compose de deux composants clés : le Décodage Parallèle Temporel-Spatial (TSPD) et l'Extrapolation de Confiance (CE).

1. Décodage Parallèle Temporel-Spatial (TSPD)

Le TSPD est un contrôleur apprenable et léger, conçu pour prendre des décisions robustes de fixation de jetons en analysant à la fois l'historique de la trajectoire d'un jeton et sa position spatiale au sein de la séquence.

• Caractéristiques d'entrée : Pour chaque jeton à chaque étape, le TSPD construit un vecteur de caractéristiques incluant :
  • Preuve instantanée : Confiance Top-1 et entropie prédictive.
  • Tendance temporelle : Momentum dérivé d'une moyenne mobile exponentielle de la confiance.
  • Contexte spatial : La position relative du jeton au sein du bloc de décodage.
• Architecture : Il utilise un contrôleur récurrent partagé (par exemple, un LSTM ou un GRU léger) pour maintenir un état caché pour chaque position de jeton, capturant la trajectoire de débruitage au fil du temps.
• Mécanisme de décision : Le contrôleur produit une décision binaire « fixer ou continuer ». Contrairement aux prédicteurs basés sur des seuils, le TSPD utilise un estimateur straight-through (STE) pendant l'entraînement pour simuler des décisions de partitionnement (gating) strictes tout en rétropropageant à travers des logits continus.
• Entraînement : Le contrôleur est entraîné via une imitation d'oracle utilisant l'entropie croisée binaire. L'étiquette de l'oracle marque un jeton comme « fixable » une fois que son candidat prédit correspond au jeton cible final, permettant au modèle d'apprendre les schémas temporels complexes de convergence.

2. Extrapolation de Confiance (CE)

La CE est un module plug-in, sans entraînement, qui agit comme un mécanisme de prévision proactif pour réduire le temps d'attente lorsque la confiance augmente mais n'est pas encore décisive.

• Modélisation d'espace d'état : La CE modélise l'évolution de la confiance des jetons comme un processus d'espace d'état latent (similaire à un filtre de Kalman). Elle traite la confiance observée comme une observation bruitée d'un niveau de confiance et d'une vélocité latents.
• Prévision : Elle prévoit les tendances de confiance à court terme ainsi que l'incertitude quantifiée.
• Look-Ahead Sensible au Risque : La CE utilise une règle de sélection d'horizon sensible au risque. Elle détermine l'horizon de prévision maximal basé sur deux facteurs :
  1. Couverture du contexte gauche : Garantir que suffisamment de jetons précédents sont fixés pour fournir un contexte fiable.
  2. Bornes d'incertitude : Calculer une borne inférieure conservatrice de la confiance prévue.
• Intégration : Si la borne inférieure de la confiance prévue dépasse un seuil minimum, le système engage le jeton de manière anticipée. Ce module est compatible avec n'importe quel contrôleur de décodage en aval (TSPD, Learn2PD ou règles heuristiques) sans modifier les paramètres du modèle de base.

Contributions Clés

1. Cadre TSPD : La proposition d'un contrôleur de décodage conscient du temporel-spatial qui utilise les caractéristiques de trajectoire par jeton (confiance, entropie, momentum) et le contexte positionnel pour prendre des décisions robustes de fixation de jetons, atténuant la fragilité des heuristiques locales à l'étape.
2. Extrapolation de Confiance (CE) : L'introduction d'un module basé sur l'espace d'état, sans entraînement, qui prévoit la confiance future avec incertitude, permettant des décisions de look-ahead proactives et limitées par le risque pour accélérer la convergence.
3. Évaluation Complète : Évaluations étendues sur plusieurs backbones de dLLM (LLaDA-8B, Dream-7B, LLaDA-MoE), domaines de tâches (arithmétique, code, raisonnement) et longueurs de génération, démontrant des accélérations substantielles avec une précision quasi sans perte et une compatibilité totale avec le cache KV.

Résultats Expérimentaux

Le cadre a été évalué sur le modèle LLaDA-8B-Instruct sur quatre benchmarks : GSM8K, MATH, HumanEval et MBPP.

• Vitesse et Efficacité :
  • Sur GSM8K avec une génération de 256 jetons, la méthode a obtenu une accélération de 5,0× sans cache KV et de 11,2× lorsqu'elle est combinée avec le cache KV, atteignant 77,3 jetons par seconde (TPS).
  • Pour des générations plus longues (1024 jetons), l'accélération s'est cumulée à 58,3× (64,1 TPS), surpassant nettement les bases de référence comme Fast-dLLM et Learn2PD.
  • La méthode a montré des améliorations constantes à travers différentes architectures de modèles (Dream-7B, LLaDA-MoE) et longueurs de génération.
• Précision :
  • L'accélération a été réalisée avec une perte de précision négligeable. Dans de nombreux cas, la précision est restée comparable à la base de référence vanilla, et dans certains cas (ex: GSM8K avec cache KV), elle a été légèrement améliorée ou maintenue au même niveau (ex: 79,4 % contre 79,3 % pour Vanilla).
• Études d'Ablation :
  • Supprimer le TSPD a réduit l'accélération de 5,0× à 2,9×, confirmant que la fixation de jetons sensible à la trace est le principal moteur d'efficacité.
  • Supprimer la CE a réduit l'accélération de 5,0× à 4,4×, montrant la valeur additive de l'extrapolation proactive.
  • Le contrôleur s'est avéré robuste aux variations d'échelle et de distribution des données d'entraînement, une architecture LSTM légère à 2 couches fournissant l'équilibre optimal entre complexité et performance.
• Compatibilité :
  • La méthode est orthogonale aux optimisations au niveau du système. Combiner TSPD et CE avec les stratégies Dual Cache et Prefix Cache a généré des gains de débit supplémentaires significatifs sans conflits architecturaux.

Signification et Revendications

Les auteurs affirment que leur travail répond à l'inefficacité fondamentale de l'inférence des dLLM causée par le débruitage répétitif et les critères d'arrêt rigides. En traitant le décodage comme un problème de contrôle séquentiel exploitant les motifs temporels-spatiaux, le cadre proposé atteint un compromis vitesse-qualité supérieur par rapport aux bases de référence de l'état de l'art.

La signification de ce travail réside dans sa capacité à rendre les dLLM pratiquement viables pour les applications sensibles à la latence. La méthode nécessite un surcoût de calcul minimal (le contrôleur ajoute moins de 0,5 % de la latence par étape) et n'introduit aucun nouveau paramètre au modèle de diffusion de base. De plus, la nature sans entraînement du module d'extrapolation de confiance permet une intégration facile dans les pipelines existants. Les auteurs concluent que leur approche facilite le déploiement des dLLM dans des environnements à ressources contraintes, tels que les appareils de bord (edge devices) et les systèmes en temps réel, en réduant drastiquement la latence d'inférence tout en préservant les capacités de génération de haute qualité des modèles de diffusion.