🤖 AI

xPress: Parallel Refinement for Diffusion Drafters in Speculative Decoding

Le papier propose xPress, un raffineur causal léger qui restaure les dépendances manquantes dans les rédacteurs par diffusion en blocs grâce à un raffinement parallèle, augmentant considérablement la longueur d'acceptation et le débit de bout en bout dans le décodage spéculatif par rapport aux méthodes existantes comme dFlash.

Auteurs originaux : Zheng Wang, Davis Wertheimer, Yu Chin Fabian Lim, Mudhakar Srivatsa, Raghu K. Ganti, Minjia Zhang, Naigang Wang

Publié 2026-08-04

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Zheng Wang, Davis Wertheimer, Yu Chin Fabian Lim, Mudhakar Srivatsa, Raghu K. Ganti, Minjia Zhang, Naigang Wang

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

La course pour parler plus vite : Pourquoi l'IA a besoin d'un meilleur système de rédaction

Imaginez que vous essayiez d'écrire une histoire, mais avec une règle stricte : vous ne pouvez écrire qu'un seul mot à la fois, et vous devez attendre qu'un éditeur super intelligent vérifie votre mot avant de pouvoir écrire le suivant. C'est ainsi que fonctionnent la plupart des modèles de langage IA les plus puissants aujourd'hui. Ils sont incroyablement précis, mais aussi douloureusement lents car ils doivent vérifier chaque mot un par un. Pour accélérer cela, les scientifiques ont inventé une astuce appelée « décodage spéculatif ». Voyez cela comme un assistant rapide, bien qu'un peu moins soigneux, qui devine les prochains mots pour vous. Si l'éditeur super intelligent est d'accord avec les suppositions de l'assistant, vous pouvez écrire tous ces mots d'un coup, évitant ainsi le temps d'attente.

Le problème est que l'assistant est généralement trop pressé. Il devine des mots qui semblent corrects individuellement mais qui ne s'assemblent pas bien dans une phrase, comme un chef qui choisit des ingrédients délicieux mais oublie de vérifier s'ils se marient bien ensemble. Récemment, un nouveau type d'assistant appelé « rédacteur par diffusion » (diffusion drafter) a été créé. Au lieu de deviner les mots un après l'autre, il essaie de deviner tout un bloc de mots d'un coup, comme si l'on jetait une poignée de pièces de puzzle sur la table. C'est extrêmement rapide, mais parce qu'il jette tout en même temps, les pièces ne s'emboîtent souvent pas correctement. L'éditeur super intelligent doit en rejeter la plupart, ce qui ralentit à nouveau tout le processus. La grande question que se posent les chercheurs est la suivante : peut-on conserver la vitesse du devineur « tout d'un coup » tout en corrigeant les erreurs pour que l'éditeur accepte réellement les mots ?

L'arrivée de xPress : Le « raffineur parallèle » qui résout le puzzle

Cet article présente une solution ingénieuse appelée xPress. Les auteurs, travaillant avec des modèles comme Qwen3-8B, ont réalisé que si le « rédacteur par diffusion » est excellent pour deviner un bloc de mots, il oublie la règle cruciale selon laquelle les mots dépendent les uns des autres (la causalité). Par exemple, si le premier mot est « elle », le suivant ne devrait pas être « sont », même si « sont » est un mot courant en soi. Le rédacteur par diffusion ne le sait pas car il devine tout simultanément.

Pour corriger cela, xPress agit comme un raffineur causal léger. Imaginez que le rédacteur par diffusion jette une poignée de pièces de puzzle sur la table. xPress est une main rapide et intelligente qui regarde l'ensemble du tas d'un seul coup d'œil et replace les pièces dans le bon ordre sans les démonter une par une. Il y parvient grâce à une technique appelée décodage de Jacobi. Au lieu de réparer le puzzle pièce par pièce (ce qui est lent), xPress regarde l'image entière, effectue un ajustement rapide sur chaque pièce simultanément, puis regarde à nouveau. Il répète ce cycle de « regarder et ajuster » seulement quelques fois (généralement 4 à 6 fois) jusqu'à ce que les pièces s'emboîtent parfaitement.

Les résultats sont impressionnants. En utilisant xPress, le système peut accepter environ 30 % de mots devinés en plus en moyenne par rapport au rédacteur rapide d'origine. Dans certains tests spécifiques, comme la résolution de problèmes mathématiques, l'accélération a atteint 56 %. En mesurant la vitesse totale de parole de l'IA, xPress a rendu le processus 1,3 fois plus rapide en moyenne, et jusqu'à 1,7 fois plus rapide dans les meilleurs cas, par rapport à la méthode originale.

L'article argumente explicitement contre deux autres méthodes par lesquelles les gens ont tenté de résoudre ce problème. Une méthode consiste à construire un arbre géant de suppositions, ce qui est complexe et coûteux à exécuter. Une autre utilise une « tête de Markov » qui corrige les mots un par un selon une ligne stricte, ce qui est précis mais lent car cela perd l'avantage de vitesse du fait de deviner tout d'un coup. Les auteurs démontrent que xPress bat les deux : il est plus rapide que le correcteur étape par étape et plus simple que la méthode complexe de l'arbre.

L'étude confirme que xPress fonctionne en réinjectant de l'« information causale » (la règle selon laquelle les mots dépendent des mots précédents) dans le système rapide sans le ralentir. Ils ont testé cela sur des benchmarks de mathématiques, de codage et de discussion, constatant que xPress surpasse systématiquement la concurrence. L'article suggère que cette approche est une amélioration mesurée et solide qui permet à l'IA d'être à la fois rapide et précise, prouvant qu'on n'a pas à sacrifier la qualité pour la vitesse si l'on possède la bonne méthode pour affiner ses suppositions.

Résumé Technique : xPress : Raffinement Parallèle pour les Drafters de Diffusion en Décodage Spéculatif

1. Énoncé du Problème

Le décodage spéculatif (SD) accélère la génération autorégressive en utilisant un drafter léger pour proposer des jetons (tokens) à vérifier par un modèle cible plus important. Bien que les drafters autorégressifs (AR) (ex. : EAGLE) soient efficaces, ils souffrent d'un surcoût séquentiel lors de la génération de blocs longs. Les drafters par diffusion de blocs (ex. : dFlash) remédient à cela en générant un bloc entier de jetons en une seule passe directe parallèle, réduisant ainsi drastiquement la latence de rédaction.

Cependant, les drafters de diffusion possèdent une limitation fondamentale d'exactitude : ils prédisent les jetons sur la base de distributions marginales plutôt que de distributions conjointes. Dans un processus de diffusion en une seule passe, le jeton $k$ est échantillonné sans être conditionné par les valeurs réalisées des jetons $1 $à$ k-1$. Ce manque de causalité intra-bloc entraîne des séquences où les jetons individuels sont localement probables mais conjointement incohérents (ex. : désaccord sujet-verbe). Lors de la vérification, le modèle cible rejette ces séquences incohérentes précocement, ce qui limite sévèrement la longueur d'acceptation ( $\tau$ ) et annule le gain de vitesse potentiel du drafting parallèle.

Les solutions existantes, telles que le drafting basé sur des arbres (ex. : PRESTO, DDTree), tentent d'atténuer ce problème en élargissant l'espace des candidats sous forme d'arbre. Cependant, ces méthodes sont coûteuses en calcul, complexes à servir et contournent souvent le problème en ciblant le rappel plutôt qu'en corrigeant l'incohérence causale sous-jacente.

2. Méthodologie : xPress

Les auteurs proposent xPress, un raffineur causal léger conçu pour restaurer les dépendances causales manquantes dans les drafters de diffusion sans revenir à une boucle autorégressive lente, jeton par jeton.

2.1 Concept Central

xPress fonctionne comme un « raffineur causal » qui prend la proposition initiale de bloc du drafter de diffusion et la réconcilie par un raffinement parallèle. Au lieu de régénérer les jetons séquentiellement, xPress injecte des informations causales dans les logits du drafter de diffusion pour corriger les distributions marginales.

2.2 Architecture

Le raffineur est un module léger qui ajoute un biais de logit appris ( $\delta_k$ ) aux logits de base du drafter ( $s_k$ ) à chaque position $k$ . L'architecture se compose de quatre étapes :

Fuse (Fusion) : Projette trois entrées dans un espace de rang faible ( $r=256$ $r = 256$ ) :
- L'état caché par position du drafter ( $h_k$ ).
- Un résumé global du bloc ( $g$ ) obtenu via un pooling moyen.
- L'embedding du jeton précédent supposé ( $x_{k-1}$ ).
Mix (Mélange) : Applique un mélangeur linéaire strictement inférieur triangulaire ( $L$ ) à travers les positions du bloc. Cela permet à la position $k$ d'agréger les informations de tous les prédécesseurs ( $j \le k$ ) en une seule opération, imposant une structure causale sans calcul séquentiel.
Trans (Transformation) : Un MLP résiduel dans l'espace de rang faible.
Read (Lecture) : Projette le résultat vers la taille du vocabulaire pour produire le biais $\delta_k$ .

Le logit raffiné final est $\ell_k = s_k + \delta_k$ .

2.3 Raffinement Parallèle via le Décodage de Jacobi

Un raffineur causal direct nécessiterait une passe séquentielle de gauche à droite, réintroduisant la latence que xPress vise à éviter. Pour résoudre cela, xPress utilise le décodage de Jacobi :

Initialisation : Le bloc est amorcé avec les prédictions gloutonnes (greedy) initiales du drafter.
Itération : Toutes les positions sont mises à jour simultanément en parallèle en utilisant l'état actuel du bloc comme préfixe pour l'itération suivante.
Convergence : Le processus se répète pour $K$ itérations. Théoriquement, pour un bloc de taille $B$ , le point fixe séquentiel est atteint en au plus $B$ itérations. Empiriquement, les auteurs constatent que $K \approx 6$ itérations suffisent pour égaler la longueur d'acceptation d'un décodage séquentiel complet, car la graine du drafter est déjà une estimation de haute qualité.

2.4 Stratégie d'Entraînement

xPress est co-entraîné avec le drafter de diffusion en utilisant une fonction de perte composite :

Cross-Entropy (Teacher Forcing) : Maximise la probabilité du jeton de vérité (ground-truth) étant donné le préfixe de vérité.
Total Variation (TV) Loss : Minimise la distance entre la distribution raffinée et la distribution du modèle cible, optimisant directement le taux d'acceptation.
Consistency Loss (Perte de Cohérence) : Une perte auxiliaire appliquée à une seconde passe directe où les entrées sont les propres prédictions du drafter (auto-conditionnement). Cela aligne la distribution d'entraînement avec le processus de Jacobi lors de l'inférence, empêchant le raffineur de sur-apprendre les préfixes de référence (gold prefixes).
Drafter Anchor (Ancre du Drafter) : Un terme de perte appliqué à la sortie du drafter de base pour garantir que le modèle de diffusion sous-jacent reste de haute qualité.

3. Contributions Clés

Injection de Dépendance Causale : xPress injecte des informations causales dans les drafters de diffusion en utilisant un mécanisme minimal (projection de rang faible et mélangeur triangulaire) qui opère sur les états cachés du drafter, évitant ainsi une augmentation massive des paramètres.
Raffinement Parallèle : En utilisant le décodage de Jacobi, xPress résout les dépendances causales en parallèle. Il atteint l'effet de $B$ étapes autorégressives dépendantes en utilisant seulement $K \ll B$ itérations parallèles, préservant les avantages de vitesse du drafting parallèle.
Gains de Bout en Bout Consistants : La méthode améliore significativement la longueur d'acceptation ainsi que le débit sans modifier le processus de vérification ou nécessiter un drafter plus grand.

4. Résultats Expérimentaux

Les expériences ont été menées sur Qwen3-8B en utilisant un drafter de diffusion de bloc-16 (dFlash) sur sept benchmarks (Math : GSM8K, MATH-500, AIME25 ; Code : HumanEval, MBPP, LiveCodeBench ; Chat : MT-Bench).

Longueur d'Acceptation : xPress augmente la longueur d'acceptation moyenne d'environ 30 % par rapport au drafter de diffusion dFlash original.
- En décodage glouton ( $T=0$ ), les gains varient de 18 % à 56 %.
- En échantillonnage ( $T=1$ ), les gains varient de 19 % à 58 %.
Accélération du Débit (Throughput) :
- Glouton ( $T=0$ ) : xPress atteint une accélération end-to-end moyenne de 1,3× par rapport au drafter dFlash, avec des pics allant jusqu'à 1,7× (ex. : sur GSM8K).
- Échantillonnage ( $T=1$ ) : L'accélération s'améliore pour atteindre une moyenne de 1,46×, avec des pics de 1,6×.
Comparaison avec les Raffineurs Séquentiels : Comparé au Markov head (un raffineur causal séquentiel de DSpark), xPress égale ou dépasse la longueur d'acceptation de ce dernier tout en étant nettement plus rapide.
- À $K=4$ itérations de Jacobi, xPress égale la précision du Markov head tout en étant 1,6× plus rapide (379 µs contre 601 µs par bloc).
- Le Markov head nécessite un décodage séquentiel fixe de 15 étapes, tandis que xPress converge en moins d'étapes parallèles.
Scalabilité : Dans des scénarios de service à gros lots (vLLM, tailles de batch 16 et 32), xPress maintient son avantage. Bien que l'accélération absolue par rapport à la ligne de base AR se comprime à mesure que les lots augmentent (en raison de la vérification limitée par le calcul), xPress surpasse systématiquement le simple drafter dFlash de 1,2 à 1,4× en débit, prouvant que les bénéfices du raffineur survivent à la charge.

5. Signification et Revendications

L'article affirme que xPress traite le « décalage inhérent » entre le drafting de diffusion parallèle et la vérification séquentielle en réparant la distribution marginale elle-même plutôt qu'en élargissant l'espace de recherche.

Efficacité : xPress démontre que l'information causale peut être injectée à un coût quasi nul par rapport au surcoût du drafting. Le coût du raffinement est une fraction minime du temps de rédaction total, et la nature parallèle du raffinement évite les pénalités de latence du traitement séquentiel.
Orthogonalité : La méthode est orthogonale aux autres améliorations ; elle peut être combinée avec des drafters plus grands ou des méthodes de vérification basées sur des arbres, car elle améliore les logits de la « structure dorsale » sur lesquels ces méthodes reposent.
Généralisabilité : Les auteurs suggèrent que la philosophie consistant à réconcilier des prédictions mutuellement dépendantes via quelques itérations parallèles peut s'appliquer au-delà du décodage spéculatif, partout où les modèles émettent simultanément des sorties interdépendantes.

L'article conclut que xPress réussit à restaurer la causalité perdue dans le drafting de diffusion en une seule passe, permettant des chaînes d'acceptation plus longues et un débit plus élevé sans sacrifier l'efficacité parallèle qui rend les drafters de diffusion attractifs.