When Hidden States Drift: Can KV Caches Rescue Long-Range Speculative Decoding?

1. Énoncé du problème

Déclin à longue portée dans le décodage spéculatif :
Le décodage spéculatif accélère l'inférence des grands modèles de langage (LLM) en utilisant un modèle « rédacteur » (drafter) léger pour proposer plusieurs jetons, qu'un modèle « cible » plus important vérifie. Les rédacteurs les plus avancés actuels (par exemple, EAGLE-3, MTP) reposent sur la réutilisation des états cachés, alimentant le modèle cible avec ses états cachés internes vers le rédacteur.

• Le problème : Ces systèmes souffrent d'un déclin à longue portée : la précision de la proposition diminue progressivement à mesure que l'étape spéculative $k$ augmente.
• Explication existante : La communauté attribue cela au décalage entraînement-inférence (le rédacteur se conditionne sur ses propres états cachés dérivants lors de l'inférence, contrairement à l'entraînement). Bien que l'entraînement au moment du test (TTT) atténue ce problème, le déclin persiste.
• Le vide : L'article soutient que le décalage entraînement-inférence n'est pas la seule cause. Le problème fondamental réside dans la préservation de l'information contextuelle. Les états cachés agissent comme une « compression biaisée du contexte » optimisée pour le jeton immédiat suivant, éliminant les informations pertinentes pour les étapes spéculatives futures.

2. Hypothèse centrale et analyse théorique

Les auteurs proposent l'hypothèse de la réutilisation des KV : réutiliser le cache KV (paires Clé-Valeur) du modèle cible, plutôt que simplement les états cachés, peut préserver un contexte plus riche pour la rédaction à longue portée.

• Réutilisation des états cachés (Compression biaisée) :
  • Les états cachés cibles sont le résultat d'une agrégation par attention : $h_t = \sum \alpha_i v_i$.
  • Les poids $\alpha_i$ sont optimisés pour prédire le jeton actuel $x_{t+1}$.
  • Les jetons faiblement pertinents pour $x_{t+1}$ reçoivent des poids proches de zéro et sont effectivement éliminés dans $h_t$.
  • Conséquence : Pour les étapes spéculatives ultérieures ($x_{t+k}$), le rédacteur doit récupérer des informations qui ont été supprimées par le mécanisme d'attention de la cible, un problème mal conditionné de « récupération d'information ».
• Réutilisation des KV (Contexte explicite) :
  • Le cache KV conserve l'ensemble complet des paires $(k_i, v_i)$ avant l'agrégation.
  • Le rédacteur peut effectuer une ré-attention en utilisant ses propres requêtes futures estimées contre le cache KV complet du préfixe.
  • Conséquence : Aucune information n'est éliminée avant le traitement par le rédacteur. La difficulté passe de « récupérer des informations perdues » à « estimer avec précision les requêtes futures » (un problème d'approximation de fonction).

3. Méthodologie : Le cadre KVShot

Pour tester cette hypothèse, les auteurs introduisent KVShot, un cadre diagnostique comparant trois paradigmes de réutilisation sous le même pipeline d'entraînement au moment du test autoregressif (TTT) (utilisé par EAGLE-3) :

1. Réutilisation uniquement des états cachés : Approche standard EAGLE-3/MTP (concaténation des états cachés cibles).
2. Réutilisation uniquement des KV : Le rédacteur reçoit le cache KV cible via une attention croisée, remplaçant entièrement les états cachés.
3. Réutilisation hybride : Un mécanisme de fusion par porte où l'état caché fournit l'ancre principale, et le cache KV fournit une « correction delta par porte » pour les étapes à longue portée.

Configuration expérimentale :

• Modèle cible : Qwen3-8B.
• Ensembles de données : ShareGPT et UltraChat (passés de 70k à 280k échantillons).
• Métriques : Taux d'acceptation étape par étape ($\alpha_k$), Moyenne des jetons acceptés (MAT), et proxies d'accélération de bout en bout (E2E).

4. Résultats clés

A. Validation des prédictions (Analyse étape par étape)

Les expériences ont confirmé trois prédictions théoriques :

1. Avantage à longue portée : La réutilisation uniquement des KV se dégrade plus gracieusement que la réutilisation uniquement des états cachés à mesure que $k$ augmente. À l'étape $k=6$, un rédacteur KV-only à 4 couches a dépassé la ligne de base EAGLE-3 ($\alpha_6 = 0,495$ contre $0,469$).
2. Goulot d'étranglement de l'estimation des requêtes : La performance KV-only dépend fortement de la profondeur du rédacteur. Un rédacteur à 1 couche a mal performé (MAT 1,84) car il ne peut pas approximer des requêtes cibles complexes. L'augmentation de la profondeur à 2 couches a produit un bond massif (MAT 2,23), mais 4 couches ont offert des rendements décroissants, restant encore derrière la ligne de base EAGLE-3 à 1 couche (MAT 2,37) en termes d'efficacité globale.
3. Désavantage à courte portée : À $k=0$ ou $1$, la réutilisation des états cachés surpasse la réutilisation KV-only car les états cachés portent un contenu sémantique plus riche que les embeddings bruts utilisés pour la ré-attention KV.

B. Performance hybride

Le rédacteur hybride (combinant les deux signaux) a obtenu les meilleurs résultats étape par étape :

• MAT : 2,54 (contre 2,37 pour la ligne de base EAGLE-3).
• Rétention : Amélioration du taux de rétention à longue portée de 73,5 % à 77,3 %.
• Mécanisme : La règle delta par porte a permis avec succès au chemin KV de corriger le chemin des états cachés à des horizons plus longs sans perturber la précision à courte portée.

C. Vérification de la réalité de bout en bout

Malgré les améliorations étape par étape, l'accélération de bout en bout (E2E) a été marginale :

• MAT mesuré par HF : N'a augmenté que de 5,01 (EAGLE-3) à 5,04 (Hybride).
• Latence : Le chemin d'attention croisée supplémentaire a introduit 5 à 10 % de latence de rédaction supplémentaire.
• Conclusion : Le pipeline TTT autoregressif actuel ne peut pas exploiter efficacement le signal KV pour surmonter le coût computationnel ajouté.

5. Goulots d'étranglement structurels identifiés

L'article identifie trois raisons spécifiques pour lesquelles le pipeline actuel échoue à traduire les gains étape par étape en accélérations E2E :

1. Difficulté d'estimation des requêtes : Estimer les requêtes cibles futures nécessite de répliquer les transformations non linéaires multi-couches du modèle cible. Les rédacteurs peu profonds (1-2 couches) manquent de capacité pour approximer ces requêtes avec précision.
2. Signaux de gradient clairsemés : Dans le TTT autoregressif, les projections KV côté rédaction ($W_K, W_V$) ne sont mises à jour que par les quelques jetons générés par la rédaction (par exemple, 7 jetons) par étape, tandis que le préfixe KV est copié. Cela résulte en des gradients extrêmement clairsemés et déséquilibrés pour le chemin KV.
3. Famine de gradients induite par la porte : Dans le modèle hybride, la porte contrôlant la correction KV a tendance à s'effondrer vers zéro au début de l'entraînement (car la branche KV initialisée aléatoirement commence comme du bruit). Une fois fermée, la porte bloque les gradients vers la branche KV, l'empêchant d'apprendre à être utile. Cela crée un « état d'injection faible auto-renforçant ».

6. Importance et orientations futures

• Changement conceptuel : L'article reframe le déclin à longue portée non seulement comme un décalage entraînement-inférence, mais comme un problème de préservation de l'information. Il prouve que les états cachés sont une compression avec perte pour les tâches à long horizon.
• Outil diagnostique : KVShot fournit un terrain d'essai fondamental pour isoler les effets de représentation des artefacts d'entraînement.
• Feuille de route pour les architectures de nouvelle génération : Les résultats suggèrent que le TTT autoregressif est insuffisant pour le décodage conscient des KV. Pour réaliser le potentiel de la réutilisation des KV, le domaine doit se tourner vers des paradigmes d'entraînement par blocs (par exemple, DFlash).
  • L'entraînement par blocs génère de nombreux jetons de rédaction en parallèle, fournissant des signaux de gradient denses pour les projections KV.
  • Il permet des modèles de rédaction plus profonds capables d'estimer avec précision les requêtes sans pénalités de latence séquentielle.

Résumé : Bien que les caches KV offrent théoriquement un signal supérieur pour le décodage spéculatif à longue portée en préservant le contexte, les pipelines d'entraînement autoregressifs actuels échouent à exploiter cela en raison des limites d'estimation des requêtes et de la famine de gradients. La voie à suivre réside dans des architectures d'entraînement parallèle par blocs.