TreeFlash: Parallel AR-Approximation for Faster Speculative Decoding

Résumé Technique : TreeFlash

Énoncé du Problème

Le décodage spéculatif accélère l'inférence des modèles de langage de grande taille (LLM) en utilisant un rédacteur léger pour proposer des jetons candidats, qui sont ensuite vérifiés en parallèle par un modèle cible plus important. Bien que le rédaction basée sur des arbres (tree-based drafting) ait amélioré les taux d'acceptation en explorant simultanément plusieurs branches de candidats, une limitation fondamentale persiste dans les rédacteurs à bloc par tirage unique (one-shot block drafters) (ex. : DFlash).

Ces rédacteurs à tirage unique génèrent l'intégralité du bloc de rédaction en une seule passe avant (forward pass), mais conditionnent chaque jeton prédit $x_{t+i}$ uniquement sur le contexte du préfixe $x_{\le t}$. Ils manquent de dépendance vis-à-vis des jetons précédemment rédigés au sein du bloc ($x_{t+i-1}$). Cette conditionnalité non autorégressive provoque une divergence significative entre la distribution du rédacteur et la véritable distribution autorégressive du modèle cible à mesure que la profondeur de la rédaction augmente. Dans les contextes basés sur des arbres, ce problème est amplifié car les différentes branches partageant un préfixe commun sont contraintes de partager la même distribution marginale pour les jetons suivants, ce qui dégrade la qualité de l'arbre de rédaction et limite les gains de débit.

Méthodologie : TreeFlash

TreeFlash traite cette divergence en incorporant un mécanisme d'Approximation Autorégressive (AR) légère dans le paradigme de rédaction à tirage unique, préservant ainsi sa complexité de décodage en $O(1)$.

1. Couche d'Approximation AR

L'innovation centrale est une couche SwiGLU légère qui modifie l'état caché du rédacteur pour incorporer des informations provenant du jeton immédiatement précédent. Pour une position de rédaction $t+i$, l'état caché modifié $h'_{t+i}$ est calculé comme suit :
$$h'_{t+i} = h_{t+i} + \text{SwiGLU}(\tilde{h}_{t+i} :: e_{t+i-1})$$
où $\tilde{h}_{t+i}$ est l'état caché normalisé du rédacteur de base, et $e_{t+i-1}$ est l'embedding d'entrée normalisé du jeton précédent. Cela permet au modèle d'approximer une distribution autorégressive $q'(x_{t+i} | x_{\le t}, x_{t+i-1})$ sans génération séquentielle.

2. Construction d'Inférence en Deux Étapes

L'application naïve de l'approximateur AR de manière séquentielle annulerait les avantages d'efficacité de la rédaction à tirage unique. TreeFlash emploie une construction en deux étapes pour maintenir le parallélisme :

• Étape 1 : La distribution DFlash originale est utilisée pour construire un arbre $M$-aire supérieur (où $M$ est un hyperparamètre contrôlant le facteur de branchement). Cet arbre contient exactement $M$ jetons uniques à chaque profondeur.
• Étape 2 : L'approximateur AR est appliqué aux nœuds de cet arbre $M$-aire pour obtenir des distributions de jetons modifiées. Puisque l'arbre possède une largeur fixe $M$, toutes les distributions modifiées peuvent être calculées en $M \cdot \gamma$ évaluations parallèles, quel que soit le budget de rédaction total $B$.
• Sélection Finale : L'arbre de rédaction final est construit via l'algorithme OPT-Tree en utilisant ces distributions modifiées.

3. Stratégie d'Entraînement

• Objectif : Le modèle est entraîné en utilisant la divergence KL directe (forward KL-divergence) pour aligner le rédacteur avec la distribution du vérificateur, plutôt que la divergence croisée (cross-entropy) standard.
• Initialisation : Le squelette (backbone) DFlash est initialisé à partir d'un point de contrôle pré-entraîné, tandis que l'approximateur AR est initialisé à zéro pour garantir que le modèle commence comme une copie fidèle de DFlash.
• Entrée : Pendant l'entraînement, l'approximateur AR est conditionné sur les jetons précédents de vérité terrain (teacher forcing).
• Mise à l'échelle de la perte : Un schéma de mise à l'échelle de la perte est adopté pour pénaliser plus lourdement les erreurs dans les jetons précoces.

Principales Contributions

1. Identification du Goulot d'Étranglement : L'article identifie l'absence de conditionnalité autorégressive dans les rédacteurs à bloc par tirage unique comme un goulot d'étranglement critique, particulièrement dans les contextes d'arbres où les distributions marginales ne parviennent pas à capturer la cohérence des jetons à travers les branches profondes.
2. Approximation AR Légère : L'introduction d'un mécanisme qui approxime la conditionnalité autorégressive tout en conservant la complexité de décodage $O(1)$ des rédacteurs entièrement parallèles.
3. Gains de Performance : Démonstration que TreeFlash surpasse systématiquement les rédacteurs entièrement séquentiels (EAGLE-3) et les rédacteurs d'arbres marginaux de pointe (DFlash, DDTree) sur diverses tâches et tailles de modèles.

Résultats Expérimentaux

Les expériences ont été menées sur Qwen3 4B, Qwen3 8B et Qwen3 Coder 30B A3B à travers le raisonnement mathématique (GSM8K, MATH-500), la génération de code (HumanEval, MBPP) et le suivi d'instructions (MT-Bench).

• Accélération et Efficacité : Avec un budget de rédaction identique ($B=16$), TreeFlash améliore l'efficacité du bloc ($\tau$) de 12,4 % et l'accélération de 9,1 % par rapport à DDTree.
• Scalabilité : Les bénéfices de l'approximation AR croissent avec le budget de rédaction. À $B=64$, TreeFlash atteint une amélioration de l'efficacité du bloc de 12,6 % et de l'accélération de 9,0 % par rapport à DDTree.
• Métriques Distributionnelles :
  • Distance de Variation Totale (TVD) : TreeFlash maintient une TVD nettement plus faible que la distribution du vérificateur à mesure que la profondeur augmente. À la profondeur 15, la TVD de TreeFlash est de 0,62, contre 0,81 pour DFlash. Notamment, au-delà de la profondeur 9, la distribution de TreeFlash est mieux calibrée que la distribution marginale de vérité terrain elle-même.
  • Couverture Top-K : La couverture top-1 de TreeFlash à la profondeur 15 (0,45) correspond à la couverture top-5 de DFlash (0,44), indiquant qu'un seul candidat TreeFlash atteint la même probabilité d'acceptation que cinq candidats DFlash.
• Études d'Ablation :
  • La couche SwiGLU est essentielle ; un approximateur linéaire est moins performant.
  • Le conditionnement sur un seul jeton précédent est suffisant ; l'utilisation de deux jetons précédents n'apporte que des gains négligeables.
  • Le gel du squelette DFlash pendant l'entraînement donne des performances presque identiques au réglage fin (fine-tuning) complet, suggérant que les gains sont principalement pilotés par l'approximateur AR.

Signification et Revendications

L'article affirme que TreeFlash représente une avancée significative dans le décodage spéculatif en comblant le fossé entre l'efficacité de la rédaction parallèle à tirage unique et la cohérence de la génération autorégressive. En traitant la divergence distributionnelle qui affecte les arbres de rédaction profonds, TreeFlash permet une utilisation plus efficace du budget de rédaction, conduisant à un débit plus élevé sans compromettre la qualité de la génération. Les auteurs soulignent que si la transition vers la rédaction basée sur des arbres fournit le plus gros gain initial, l'approximation AR fournit un gain constant et complémentaire qui est plus prononcé dans les régimes à haute acceptation et avec des budgets de rédaction plus importants.

Limites Notées :

• Les expériences sont limitées à des contextes de batch unique avec l'attention SDPA ; les interactions avec les optimisations de niveau production (quantification, multi-batch) ne sont pas évaluées.
• Le modèle est initialisé à partir d'un point de contrôle DFlash existant plutôt que d'être entraîné à partir de zéro.
• L'approximateur AR est entraîné avec le teacher forcing, ce qui peut introduire un biais d'exposition lors de l'inférence.
• L'évaluation est actuellement restreinte à la famille de modèles Qwen.