🤖 machine learning

Trees from Marginals: Autoregressive drafting with factorized priors

Le document présente Weaver, un adaptateur autorégressif léger qui reconstruit les dépendances conditionnelles à partir de marginales de brouillon factorisées pour permettre un décodage spéculatif basé sur des arbres efficace, atteignant une accélération de 4,37 fois par rapport au décodage autorégressif standard grâce à un nouvel algorithme de vérification sans retour en arrière et des noyaux CUDA optimisés.

Auteurs originaux : Yuma Oda, Ryan Mathieu, Roman Knyazhitskiy, Artur Chakhvadze

Publié 2026-07-09

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Yuma Oda, Ryan Mathieu, Roman Knyazhitskiy, Artur Chakhvadze

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 d'écrire une histoire avec un bibliothécaire très intelligent, mais lent (le modèle IA). Chaque fois que vous lui demandez le mot suivant de l'histoire, le bibliothécaire doit s'arrêter, réfléchir intensément, consulter toute sa bibliothèque de livres, puis vous chuchoter le mot suivant. C'est ainsi que fonctionne l'IA actuelle : un mot à la fois, une étape à la fois. C'est précis, mais c'est lent.

Cette publication présente une nouvelle façon de rendre ce bibliothécaire beaucoup plus rapide sans perdre en précision. Ils appellent leur méthode « Trees from Marginals » (ou DFlash-TfM). Voici comment cela fonctionne, décomposé en analogies simples.

Le Problème : La limite du « Jeu de Devinettes »

Pour accélérer les choses, les chercheurs ont inventé une astuce appelée Décodage Spéculatif (Speculative Decoding).

L'ancienne méthode : Un assistant rapide, un « rédacteur » (le drafter), devine les prochains mots. Ensuite, le bibliothèreur lent (le « vérificateur ») vérifie si ces devinettes sont correctes. S'il les accepte, il les valide tous d'un coup. Sinon, le bibliothécaire corrige l'erreur et recommence.
Le problème des rédacteurs « factorisés » : Certains assistants sont super rapides car ils devinent tous les mots suivants d'un coup, en ignorant la façon dont ils sont connectés entre eux. C'est comme un chef qui devinerait les trois prochains ingrédients d'une soupe sans avoir goûté les précédents.
- Le piège : À mesure que la liste de devinettes s'allonge, le chef devient moins bon pour deviner la séquence. La première devinette peut être juste, mais la troisième est généralement fausse car elle n'a pas tenu compte des deux premières. Cela limite le nombre de mots pouvant être acceptés à la fois.

La Solution : L'assistant « Tisserand »

Les auteurs ont créé un nouveau système qui combine la vitesse du chef rapide avec la logique d'un éditeur méticuleux. Ils appellent ce nouvel éditeur Weaver.

La « Liste Courte Top-K » : D'abord, l'assistant rapide (DFlash) fait une devinette rapide et approximative et fournit une liste courte des 512 mots les plus probables pour l'emplacement suivant. C'est comme un chef disant : « Je pense que le prochain ingrédient fait probablement partie de ces 512 épices. »
Le travail de Weaver : Au lieu de deviner aveuglément, le Weaver (une IA minuscule et légère) examine cette liste courte. Il agit comme un éditeur intelligent qui dit : « D'accord, si le premier mot était 'sel', alors le mot suivant est presque certainement 'poivre', et non 'sucre'. »
Construire un Arbre : Le Weaver ne se contente pas de faire une suite de devinettes en ligne droite. Il construit un arbre.
- Imaginez un arbre généalogique. La racine est la phrase actuelle.
- Le Weaver crée des ramifications, créant différents chemins possibles pour l'histoire (par exemple, « Le chat s'est assis sur le tapis » contre « Le chat s'est assis sur le sol »).
- Parce que le Weaver est petit et ne regarde que la liste courte fournie par l'assistant rapide, il est incroyablement rapide pour construire cet arbre de possibilités.

La Vérification : Vérifier l'Arbre

Maintenant, le bibliothécaire lent doit vérifier cet arbre de devinettes.

L'ancien problème : Si le bibliothécaire utilise un système de mémoire « récurrent » standard (comme les couches Gated Delta Net dans l'IA moderne), vérifier un arbre est généralement un cauchemar. C'est comme essayer de parcourir chaque branche d'un arbre une par une pour voir quel chemin est le bon. C'est lent.
La nouvelle astuce : Les auteurs ont inventé un raccourci mathématique spécial (un algorithme « sans retour en arrière » ou rollback-free).
- Au lieu de parcourir chaque branche, ils utilisent une résolution triangulaire masquée (masked triangular solve). Considérez cela comme une carte magique qui permet au bibliothécaire de regarder la structure entière de l'arbre d'un seul coup et de savoir instantanément quel chemin est le bon, sans avoir à recalculer l'état de la mémoire pour chaque branche.
- C'est comme avoir un GPS qui met instantanément en évidence l'itinéraire correct sur une carte complexe sans que vous ayez besoin de parcourir chaque impasse au préalable.

Le Résultat : Vitesse et Efficacité

En combinant ces idées, le système obtient deux victoires majeures :

Plus de mots acceptés : Parce que le Weaver corrige les erreurs de logique de l'assistant rapide, le bibliothécaire accepte des chaînes de mots plus longues (jusqu'à 77 % de plus que la méthode précédente la plus performante).
Accélération massive : Tout le processus est si efficace que l'IA génère du texte 4,37 fois plus vite que la méthode lente standard. Elle bat également la méthode précédemment la plus rapide d'environ 25 %.

Analogie de Synthèse

IA Standard : Un escargot écrivant une histoire, une lettre à la fois, en vérifiant chaque lettre dans un dictionnaire.
Ancienne méthode rapide : Un lecteur rapide devinant tout le paragraphe suivant, mais se trompant souvent au milieu du paragraphe parce qu'il n'a pas prêté attention au début.
Cette nouvelle méthode (Weaver) : Un lecteur rapide qui choisit rapidement les 500 mots les plus probables qui pourraient convenir, et un minuscule éditeur super intelligent (Weaver) qui organise instantanément ces mots en un arbre de phrases logiques. Une « carte magique » (le nouveau noyau ou kernel) vérifie ensuite tout l'arbre instantanément pour voir quel chemin est réel.

Le résultat est une IA qui écrit aussi vite qu'un lecteur rapide, mais avec la précision d'un éditeur méticuleux, rendant les interactions beaucoup plus instantanées et réactives.

Résumé Technique : Arbres à partir de Marginales (DFlash-TfM)

Énoncé du Problème

Les modèles de langage de grande taille (LLM) autorégressifs sont intrinsèquement séquentiels, générant un jeton par étape. Cette nature séquentielle crée un goulot d'étranglement lié à la mémoire où la latence est dominée par la bande passante de la DRAM plutôt que par le calcul, ce qui entraîne une faible efficacité énergétique et une interactivité limitée (jetons décodés par seconde par séquence). Bien que le décodage spéculatif offre une solution en utilisant un modèle « drafter » (rédacteur) rapide pour proposer plusieurs jetons pour une vérification parallèle par le modèle « verifier » (vérificateur) principal, les approches existantes font face à des compromis :

Drafters Factorisés : Des modèles comme DFlash prédisent les jetons futurs comme des marginales conditionnellement indépendantes en parallèle. Bien qu'efficaces, leur hypothèse d'indépendance provoque une dégradation rapide des taux d'acceptation à mesure que le budget spéculatif (longueur du draft) augmente, car ils ne parviennent pas à capturer les dépendances conditionnelles entre les jetons proposés.
Drafters Autorégressifs : Ceux-ci capturent correctement les dépendances mais induisent un coût élevé dû aux projections sur l'ensemble du vocabulaire à chaque étape, annulant les gains de vitesse.
Vérification sur Modèles Récurrents : Les algorithmes standards de vérification d'arbres reposent sur des masques d'attention. Cependant, les modèles modernes utilisant des couches Gated Delta Net (GDN) (par exemple, Qwen3.6) utilisent des transitions d'état non commutatives qui ne supportent pas de simple masquage restreint aux ancêtres, faisant de la vérification efficace d'arbres en un seul passage un problème ouvert.

Méthodologie : DFlash-TfM

Les auteurs proposent DFlash-TfM (Trees from Marginals), une architecture hybride qui combine l'efficacité parallèle des drafters factorisés avec la précision conditionnelle des modèles autorégressifs, ainsi qu'un nouveau noyau de vérification au niveau système.

1. L'Adaptateur Weaver

L'innovation centrale est Weaver, un adaptateur autorégressif léger (56,7 millions de paramètres) qui opère sur un vocabulaire contraint.

Entrée : Weaver prend les états cachés d'un drafter factorisé pré-entraîné (DFlash) et du verifier.
Mécanisme : Au lieu de projeter vers le vocabulaire complet, Weaver extrait les jetons marginaux Top- $K$ prédits par DFlash. Il prédit ensuite de manière autorégressive la distribution résiduelle sur ce petit ensemble de candidats ( $K=512$ ).
Avantage : Cela restaure les dépendances conditionnelles entre les jetons (par exemple, le jeton $t+1$ dépend du jeton réalisé $t$ ) sans le coût de bande passante mémoire d'une projection sur le vocabulaire complet. Le modèle construit un arbre de proposition où les branches sont échantillonnées à partir de ces résidus conditionnés.

2. Construction d'Arbre

Les auteurs modifient l'algorithme de construction d'arbre DySpec pour améliorer le débit. Au lieu d'une expansion strictement séquentielle de type "best-first", ils extraient les $w$ nœuds supérieurs de la file d'attente des candidats et effectuent $w$ expansions simultanément. Cela équilibre la bande passante mémoire et la pression de calcul, permettant la construction de l'arbre de draft complet en $\lceil B/w \rceil$ opérations séquentielles.

3. Vérification d'Arbre sans Rollback pour GDN

Pour supporter les modèles avec des couches Gated Delta Net, les auteurs dérivent un algorithme de vérification d'arbre sans rollback :

Défi : Les couches GDN mettent à jour l'état via $S_t = \alpha_t(I - \beta_t k_t k_t^\top)S_{t-1} + \beta_t k_t v_t^\top$ . La nature non commutative de la mise à jour de la matrice empêche les optimisations classiques de produit cumulatif utilisées dans les modèles d'espace d'état diagonaux.
Solution : Les auteurs utilisent une forme par blocs (chunked form) de la récurrence linéaire. Ils formulent la vérification comme une résolution triangulaire masquée sur la structure de l'arbre.
- L'arbre est trié topologiquement.
- L'algorithme calcule les matrices d'interaction ( $X$ et $Y$ ) représentant les relations d'ancêtre.
- Il résout un système linéaire $(I+X)U = \beta V$ pour calculer les sorties sans mettre à jour spéculativement l'état récurrent.
- L'engagement de l'état est retardé jusqu'à l'itération de décodage suivante, une fois que la branche acceptée est connue. Cela élimine le besoin de rollback (retour en arrière) de l'état.
Implémentation : Ceci est implémenté sous la forme d'un noyau CUDA fusionné dans SGLang, qui calcule la matrice de Gram, résout le système via une substitution directe par blocs (tiled forward substitution) et applique l'étape de sortie en un seul passage.

Contributions Clés

Architecture Hybride : Introduction de DFlash-TfM, qui utilise un léger adaptateur autorégressif (Weaver) pour conditionner les marginales factorisées, levant le plafond du taux d'acceptation des drafters purement factorisés à de longues profondeurs.
Insight Théorique : Analyse montrant que le taux d'acceptation des drafters factorisés est limité par la distance de Variation Totale (TV) entre les distributions marginales et conditionnelles, et que le conditionnement sur les jetons réalisés (via Weaver) permet au modèle hybride de dépasser les limites théoriques de tout drafter uniquement basé sur les marginales.
Optimisation Système : Dérivation d'un nouvel algorithme de vérification d'arbre sans rollback pour les couches Gated Delta Net, permettant le décodage spéculatif efficace pour des modèles comme Qwen3.6 qui étaient auparavant difficiles à optimiser pour les méthodes basées sur les arbres.
Implémentation : Noyaux CUDA optimisés pour SGLang qui fusionnent les étapes de vérification, réduisant considérablement la latence par rapport aux bases récurrentes par branche.

Résultats Expérimentaux

Évalué sur Qwen3.6-27B (bfloat16) utilisant un GPU NVIDIA B200 sur des benchmarks de chat, de mathématiques et de code :

Accélération : DFlash-TfM atteint une accélération de 4,37× par rapport au décodage autorégressif standard.
Comparaison avec les Baselines : Il surpasse une baseline DFlash hautement optimisée de 24,7 % et une baseline DDTree de 10 à 30 % selon le jeu de données.
Longueur d'Acceptation : La méthode augmente la longueur moyenne d'acceptation (MAL) de 77 % par rapport à la chaîne DFlash et de 32 % par rapport à DDTree à taille de arbre égale.
Efficacité de Vérification : Le noyau fusionné passe à l'échelle bien mieux que la vérification récurrente par branche. Pour une taille d'arbre de 128 jetons, le noyau fusionné est 7,1× plus rapide que la baseline récurrente. L'étape de vérification ne représente qu'environ 12 % du temps de décodage total.

Signification et Revendications

L'article affirme que la limitation principale des drafters factorisés est structurelle (hypothèse d'indépendance) plutôt qu'un manque de capacité du modèle. En introduisant un adaptateur autorégressif léger qui opère sur un vocabulaire contraint, les auteurs démontent qu'il est possible de restaurer les dépendances conditionnelles et d'atteindre des taux d'acceptation élevés sans subir l'overhead des projections sur le vocabulaire complet.

De plus, ce travail répond à une lacune critique du système : permettre le décodage spéculatif efficace pour les modèles avec des couches récurrentes non-diagonales (GDN). Le noyau de vérification sans rollback proposé élimine le besoin de rollback d'état, rendant le décodage spéculatif basé sur les arbres viable pour une classe plus large d'architectures modernes. Les auteurs concluent que ces contributions combinées au modèle et au système représentent une amélioration de l'état de l'art pour l'interactivité de l'inférence LLM locale ou côté serveur.