Cost-Aware Diffusion Draft Trees for Speculative Decoding

Résumé Technique : Arbres de Diffusion de Drafts Sensibles au Coût pour le Décodage Spéculatif

Énoncé du Problème
Le décodage spéculatif accélère l'inférence des modèles de langage autorégressifs en utilisant un modèle de draft (drafter) léger pour proposer des jetons candidats, qui sont ensuite vérifiés en parallèle par le modèle cible. Bien que la vérification basée sur des arbres permette la vérification simultanée de multiples continuations candidates, les méthodes existantes pour construire ces arbres de draft font face à un désalignement fondamental dans leurs objectifs d'optimisation.

Les approches actuelles, telles que DDTree, construisent des arbres de draft pour maximiser la longueur d'acceptation attendue sous un budget de nœuds fixe. Cependant, les auteurs observent que la longueur d'acceptation est non décroissante par rapport au budget ; des arbres plus grands produisent toujours des longueurs d'acceptation plus longues, n'offrant aucune base principled pour sélectionner une taille de budget optimale. De plus, le budget optimal n'est pas statique ; il varie considérablement d'un tour de décodage à l'autre en fonction de la confiance du drafter et de la longueur du contexte actuel. Les méthodes existantes qui adaptent la forme de l'arbre par étape reposent toujours sur un budget total de nœuds fixe et sélectionné hors ligne, échouant à répondre à la variabilité par tour entre la couverture et le coût de vérification.

Méthodologie : CaDDTree
Le papier introduit CaDDTree (Cost-aware Diffusion Draft Tree), un cadre qui optimise directement le débit de jetons (jetons générés par unité de temps attendus) plutôt que la longueur d'acceptation. La méthode opère sur des drafters de diffusion par blocs (ex: DFlash), qui génèrent des distributions marginales par position en un seul passage direct (forward pass).

1. Objectif de Débit : Les auteurs modélisent le coût par tour comme $C(T; \ell) = C_d + C_v(|T|; \ell)$, où $C_d$ est le coût de rédaction et $C_v$ est le coût de vérification dépendant de la taille de l'arbre $|T|$ et de la longueur du contexte $\ell$. L'objectif est de maximiser le ratio des jetons acceptés attendus ($1 + \Phi(T)$) par rapport au coût total.
2. Décomposition : Les auteurs prouvent que la maximisation du débit sur tous les arbres valides se réduit à une recherche unidimensionnelle sur le budget de nœuds $n$. Pour tout $n$ fixé, la structure d'arbre optimale est l'arbre "greedy best-first" (sélectionnant les nœuds avec les probabilités de préfixe les plus élevées), ce qui est déjà la construction utilisée par DDTree.
3. Unimodalité et Arrêt Glouton (Greedy Stopping) : Une contribution théorique clé est la preuve que, sous une condition de convexité légère sur la fonction de coût de vérification, l'objectif de débit de substitution est unimodal par rapport à la taille de l'arbre. Cette propriété permet un algorithme glouton efficace : l'algorithme ajoute incrémentalement des nœuds à l'arbre et s'arrête dès que le gain marginal de débit devient non positif.
4. Budget Adaptatif : Contrairement aux travaux précédents, CaDDTree ne nécessite pas de recherche par grille hors ligne pour un budget fixe. Au lieu de cela, il adapte le budget $n^*$ pour chaque tour de décodage en fonction de la distribution par position actuelle et de l'estimation du coût de vérification (qui varie avec la longueur du contexte).
5. Profilage des Coûts : La fonction de coût de vérification $C_v(n; \ell)$ est estimée une seule fois hors ligne via une régression convexe sur les latences mesurées à travers diverses longueurs de contexte et nombres de nœuds. Cela garantit que le modèle de coût reste convexe et monotone, satisfaisant la garantie d'unimodalité.

Contributions Clés

• Changement d'Objectif : Propose l'optimisation directe du débit comme critère principal pour la construction de l'arbre de draft, remplaçant la maximisation de la longueur d'acceptation mal alignée.
• Garantie Théorique : Prouve que l'objectif de débit est unimodal sous des coûts de vérification convexes, permettant une règle d'arrêt gloutonne prouvablement optimale sans entraînement supplémentaire.
• Mécanisme Adaptatif : Introduit une méthode qui sélectionne dynamiquement la taille optimale de l'arbre par tour, éliminant ainsi le besoin de réglage manuel ou hors ligne du budget.
• Implémentation : Instancie le cadre en utilisant des drafters de diffusion par blocs, exploitant leurs distributions marginales en un seul passage pour construire les arbres efficacement.

Résultats Expérimentaux
La méthode a été évaluée sur Qwen3-4B et Qwen3-8B à travers huit benchmarks couvrant le raisonnement (MATH-500, GSM8K, AIME 2025), le codage (HumanEval, MBPP, LiveCodeBench) et le suivi d'instructions (MT-Bench, Alpaca).

• Performance : CaDDTree égale ou surpasse DDTree avec sélection de budget oracle (le meilleur budget fixe trouvé via recherche par grille) sur presque toutes les tâches.
• Efficacité : La méthode atteint des temps de génération par jeton plus faibles (débit plus élevé) par rapport aux baselines autorégressives et au DDTree à budget fixe. Par exemple, sur Qwen3-8B avec MATH-500, CaDDTree atteint un temps par jeton de 4,60 ms comparé à 5,83 ms pour la baseline autorégressive.
• Adaptativité : L'analyse montre que le budget optimal varie considérablement par tour. Sur des tâches de raisonnement comme MATH-500, le budget sélectionné fluctue largement (écart-type $\approx 42$) selon la confiance du drafter, s'écartant souvent significativement du budget oracle fixe. Sur les tâches de suivi d'instructions, le budget est plus stable mais reste distinct des points de grille fixes.
• Robustesse : Les études d'ablation indiquent que la méthode est robuste aux perturbations des estimations de paramètres de coût, suggérant qu'un étalonnage précis hors ligne n'est pas strictement nécessaire.

Signification
Le papier affirme que CaDDTree fournit une approche rigoureuse et sensible au coût pour le décodage spéculatif, résolvant la tension entre la couverture de l'arbre et la latence de vérification. En prouvant l'unimodalité de l'objectif de débit, les auteurs permettent un algorithme efficace et sans entraînement qui s'adapte automatiquement à la difficulté spécifique de chaque étape de décodage. Cela élimine le besoin de sélection heuristique du budget et démontre qu'une construction d'arbre dynamique et sensible au coût peut systématiquement surpasser les baselines statiques optimisées par oracle à travers diverses tâches et tailles de modèles.