Résumé technique : FlexDraft

Énoncé du problème

Le décodage spéculatif accélère l'inférence des grands modèles de langage (LLM) en utilisant un rédacteur léger pour proposer plusieurs jetons candidats en vue d'une vérification parallèle par un modèle cible. Bien qu'efficace, le décodage spéculatif séquentiel conventionnel souffre d'une attente mutuelle entre les phases de rédaction et de vérification, ainsi que d'une surcharge élevée d'accès mémoire due aux échanges répétés d'états.

Les méthodes récentes de décodage spéculatif parallèle tentent de chevaucher la rédaction et la vérification au sein d'un seul passage avant du modèle cible pour atténuer l'attente mutuelle. Cependant, ce paradigme introduit deux limitations fondamentales qui dégradent les performances, en particulier pour les grandes tailles de lot :

1. Incertitude du jeton bonus : Dans le décodage parallèle, les jetons de rédaction sont générés avant que le résultat de la vérification (spécifiquement le « jeton bonus » échantillonné au premier point de rejet) ne soit connu. Cela entraîne un décalage entre la distribution de rédaction et les contraintes de vérification, réduisant les taux d'acceptation.
2. Incertitude de la longueur d'acceptation : Puisque le nombre de jetons acceptés est inconnu pendant la rédaction, les méthodes parallèles doivent préparer des branches candidates pour tous les préfixes acceptés possibles. Cela force le modèle cible à traiter un nombre croissant de « jetons masqués » (rédactions futures) parallèlement aux jetons de vérification, provoquant une surcharge computationnelle qui évolue de manière quadratique ($O(N^2)$) avec la longueur du bloc de rédaction. Cette redondance fait s'effondrer les gains de débit aux grandes tailles de lot.

Méthodologie : FlexDraft

FlexDraft est un cadre de décodage spéculatif sans perte conçu pour s'adapter de manière flexible aux tailles de lot variables tout en préservant la distribution du modèle cible. Il répond aux limitations susmentionnées grâce à trois composants principaux :

1. Ajustement de l'attention (Attn Tuning)

Pour permettre la rédaction par diffusion de blocs (génération de plusieurs jetons en parallèle) sans réentraînement complet du modèle ni dégradation de la qualité :

• Architecture : Le cadre conserve le chemin autorégressif du modèle cible (couches FFN et projecteurs d'attention d'origine) figé.
• Paramètres entraînables : Seuls les projecteurs d'attention ($Q, K, V, O$) des dernières couches sont ajustés spécifiquement pour les jetons masqués.
• Mécanisme : Les jetons masqués sont ajoutés au préfixe propre. Dans les couches profondes, ces jetons masqués s'attendent de manière bidirectionnelle au sein de leur bloc et vers le préfixe propre, tout en restant isolés des autres blocs. Cela permet au rédacteur d'hériter des connaissances préentraînées du modèle cible via les FFN figés tout en apprenant à prédire plusieurs jetons futurs en parallèle en utilisant un nombre minimal de paramètres supplémentaires (environ 6 % de la taille du modèle cible).

2. Calibration guidée par le bonus

Pour résoudre le décalage causé par l'incertitude du jeton bonus :

• Mécanisme : Un Perceptron Multicouche (MLP) léger est appliqué après la tête de modélisation du langage (LM).
• Entrée : Le MLP prend la concaténation de l'embedding du jeton bonus résolu et de l'état caché de la position du masque actuelle.
• Sortie : Il produit un biais de calibration additif pour les logits de rédaction. Cela oriente la distribution de rédaction pour qu'elle soit cohérente avec le jeton bonus déterminé par le vérificateur, alignant efficacement la rédaction sur les contraintes de vérification même si la rédaction a été générée sans connaissance du jeton bonus.

3. Décodage flexible

Pour gérer le compromis entre l'attente mutuelle (petites tailles de lot) et le calcul redondant (grandes tailles de lot) :

• Commutation dynamique : Le système commute dynamiquement entre Rédaction et vérification parallèles (pour les petites tailles de lot) et Rédaction séquentielle puis vérification (pour les grandes tailles de lot) en fonction d'un seuil de taille de lot.
• Vérification sélective (petites tailles de lot) : En mode parallèle, le système estime la probabilité des longueurs acceptées en utilisant la confiance cumulative de rédaction. Il élague les branches candidates à faible confiance, ne vérifiant qu'un sous-ensemble de préfixes à haute confiance pour réduire la surcharge du passage avant cible sans modifier la longueur du bloc de rédaction.
• Exécution découplée (grandes tailles de lot) : En mode séquentiel, une fois la longueur acceptée résolue, le rédacteur ne génère que la branche unique conditionnée par le préfixe vérifié. Cela réutilise le cache KV de l'étape de vérification cible, éliminant les branches redondantes qui causent une surcharge de $O(N^2)$.

Contributions clés

1. Rédaction par diffusion de blocs sans perte : Une méthode permettant la génération parallèle de jetons en ajustant uniquement les projecteurs d'attention des dernières couches, préservant la distribution cible et évitant un préentraînement continu coûteux.
2. Calibration guidée par le bonus : Un mécanisme pour atténuer le décalage rédaction-vérification en calibrant les logits de rédaction sur la base du jeton bonus résolu, améliorant considérablement l'alignement.
3. Exécution adaptative à la taille de lot : Un cadre unifié qui commute entre les stratégies parallèles et séquentielles et emploie un élagage basé sur la confiance pour empêcher l'effondrement de l'accélération aux grandes tailles de lot.

Résultats expérimentaux

Des expériences ont été menées sur des modèles Qwen3 (1,7B, 4B et 8B) sur des benchmarks incluant GSM8K, MATH, HumanEval, MBPP et MT-Bench.

• Performance : FlexDraft atteint une accélération moyenne de 4,59× par rapport au décodage autorégressif sur Qwen3-8B sans perte de qualité.
• Comparaison : Il surpasse des bases solides, notamment EAGLE-3, DFlash, BiTA et Apple MTP. Spécifiquement, sur Qwen3-8B, FlexDraft atteint une accélération de 5,88× sur GSM8K et 5,79× sur MATH, surpassant DFlash (4,55× et 5,50× respectivement).
• Robustesse à la taille de lot : Contrairement aux méthodes parallèles existantes qui souffrent d'un effondrement de l'accélération aux grandes tailles de lot en raison de calculs redondants, FlexDraft maintient une accélération supérieure sur toutes les tailles de lot (1 à 16) en passant à une exécution séquentielle et en élaguant les branches à faible confiance si nécessaire.
• Études d'ablation :
  • L'ajustement de l'attention seul fournit des gains significatifs en réutilisant les connaissances cibles.
  • La calibration guidée par le bonus améliore davantage l'accélération en alignant les rédactions sur les contraintes de vérification.
  • La vérification sélective réduit les calculs redondants du passage avant cible, contribuant à l'accélération finale.
  • L'utilisation de 10 couches pour le modèle de rédaction s'est révélée offrir le meilleur compromis entre la longueur d'acceptation et la surcharge de décodage.

Importance

L'article affirme que FlexDraft offre une voie pratique vers une inférence efficace des LLM dans des conditions de service diversifiées. En combinant une adaptation légère des paramètres avec des stratégies d'exécution dynamiques, il répond avec succès aux compromis inhérents du décodage spéculatif parallèle. Le cadre atteint une accélération élevée sans compromettre la distribution du modèle cible (sans perte) et évite la dégradation de qualité souvent associée aux méthodes nécessitant un préentraînement continu extensif. La capacité à maintenir un débit élevé sur des tailles de lot variables le rend particulièrement adapté aux environnements de production où les caractéristiques de la charge de travail fluctuent.