Résumé technique : Éviction de cache KV stochastique sensible à la valeur pour les modèles de raisonnement

1. Énoncé du problème

Les modèles de raisonnement (par exemple, Qwen3, o1 d'OpenAI) atteignent une grande précision en générant de longues chaînes de pensée avant de produire une réponse finale. Cependant, cette capacité crée un goulot d'étranglement significatif en termes de mémoire et de calcul lors de la phase de décodage. À mesure que la longueur de la séquence augmente, le cache Clé-Valeur (KV) requis pour stocker les représentations de chaque jeton passé entraîne un surcoût substantiel.

Les solutions existantes se divisent en deux catégories :

• Méthodes basées sur la sélection : Elles conservent l'intégralité du cache KV mais n'activent qu'un sous-ensemble épars de jetons lors du calcul de l'attention. Bien qu'accurates, leur empreinte mémoire croît linéairement avec la longueur de la séquence ($O(T)$), ne résolvant pas le goulot d'étranglement de la mémoire.
• Méthodes basées sur l'éviction : Elles rejettent définitivement les paires KV de faible importance une fois qu'un budget prédéfini est atteint, offrant une empreinte mémoire statique et un meilleur débit. Cependant, les méthodes d'éviction actuelles souffrent d'une dégradation significative de la précision sur les tâches de raisonnement par rapport aux alternatives de sélection, provoquant souvent des boucles de raisonnement répétitives ou la génération de sorties non sensiques.

L'article identifie que les stratégies d'éviction actuelles ne tiennent pas compte de deux facteurs critiques : l'influence disproportionnée des états de valeur de grande magnitude et le besoin de diversité stochastique dans les jetons conservés.

2. Méthodologie : VASE

Les auteurs proposent VASE (Value-Aware Stochastic KV Cache Eviction), un cadre d'éviction sans entraînement conçu pour combler l'écart entre efficacité et précision. VASE opère au sein d'un cadre d'éviction périodique (utilisant un budget persistant $K$ et un tampon récent $B$) et introduit deux mécanismes de base :

A. Protection des états de valeur de grande magnitude

Les auteurs observent que les états de valeur dans les modèles de raisonnement présentent une distribution fortement asymétrique, où une petite fraction de jetons possède des vecteurs de magnitude anormalement élevée (mesurée par l'étendue $Range(v) = \max(v) - \min(v)$).

• Constat : L'éviction de ces valeurs de grande magnitude provoque un effondrement catastrophique de la précision (par exemple, passant d'environ 88 % à 14 % sur GSM8K) et induit des boucles répétitives où le modèle réexamine indéfiniment le contexte sans parvenir à une conclusion.
• Mécanisme : VASE réserve une partie spécifique du budget de jetons ($N_v$) pour conserver inconditionnellement les $N_v$ jetons ayant les plus grandes magnitudes de valeur. Cela garantit que les vecteurs de valeur les plus influents ne sont jamais écartés.

B. Introduction de la stochasticité

Les méthodes d'éviction actuelles utilisent souvent une sélection top-$k$ déterministe, ce qui peut entraîner un manque de diversité dans le cache conservé.

• Constat : L'introduction de la stochasticité améliore la précision en assurant une couverture plus représentative de l'ensemble du contexte.
• Mécanisme : Au lieu de sélectionner déterministement les jetons les mieux classés, VASE emploie un échantillonnage stochastique pondéré.
  • VASE-AttnV : Combine la réservation sensible à la valeur avec un échantillonnage stochastique basé sur les scores d'attention (dérivés de SnapKV).
  • VASE-DKV : Adapte la méthode CurDKV (qui utilise les scores de levier de la décomposition de matrice CUR) en rééchantillonnant la matrice de projection gaussienne $G$ à chaque étape d'éviction. Cela empêche les jetons ayant des représentations spécifiques d'être systématiquement assignés des scores faibles et évictés de façon permanente.

3. Contributions clés

1. Identification des facteurs critiques : L'article établit que (1) les états de valeur de grande magnitude sont cruciaux pour maintenir la progression du raisonnement et prévenir les boucles répétitives, et (2) la stochasticité dans les décisions d'éviction augmente considérablement la précision en augmentant la diversité du cache.
2. Cadre VASE : Une nouvelle recette d'éviction sans entraînement qui intègre la protection de la magnitude des états de valeur et l'échantillonnage stochastique. C'est la première méthode d'éviction qui combine scoring basé sur les clés, scoring basé sur les valeurs et promotion de la diversité.
3. Lien avec la quantification : Les auteurs démontrent que les états de valeur à large étendue sont également la source principale d'erreur de reconstruction dans la quantification du cache KV par jeton, suggérant que les enseignements de VASE se généralisent à d'autres techniques de compression.

4. Résultats expérimentaux

Les auteurs ont évalué VASE sur Qwen3-4B et Qwen3-14B à travers six tâches de raisonnement (AIME25/26, HMMT25, GPQA-Diamond, MATH, LiveCodeBench-v6) avec un ratio de compression du cache KV de 4×.

• Précision vs Méthodes de sélection : VASE-AttnV a atteint une précision moyenne plus élevée que la méthode de sélection la plus forte (SeerAttention-R) sur les deux tailles de modèles, tout en maintenant une empreinte mémoire statique.
  • Qwen3-4B : VASE-AttnV (59,09 %) a surpassé SeerAttention-R (58,81 %) et la base d'éviction la plus forte R-KV (54,69 %) de 4,4 %.
  • Qwen3-14B : VASE-AttnV (65,81 %) a égalé SeerAttention-R (65,37 %) et a surpassé R-KV (60,90 %) de 4,9 %.
• Études d'ablation :
  • Sensibilité à la valeur : La réservation de créneaux pour les valeurs de grande magnitude a amélioré la précision sur GSM8K jusqu'à 16,2 % par rapport aux bases de référence.
  • Stochasticité : L'ajout de l'échantillonnage stochastique à CurDKV a amélioré la précision de 9,2 % sur Qwen3-14B.
• Efficacité : VASE-DKV a atteint le débit le plus élevé (3,1× plus rapide que le modèle de base complet à 16K jetons) et l'utilisation de la mémoire de pointe la plus faible parmi toutes les méthodes testées.
• Génération de code : Sur LiveCodeBench, les méthodes VASE ont nettement surpassé la méthode de sélection SeerAttention-R, qui éprouvait des difficultés face aux changements de domaine.

5. Signification et affirmations

L'article affirme que VASE comble avec succès l'écart entre efficacité et précision qui a historiquement entravé les méthodes d'éviction de cache KV. En privilégiant les états de valeur de grande magnitude et en introduisant la stochasticité, VASE permet aux modèles de raisonnement de fonctionner avec une empreinte mémoire statique sans sacrifier la précision habituellement associée aux approches de cache complet ou de sélection.

Les auteurs soulignent que leurs découvertes concernant l'importance de la magnitude des états de valeur ont des implications plus larges au-delà de l'éviction, spécifiquement pour la quantification du cache KV, où les valeurs à large étendue sont identifiées comme une source primaire d'erreur. Ils suggèrent que les futures méthodes d'inférence économes en mémoire devraient envisager des approches à précision mixte qui protègent ces états critiques de haute magnitude.

En fin de compte, VASE fournit une recette simple, efficace et sans entraînement pour supporter FlashAttention2 et permettre une inférence scalable pour les modèles de raisonnement à longue chaîne.