Résumé Technique : BASTION

Énoncé du Problème

Les modèles de langage de grande taille (LLM) souffrent de coûts de service élevés en raison de la nature séquentielle du décodage autoregressif (AR), qui nécessite un passage avant distinct du modèle cible pour chaque token généré. Le décodage spéculatif remédie à cela en utilisant un brouilleur léger pour proposer plusieurs tokens candidats, que le modèle cible vérifie en parallèle.

Bien que les récentes avancées en matière de brouilleurs par diffusion par blocs (par exemple, DFlash, TiDAR) permettent la prédiction parallèle de plusieurs tokens futurs en une seule étape, elles introduisent une limitation critique : ces modèles génèrent des marginales positionnelles plutôt que des séquences entièrement conditionnées. Par conséquent, s'engager sur un seul chemin glouton (top-1) aboutit souvent à des séquences incohérentes sous la distribution du modèle cible. Les méthodes existantes de décodage spéculatif basées sur des arbres (par exemple, EAGLE-3) reposent généralement sur des topologies d'arbres statiques et prédéfinies heuristiquement qui ne s'adaptent ni à la requête spécifique ni aux contraintes matérielles, conduisant à des compromis sous-optimaux entre la qualité du brouillage et la surcharge de vérification.

Méthodologie : BASTION

BASTION (Décodage Spéculatif Conscient du Budget avec Brouillage par Diffusion par Blocs en Arbre) est un cadre sans entraînement qui construit dynamiquement des topologies d'arbres dépendantes de la requête pour maximiser l'accélération du décodage. Il intègre trois composants synergiques :

1. Brouillage par Diffusion par Blocs en Arbre

Contrairement aux brouilleurs autoregressifs qui génèrent des tokens séquentiellement, BASTION utilise un brouilleur par diffusion par blocs pour produire des distributions marginales positionnelles $\{q_k(\cdot)\}$ en un seul passage avant.

• Découplage de la Topologie : Le coût de brouillage ($T_{draft}$) est invariant par rapport à la structure de l'arbre (chaînes, arbres larges ou arbres profonds) car le bloc est traité en parallèle.
• Indépendance des Marginales : Les candidats à différentes positions sont indépendants les uns des autres étant donné le contexte validé, ce qui permet de les combiner librement en un arbre de préfixes $\mathcal{T}$.
• Notation des Chemins : Chaque nœud $i$ dans l'arbre à la profondeur $d(i)$ se voit attribuer un score de chemin $\rho(i) = \prod_{k=1}^{d(i)} q_k(x_{ik})$, représentant la probabilité conjointe du chemin sous la distribution du brouilleur.

2. Surrogat d'Acceptation et Expansion Meilleur-D'Abord

Pour déterminer quels nœuds inclure dans l'arbre de vérification sans encourir le coût de la vérification par le modèle cible, BASTION emploie un surrogat de confiance de chemin.

• Estimation par Surrogat : La longueur acceptée attendue $\mathbb{E}[A(\mathcal{T})]$ est approximée par la somme des scores de chemin : $\hat{A}(\mathcal{T}) = \sum_{i \in \mathcal{T}} \rho(i)$. Cela repose sur le postulat que les chemins à haute probabilité du brouilleur s'alignent avec la trajectoire préférée du modèle cible.
• Construction Optimale : L'article démontre qu'une stratégie d'expansion meilleur-d'abord est prouvée optimale pour maximiser $\hat{A}(\mathcal{T})$ sous un budget fixe. L'algorithme ajoute itérativement le nœud valide ayant le $\rho(i)$ le plus élevé, garantissant que l'arbre se développe d'abord le long des trajectoires les plus prometteuses. Cela produit une relation concave entre la taille de l'arbre et l'acceptation attendue, présentant des rendements décroissants.

3. Contrôleur de Budget en Ligne Conscient du Matériel

BASTION détermine dynamiquement la taille optimale de l'arbre $N$ à chaque étape de décodage en équilibrant le gain marginal en longueur d'acceptation contre l'augmentation marginale de la latence de vérification.

• Modélisation de la Latence : Le cadre utilise un modèle de plafond conscient du matériel pour estimer la latence de vérification $T_{verify}(N)$. Ce modèle prend en compte la bande passante mémoire et les contraintes de calcul, calibré via des temps d'exécution observés pour corriger la fusion de noyaux et les frais généraux de planification.
• Arrêt Adaptatif : Le contrôleur étend l'arbre jusqu'à ce que l'accélération estimée $S(N) = \frac{\hat{A}(\mathcal{T}_N) \cdot L_{AR}}{C(N)}$ commence à diminuer. Cela garantit que le système arrête l'expansion exactement lorsque le coût de vérification de nœuds supplémentaires l'emporte sur le bénéfice d'une séquence acceptée plus longue.

Contributions Clés

1. Brouillage par Diffusion par Blocs en Arbre : Introduction d'un algorithme de construction d'arbres de préfixes meilleur-d'abord qui exploite les scores de chemin des logits du brouilleur pour construire des arbres de candidats de haute qualité et fermés par préfixe à partir de logits de diffusion par blocs, dépassant le décodage glouton sur un seul chemin.
2. Contrôleur de Budget en Ligne Conscient du Matériel : Un contrôleur novateur qui détermine dynamiquement la taille optimale de l'arbre à chaque étape en équilibrant les gains d'acceptation attendus contre la latence de vérification calibrée matériellement, éliminant le besoin d'arbres statiques de taille fixe.
3. Validation Empirique Étendue : Évaluation complète sur divers modèles cibles (Qwen3, Llama-3.1), benchmarks (mathématiques, code, chat) et architectures GPU (A100, H100, A6000, RTX PRO 6000), démontrant des accélérations significatives sans ajustement par configuration.

Résultats Expérimentaux

Évalué sur un GPU NVIDIA A100 avec décodage glouton (température 0) :

• Accélération : BASTION réalise une accélération moyenne de 6,61× par rapport au décodage autoregressif standard sur Qwen3-8B.
• Comparaison avec les Baselines :
  • Surpasse EAGLE-3 (décodage par arbre basé sur AR état de l'art) de 2,45× (6,61× contre 2,70× en moyenne).
  • Surpasse DFlash (diffusion par blocs sur un seul chemin état de l'art) de 1,39× (6,61× contre 4,76× en moyenne).
• Longueur d'Acceptation : BASTION atteint des Longueurs Moyennes Acceptées (AAL) significativement plus grandes que les baselines sur un seul chemin (par exemple, 8,48 contre 6,24 sur Qwen3-8B), confirmant que la vérification structurée en arbre capture avec succès des continuations valides plus longues.
• Généralisation : Les gains de performance sont cohérents à travers différentes générations de GPU (y compris les architectures H100 et Blackwell) et familles de modèles, indiquant que la méthode n'est pas liée à des configurations matérielles ou de modèles spécifiques.

Importance et Revendications

L'article revendique que BASTION représente une avancée significative dans le décodage spéculatif en résolvant la tension entre l'efficacité parallèle des brouilleurs par diffusion par blocs et la cohérence séquentielle requise par les modèles cibles. En remplaçant les formes d'arbres statiques et heuristiques par une construction dynamique et consciente du budget, BASTION utilise pleinement les ressources de calcul parallèle disponibles tout en évitant les pénalités de latence d'une sur-vérification des chemins à faible probabilité.

Crucialement, la méthode est sans entraînement et ne nécessite aucun ajustement par configuration, la rendant immédiatement déployable sur diverses configurations matérielles et de modèles. Les auteurs affirment que leur approche comble efficacement le fossé entre le potentiel théorique du brouillage par diffusion par blocs et l'inférence pratique à haut débit, réalisant des accélérations quasi optimales en alignant adaptativement la structure du brouillage avec les contraintes système en temps réel.