🤖 AI

PRESTO: Prefix-Aligned Tree Drafting for Diffusion Speculative Decoding

Le papier introduit PRESTO, un cadre fondé sur des principes qui améliore le décodage spéculatif basé sur la diffusion en implémentant une notation alignée sur les préfixes et une recherche arborescente basée sur la priorité afin de résoudre le décalage entre les marginales de diffusion et la vérification autorégressive, améliorant ainsi considérablement le débit de bout en bout.

Auteurs originaux : Zheng Wang, Zhifan Ye, Qi Cheng, Yonggan Fu, Ziyan Wang, Feng Zhu, Haozhe Zhao, Jan Kautz, Pavlo Molchanov, Humphrey Shi, Minjia Zhang

Publié 2026-07-28

📖 11 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Zheng Wang, Zhifan Ye, Qi Cheng, Yonggan Fu, Ziyan Wang, Feng Zhu, Haozhe Zhao, Jan Kautz, Pavlo Molchanov, Humphrey Shi, Minjia Zhang

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 prédire le mot suivant dans une histoire. Pendant longtemps, les ordinateurs les plus intelligents (appelés Grands Modèles de Langage) faisaient cela un mot à la fois, comme une personne lisant un livre à voix haute, s'arrêtant après chaque mot pour réfléchir à ce qui vient ensuite. C'est précis, mais c'est lent. Récemment, des scientifiques ont découvert une nouvelle façon d'écrire ces histoires en utilisant des modèles de « diffusion ». Pensez à ces modèles comme à un sculpteur qui commence par un bloc de marbre et sculpte toute la statue d'un coup, plutôt que de tailler chaque petit morceau un par un. Cela permet à l'ordinateur de deviner de nombreux mots simultanément, ce qui est incroyablement rapide.

Cependant, il y a un piège. Quand vous devinez de nombreux mots à la fois, vous pouvez en rater quelques-uns. Pour corriger cela, il existe une astuce ingénieuse appelée « décodage spéculatif ». C'est comme avoir un jeune assistant rapide qui devine les prochains mots, puis un patron super intelligent et lent qui vérifie si ces devinettes sont correctes. Si le patron est d'accord, il accepte tout le lot de mots instantanément, ce qui fait gagner un temps fou. Le problème est que l'assistant junior (le modèle de diffusion) est excellent pour deviner des mots individuels, mais il ne sait pas toujours comment ces mots s'assemblent dans un ordre spécifique. C'est comme si l'assistant était excellent pour choisir les ingrédients individuels d'un gâteau, mais qu'il ne savait pas toujours quelle combinaison d'ingrédients aura un bon goût.

C'est là qu'intervient un nouvel article de recherche. Les chercheurs, dirigés par Zheng Wang et ses collègues, ont réalisé que la façon actuelle d'utiliser ces assistants rapides laisse beaucoup de vitesse sur la table. Ils ont découvert que, bien que l'assistant puisse générer une immense variété de combinaisons de mots, la méthode actuelle ne vérifie qu'un seul chemin, comme si l'on marchait dans un seul couloir en espérant que la porte au bout soit ouverte. Les auteurs proposent un nouveau système appelé PRESTO (Prefix-Aligned Tree Drafting). Au lieu de marcher dans un seul couloir, PRESTO construit un arbre de possibilités, explorant de nombreux chemins à la fois. Mais voici la magie : il corrige une faille fondamentale dans la manière dont la confiance de l'assistant est mesurée. La confiance originale de l'assistant est « aveugle au préfixe » (prefix-blind), ce qui signifie qu'elle ne se soucie pas de savoir quels mots sont venus avant. PRESTO ajoute un score « aligné sur le préfixe », qui agit comme une boussole, garantissant que les chemins choisis pour l'exploration sont ceux qui sont les plus susceptibles d'être acceptés par le patron.

Le résultat est un système nettement plus rapide. Dans leurs tests, PRESTO a aidé l'ordinateur à accepter plus de mots lors de chaque tour de devinette. Sur certaines des meilleures configurations existantes, il a rendu le processus 1,5 fois plus rapide. Sur d'autres configurations, il a fourni une accélération de 1,12 fois. L'article suggère qu'en traitant le processus de devinette comme une aventure de ramification d'arbre plutôt que comme une ligne droite, et en s'assurant que les branches choisies correspondent bien à l'histoire en cours, nous pouvons obtenir le meilleur des deux mondes : la vitesse de la diffusion et la précision de la vérification minutieuse.

Le Problème : Le Piège du « Chemin Unique »

Pour comprendre pourquoi PRESTO est nécessaire, imaginez que vous jouez à un jeu de « Mad Libs » avec un ami qui essaie de deviner les mots manquants. Votre ami est un modèle de diffusion. Il est incroyable pour regarder un espace vide et dire : « Je parie que le mot ici est "chat" ! » ou « Peut-être que c'est "chien" ? » ou « Ou peut-être "fusée" ? ». Il peut crier toutes ces options en même temps.

Cependant, la méthode actuelle d'utilisation de cet ami est très rigide. Elle prend la meilleure supposition de l'ami, l'écrit, puis demande au « patron » (le modèle cible) si cette supposition est correcte. Si le patron dit « Non », tout est jeté, et vous devez recommencer. Si le patron dit « Oui », vous passez au mot suivant et recommencez. C'est ce qu'on appelle le drafting linéaire (linear drafting). C'est comme marcher dans une forêt et ne regarder que le chemin directement devant soi.

Les auteurs ont observé que cette approche est inefficace. Parce que le modèle de diffusion génère de nombreuses options simultanément, il existe un immense « espace combinatoire » de possibilités. C'est comme avoir une carte avec mille sentiers différents, mais que vous n'ayez le droit d'en suivre qu'un seul. L'article montre qu'en s'en tenant à un seul chemin, le système manque de nombreuses routes valides. En fait, sur des problèmes mathématiques comme GSM8K, la méthode actuelle acceptait environ 6,5 mots en moyenne, mais les chercheurs ont calculé que s'ils pouvaient vérifier tous les meilleurs chemins, ils auraient pu en accepter près de 10. C'est un écart énorme !

Le Décalage : La Boussole « Aveugle »

Les chercheurs ont creusé plus loin et ont découvert une raison spécifique pour laquelle le simple fait de vérifier plus de chemins (construire un arbre) ne fonctionnait pas parfaitement avec les anciennes méthodes. Ils ont identifié un « décalage fondamental ».

Dans le monde de l'IA standard (les modèles autorégressifs), le score de confiance pour un mot dépend fortement des mots qui l'ont précédé. Si la phrase est « Le chat est assis sur le... », le modèle sait que « tapis » est un mot très probable, mais que « pizza » ne l'est pas. C'est ce qu'on appelle être aligné sur le préfixe (prefix-aligned).

Mais les modèles de diffusion fonctionnent différemment. Ils génèrent une probabilité « marginale » pour chaque position de manière indépendante. C'est comme si le modèle disait : « À la position 5, "chat" a 80 % de chances d'apparaître », sans se soucier de savoir si la position 4 était « Le » ou « Le rapide chat brun ». C'est ce qu'on appelle être aveugle au préfixe (prefix-blind).

Lorsque vous essayez de construire un arbre de suppositions en utilisant ces scores aveugles, vous obtenez un problème de classement. Vous pourriez choisir un chemin qui semble excellent pour le premier mot mais terrible pour le second parce que le modèle n'a pas réalisé que le premier mot a changé le contexte. C'est comme un GPS qui vous donne des directions basées uniquement sur la rue actuelle, ignorant le fait que vous venez de tourner à gauche et que vous êtes maintenant sur une rue à sens unique. L'article soutient que l'utilisation de ces scores aveugles pour construire un arbre conduit à un « classement de chemin peu fiable », ce qui signifie que le système explore les mauvages branches et perd du temps.

La Solution : PRESTO

PRESTO (Prefix-Aligned Scoring and priority-based Tree search for diffusion Speculative decOding) résout cela en ajoutant une « correction » aux scores du modèle de diffusion.

Scoring Aligné sur le Préfixe (Prefix-Aligned Scoring) : Les auteurs ont réalisé qu'ils devaient combiner le signal « marginal » fort du modèle de diffusion (la probabilité d'un mot seul) avec un signal « conditionné par le préfixe » (sa probabilité étant donné les mots précédents). Ils ont créé une nouvelle formule de score qui multiplie la probabilité de diffusion par un facteur de correction dérivé d'un modèle n-gramme simple (un outil léger qui observe les combinaisons de mots). Cela crée un score qui respecte le flux de l'histoire.
Recherche en Arbre basée sur la Priorité (Priority-Based Tree Search) : Au lieu de simplement choisir le meilleur chemin, PRESTO construit un arbre. Il utilise les nouveaux scores corrigés pour décider quelles branches faire croître. Il donne la priorité aux chemins qui ont le plus grand potentiel d'être acceptés par le patron. C'est comme un randonneur qui, au lieu de marcher tout droit, regarde une carte et choisit le sentier le plus susceptible de mener au sommet, même si ce sentier n'est pas le plus évident au départ.

L'article teste deux façons de faire croître cet arbre : la Recherche en faisceau (Beam Search - garder un nombre fixe de chemins principaux à chaque étape) et la Recherche de type "Best-First" (Best-First Search - toujours étendre le meilleur chemin trouvé jusqu'à présent). Ils ont constaté que pour leur configuration spécifique, la Recherche en faisceau fonctionnait aussi bien que la recherche Best-First plus complexe, ils ont donc opté pour l'option la plus simple et la plus efficace.

Les Résultats : Plus Rapides et Plus Intelligents

Les auteurs ont mis PRESTO à l'épreuve sur diverses tâches, notamment des problèmes mathématiques (GSM8K, Math500), des défis de codage (HumanEval, LiveCodeBench) et des conversations de chat. Ils ont utilisé deux types de systèmes différents :

Des drafters de diffusion dédiés : Un petit modèle de diffusion rapide qui devine pour un modèle autorégressif plus grand (comme dFlash).
Des LLM de diffusion auto-spéculatifs : Un seul modèle de diffusion qui devine et vérifie son propre travail (comme Nemotron-Labs-Diffusion).

Les résultats sont cohérents sur toute la ligne. PRESTO augmente systématiquement la Longueur de l'Acceptation Moyenne (Average Acceptance Length), qui est le nombre de mots que le patron accepte en une seule fois.

Sur le système dFlash (utilisant Qwen3-8B), la longueur d'acceptation moyenne est passée d'environ 6,6 mots à 9,6 mots. Cela s'est traduit par une accélération globale de 1,5 fois.
Sur le système Nemotron-Labs-Diffusion, la longueur d'acceptation est passée de 8,8 à 9,9 mots, entraînant une accélération de 1,12 fois.

Plus impressionnant encore, l'article montre que PRESTO fonctionne même lorsque le système est « stochastique » (aléatoire), ce qui est habituellement plus difficile à prédire. Dans ces cas, l'accélération est encore plus prononcée, certains benchmarks voyant un débit presque doublé.

Les auteurs ont également vérifié le « coût » de cette nouvelle méthode. Ils ont trouvé que le travail supplémentaire requis pour construire l'arbre et calculer les nouveaux scores était minime — moins de 4 % du temps total. La grande majorité du temps (plus de 90 %) est toujours consacrée à la vérification réelle par le modèle patron. Cela signifie que PRESTO est une mise à niveau hautement efficace qui ne ralentit pas le système avec une surcharge inutile.

Ce que PRESTO N'EST PAS

Il est important de noter ce que l'article ne prétend pas. Les auteurs déclarent explicitement qu'appliquer une structure d'arbre naïve (sans leur scoring aligné sur le préfixe) est sous-optimal. Si vous prenez simplement les scores bruts du modèle de diffusion et construisez un arbre, vous ne tirez pas pleinement parti de l'avantage car à cause de la nature « aveugle » des scores. PRESTO vise spécifiquement à corriger ce décalage de scoring.

De plus, l'article ne prétend pas avoir résolu entièrement le problème des modèles de diffusion. Ils reconnaissent que leur méthode repose sur un « signal aligné sur le préfixe traçable » (comme le modèle n-gramme qu'ils ont utilisé) pour effectuer le gros du travail de correction. Ils suggèrent que les travaux futurs pourraient explorer des signaux encore plus riches, mais pour l'instant, leur correction simple est suffisante pour voir des gains massifs.

Pourquoi cela compte

Dans la course à rendre l'IA plus rapide et plus efficace, chaque gain de vitesse compte. Le décodage spéculatif est un sujet brûlant car il nous permet d'utiliser les meilleurs modèles, les plus précis, sans payer le plein prix en temps. Cependant, les méthodes actuelles étaient limitées par le fait qu'elles traitaient les modèles de diffusion comme des machines linéaires, ignorant leur capacité unique à générer de nombreuses options à la fois.

PRESTO change la donne en traitant les modèles de diffusion comme les explorateurs multi-chemins qu'ils sont. En alignant le scoring sur la façon dont le patron vérifie le travail, il débloque tout le potentiel de la génération parallèle des modèles de diffusion. Le résultat est un système qui n'est pas seulement légèrement plus rapide, mais nettement plus efficace, permettant de générer du texte, de résoudre des problèmes mathématiques et d'écrire du code à des vitesses qui étaient auparavant jugées impossibles pour ces types de modèles. Comme le disent les auteurs, ils ont transformé une marche sur un « chemin unique » en une expédition basée sur un « arbre », garantant que chaque pas fait est un pas vers la bonne réponse.

Résumé Technique : PRESTO

Problématique

Les modèles de langage à diffusion (dLLM) offrent une alternative prometteuse aux modèles autorégressifs (AR) en permettant une génération de jetons en parallèle. Des recherches récentes ont établi que les dLLM sont d'excellents rédacteurs pour le décodage spéculatif (SD), capables de générer des blocs entiers de jetons en un seul passage direct (forward pass). Cependant, les méthodes de rédaction basées sur la diffusion existantes reposent principalement sur une rédaction linéaire (à chemin unique).

Cette approche linéaire est sous-optimale pour les modèles de diffusion car la génération par diffusion produit naturellement un riche espace combinatoire de jetons candidats à travers de multiples positions simultanément. En restreignant l'exploration à une seule trajectoire, la rédaction linéaire ne parvient pas à exploiter cet espace de candidats, ce qui limite intrinsèquement la longueur d'acceptation et l'efficacité du décodage.

De plus, un décalage fondamental existe entre le score de rédaction par diffusion et la vérification basée sur le préfixe des modèles AR :

Les marginales de diffusion sont aveugles au préfixe : Les rédacteurs de diffusion génèrent des probabilités marginales par position ( $q_i$ ) en parallèle, sans dépendances strictes vis-à-vis du préfixe précédent.
La vérification AR est conditionnelle au préfixe : Le modèle cible vérifie les jetons de manière séquentielle, où la probabilité d'accepter un jeton dépend fortement du préfixe spécifique accepté jusqu'alors.
La conséquence : Appliquer naïvement une rédaction basée sur un arbre (qui repose sur le classement des chemins) aux modèles de diffusion en utilisant les scores marginaux bruts conduit à un classement de chemins peu fiable. Un jeton qui est hautement probable de manière isolée peut devenir peu probable une fois conditionné par un préfixe spécifique, ce qui pousse la construction de l'arbre à prioriser des chemins qui sont ensuite rejetés par le vérificateur.

Méthodologie : PRESTO

Les auteurs proposent PRESTO (Prefix-aligned Scoring and priority-based Tree search for diffusion Speculative decOding), un cadre conçu pour résoudre le décalage entre la rédaction par diffusion et la vérification AR tout en permettant une exploration efficace basée sur des arbres.

1. Score aligné sur le préfixe (Prefix-Aligned Scoring)

Pour remédier au décalage structurel, PRESTO introduit une fonction de score alignée sur le préfixe qui corrige de manière minimale la marginale de la diffusion avec un signal dépendant du préfixe.

Le Score : Pour un jeton $t$ $t$ à la position $d$ $d$ étant donné un préfixe intra-bloc $c_d$ $c_{d}$ , le score est défini par :
$s_{d,t}(c_d) = \log q_d(t) + \lambda_d \log \rho_d(t | c_d)$
Où :
- $q_d(t)$ est la probabilité marginale du rédacteur de diffusion.
- $\rho_d(t | c_d)$ est un signal traçable, conditionné par le préfixe (instancié dans l'article à l'aide d'un modèle $n$ -gramme léger).
- $\lambda_d$ est un hyperparamètre contrôlant la force de la correction du préfixe.
Raisonnement : Cette formulation préserve le signal d'acceptation global calibré du modèle de diffusion ( $q_d$ ) tout en incorporant un terme de correction ( $\rho_d$ ) qui s'aligne sur la nature conditionnée par le préfixe de la vérification AR. La distribution de substitution jointe résultante est alignée sur le préfixe, permettant un classement de chemin précis.

2. Construction d'arbre basée sur la priorité

En utilisant cette nouvelle fonction de score, PRESTO emploie une recherche en arbre basée sur la priorité pour construire l'arbre de rédaction, plutôt qu'une simple sélection linéaire.

Objectif : Maximiser la longueur d'acceptation attendue en sélectionnant les $B$ meilleurs préfixes ayant les scores de chemin cumulés les plus élevés.
Algorithmes : Le cadre prend en charge deux stratégies d'expansion :
- Recherche à priorité maximale (BFS - Best-First Search) : Maintient une file de priorité globale de tous les nœuds de frontière, en développant le nœud unique ayant le score le plus élevé à chaque étape.
- Recherche par faisceau avec rétention globale (Beam Search with Global Retention) : Développe les $W$ meilleurs nœuds à chaque profondeur mais conserve les nœuds hors du faisceau pour la sélection finale.
Implémentation : Les auteurs adoptent la recherche par faisceau comme stratégie par défaut pour l'efficacité pratique, opérant dans un régime limité par la mémoire où les deux stratégies produisent des résultats comparables.

3. Extension aux dLLM auto-spéculatifs

PRESTO est conçu comme un cadre général applicable à la fois aux :

Rédacteurs de diffusion dédiés : Où un modèle de diffusion distinct rédige les jetons pour un modèle AR cible.
dLLM auto-spéculatifs : Où le même modèle assure à la fois la rédaction et la vérification.
Dans le cadre de l'auto-spéculation, (plus précisément la spéculation quadratique), PRESTO s'intègre dans le pipeline "vérifier-et-rédiger". Il réutilise le cache KV pour le préfixe engagé tout en débruitant simultanément les emplacements de jetons masqués en fin de séquence pour le tour suivant de rédaction basée sur l'arbre. Cela élimine le gaspillage de calcul sur les branches rejetées, courant dans la rédaction spéculative linéaire.

Contributions Clés

Identification du décalage : L'article identifie que les probabilités marginales de diffusion sont insuffisantes pour la rédaction basée sur des arbres car elles manquent de conditionnement par préfixe, conduisant à un classement de chemin sous-optimal.
Score aligné sur le préfixe : La conception d'une fonction de score principielle qui combine les marginales de diffusion avec des signaux dépendants du préfixe pour créer une distribution de substitution compatible avec la vérification AR.
Rédaction par diffusion basée sur l'arbre : La formalisation d'un cadre de décodage spéculatif basé sur l'arbre pour les modèles de diffusion qui utilise une expansion basée sur la priorité pour maximiser la longueur d'acceptation.
Généralisabilité : L'extension de ce cadre aux rédacteurs de diffusion dédiés et aux dLLM auto-spéculatifs, en faisant la première solution de rédaction basée sur l'arbre pour ces derniers.

Résultats Expérimentaux

Des expériences approfondies ont été menées sur divers benchmarks (Math, Code, Chat) utilisant des modèles tels que Qwen3 (4B, 8B, 30B) et Nemotron-Labs-Diffusion-8B.

Rédacteurs de diffusion dédiés (dFlash) :
- PRESTO a obtenu une accélération du débit de bout en bout moyenne de 1,5× par rapport aux lignes de base de rédaction linéaire.
- La longueur d'acceptation moyenne a augmenté de manière significative (par exemple, de ~6,5 à ~9,4 sur GSM8K pour Qwen3-4B).
- Les gains ont été constants à travers les échelles de modèles et les températures de décodage (déterministe et stochastique).
dLLM auto-spéculatifs (Nemotron-Labs-Diffusion) :
- Mode Linéaire : A obtenu une accélération du débit moyenne de 1,12× et a amélioré la longueur d'acceptation de 8,8 à 9,9 (déterministe).
- Mode Quadratique : A obtenu une accélération du débit moyenne de 2,14× et a presque doublé la longueur d'acceptation (de 2,4 à 4,7) sous décodage stochastique, là où la rédaction à chemin unique échoue généralement en raison d'un décalage de trajectoire.
Surcharge (Overhead) : La construction de l'arbre et les opérations connexes ont introduit une surcharge négligeable (0,7 % à 4 % de la latence par étape), car la vérification du modèle cible reste le coût dominant.

Signification et Revendications

L'article affirme que PRESTO est le premier cadre à appliquer avec succès la rédaction basée sur l'arbre aux rédacteurs de diffusion en résolvant le décalage fondamental entre le score marginal et la vérification basée sur le préfixe. Sa signification réside dans :

Libérer le potentiel combinatoire : Il exploite pleinement l'espace de candidats multiples naturellement produit par les modèles de diffusion, que les méthodes linéaires laissent inexploré.
Adaptation principielle : Il fournit une méthode théoriquement fondée (score aligné sur le préfixe) pour adapter les modèles de diffusion à la spéculation sans nécess avoir besoin de réentraîner le rédacteur.
Large applicabilité : Il sert de remplacement direct (drop-in replacement) pour la rédaction linéaire dans les pipelines dédiés et auto-spéculatifs, offrant des gains de débit substantiels sans compromettre la nature sans perte (lossless) du décodage spéculatif.

Les auteurs notent que bien que le cadre soit efficace, il repose sur un signal aligné sur le préfixe traçable (actuellement un modèle $n$ -gramme) et que les travaux futurs pourraient explorer des signaux plus riches ou des optimisations à l'échelle de la production dans des frameworks comme vLLM.