Multi-Drafter Speculative Decoding with Alignment Feedback

1. Contexte et Problématique

Le Décodage Spéculatif (SD) :
Le décodage spéculatif est une technique clé pour accélérer l'inférence des grands modèles de langage (LLM). Elle utilise un modèle plus petit (le "brouilleur" ou drafter) pour prédire une séquence de jetons futurs, qui sont ensuite vérifiés en parallèle par le modèle cible (target LLM). Seuls les jetons alignés avec les prédictions du modèle cible sont acceptés, ce qui préserve la qualité de génération tout en réduisant la latence.

Limites des approches actuelles :
La plupart des méthodes SD existantes reposent sur un seul brouilleur. Cette approche présente des limitations majeures :

• Manque de généralisation : Un brouilleur entraîné pour une tâche spécifique (ex: code) performe mal sur d'autres tâches (ex: traduction ou raisonnement mathématique).
• Biais d'entraînement : Les brouilleurs spécialisés souffrent de biais inhérents à leurs données d'entraînement, ce qui réduit leur taux d'acceptation sur des entrées hors distribution (OOD) ou des requêtes dynamiques.
• Rigidité : Les systèmes statiques ou les ensembles fixes de brouilleurs ne s'adaptent pas aux changements de contexte au sein d'une même génération.

Le défi :
Comment intégrer dynamiquement plusieurs brouilleurs hétérogènes dans un pipeline SD pour maximiser le taux d'acceptation des jetons, tout en minimisant la surcharge computationnelle et en s'adaptant aux variations de la distribution des requêtes ?

2. Méthodologie : METASD

Les auteurs proposent METASD, un cadre unifié qui traite la sélection de brouilleurs comme un problème d'apprentissage en ligne via des Bandits Multi-Armes (MAB).

A. Formulation du Problème

Le processus est modélisé comme un problème de bandit où :

• Les bras (Arms) : Chaque brouilleur disponible dans une pool hétérogène.
• La round (Tour) : Chaque cycle de décodage spéculatif (brouillage, vérification, acceptation).
• L'objectif : Maximiser le nombre de jetons acceptés par tour, ce qui équivaut à minimiser le temps d'arrêt (nombre total de tours nécessaires pour générer une séquence de longueur $B$).

Les auteurs définissent une nouvelle métrique de regret basée sur le temps d'arrêt (Stopping Time Regret) :
$$REG(\pi, B) = E[\tau(\pi, B)] - E[\tau(\pi^*, B)]$$
Où $\tau$ est le nombre de tours nécessaires pour générer $B$ jetons. Minimiser ce regret équivaut à maximiser l'efficacité du décodage.

B. Rétroaction d'Alignement et Récompense (Reward)

Le cœur de METASD réside dans la définition d'une récompense basée sur l'alignement entre le brouilleur et le modèle cible. Deux signaux sont comparés :

1. Block Efficiency (BE) : Le nombre moyen de jetons acceptés.
2. Block Divergence (BD) : La mesure de la distance de variation totale (TV distance) entre les distributions de probabilité du brouilleur et du modèle cible sur un bloc de jetons.

Choix de la récompense BD :
Les auteurs démontrent théoriquement et empiriquement que la récompense BD est supérieure à la BE.

• Signal plus informatif : La BD offre une variance plus faible et une meilleure discrimination entre les brouilleurs optimaux et sous-optimaux.
• Non-sparse : Contrairement à la BE qui peut être nulle (si aucun jeton n'est accepté), la BD fournit toujours une rétroaction non nulle, permettant un apprentissage continu.
• Théorème 1 : La BD réduit le "feedback signal" nécessaire pour identifier le meilleur bras, conduisant à une convergence plus rapide.

C. Algorithme : MetaSD-UCB

L'implémentation principale utilise l'algorithme Upper Confidence Bound (UCB) adapté au contexte SD.

• Exploration vs Exploitation : L'algorithme sélectionne le brouilleur $i$ qui maximise la somme de la récompense moyenne estimée et d'un terme d'incertitude (basé sur le nombre de fois où le brouilleur a été choisi).
• Adaptabilité : Le système réinitialise le bandit pour chaque nouvelle requête (gestion de la non-stationnarité inter-requêtes) et peut être étendu à des algorithmes de bandits non stationnaires (ex: Discounted-UCB) pour gérer les changements de contexte au sein d'une même génération.
• Gestion du Cache KV : Pour minimiser le coût de commutation entre brouilleurs, les poids sont préchargés en DRAM, et le recalcul du cache KV est optimisé (seule une couche est recalculée pour les brouilleurs de type Eagle).

3. Contributions Clés

1. Framework Unifié (METASD) : Première approche intégrant dynamiquement plusieurs brouilleurs hétérogènes dans le SD, fonctionnant aussi bien en configuration "boîte noire" (brouilleurs indépendants) que "boîte blanche" (brouilleurs utilisant les représentations latentes du modèle cible, ex: Eagle).
2. Analyse Théorique :
   • Établissement de bornes de regret logarithmiques pour MetaSD-UCB.
   • Démonstration que l'utilisation de la récompense Block Divergence (BD) offre une borne de regret plus serrée que la Block Efficiency (BE).
   • Preuve que minimiser le regret de temps d'arrêt est équivalent à maximiser le nombre de jetons acceptés.
3. Performance Supérieure : Résultats expérimentaux montrant que METASD surpasse systématiquement les approches à brouilleur unique, les ensembles statiques et les méthodes de routage basées sur la classification (MoE).

4. Résultats Expérimentaux

Les expériences ont été menées sur des tâches variées (Code, Traduction, Résumé, QA, Math) et des langues multiples (Japonais, Russe, Allemand, Français, Chinois vers l'Anglais), utilisant des modèles cibles comme Vicuna-7B et des brouilleurs spécialisés.

• Accélération (Speedup) :
  • Boîte Noire (MetaSpS) : METASD-UCB atteint des accélérations allant de 1.5x à 2.3x par rapport au décodage autorégressif standard, surpassant souvent le meilleur brouilleur spécialisé statique.
  • Boîte Blanche (MetaEagle) : Avec les brouilleurs Eagle, les gains sont encore plus significatifs, atteignant jusqu'à 3.7x sur certaines tâches, dépassant les méthodes de pointe comme Medusa ou Rescored-BPD.
• Robustesse :
  • Requêtes Perturbées : METASD maintient une performance élevée même lorsque les prompts sont reformulés par un LLM, là où les routeurs statiques (basés sur des embeddings) échouent.
  • Hors Domaine (OOD) : Sur des tâches non vues pendant l'entraînement des brouilleurs spécialisés (ex: Finance, RAG), METASD généralise mieux qu'un brouilleur "One-Size-Fits-All" (OFA).
• Efficacité Mémoire : L'approche n'augmente pas la bande passante VRAM car un seul brouilleur est actif à la fois. L'augmentation de l'utilisation DRAM est minime (~2 Go pour un modèle 7B).

5. Signification et Impact

Ce travail représente une avancée significative pour l'inférence efficace des LLM :

• Démocratisation des Brouilleurs : METASD permet d'exploiter efficacement des pools de brouilleurs pré-entraînés hétérogènes (disponibles sur des dépôts publics) sans nécessiter de réentraînement coûteux ni de fine-tuning complexe.
• Adaptabilité Dynamique : En passant d'une sélection statique à une sélection dynamique au niveau du jeton, le système s'adapte aux variations de contexte, rendant l'inférence plus robuste face à la diversité des requêtes utilisateurs.
• Théorie et Pratique : La démonstration théorique de l'efficacité de la récompense BD et la validation empirique sur des scénarios réels (multilingue, OOD, prompts perturbés) offrent une feuille de route solide pour les futurs systèmes de décodage spéculatif.

En résumé, METASD résout le compromis entre la spécialisation des modèles et la généralisation requise pour des applications LLM réelles, en utilisant l'apprentissage par bandit pour orchestrer intelligemment une flotte de modèles de décodage.