Résumé Technique : MOMENTKV

Problématique

Le décodage autorégressif dans les modèles de langage basés sur les Transformers repose sur le cache Clé-Valeur (KV), qui stocke les projections précédemment calculées pour les réutiliser lors des étapes de décodage. À mesure que la longueur des séquences atteint des dizaines ou des centaines de milliers de jetons, l'empreinte mémoire du cache KV croît linéairement, devenant le principal goulot d'étranglement pour l'inférence à contexte long.

L'éviction du cache KV traite ce problème en conservant un sous-ensemble de taille fixe de paires clé-valeur et en écartant le reste. Bien que les méthodes existantes aient amélioré les stratégies de sélection de jetons (par exemple, basées sur les poids d'attention cumulés ou le scoring sensible aux valeurs), elles partagent une limitation fondamentale : l'inférence post-éviction renormalise l'attention exclusivement sur les jetons conservés, écartant complètement l'ensemble évincé.

Les auteurs identent que la principale source de dégradation de la sortie n'est pas simplement la masse d'attention résiduelle sur les jetons évincés (que les méthodes existantes minimisent déjà), mais un décalage directionnel entre l'ensemble de jetons conservés et l'ensemble évincé. L'analyse empirique révèle que, dans la pratique, les jetons évincés sont souvent presque orthogonaux à ceux conservés. Par conséquent, même une petite quantité de masse d'attention évincée peut produire une erreur directionnelle substantielle que la renormalisation standard ne peut récupérer. Cela crée un « écart directionnel » où l'amélioration de la seule sélection de jetons produit des rendements décroissants car la rétention plus sélective concentre l'ensemble conservé sur des directions spécifiques alignées avec la requête, élargissant davantage l'écart avec l'ensemble évincé.

Méthodologie : MOMENTKV

Pour remédier à ce décalage directionnel, les auteurs proposent MOMENTKV, un cadre qui maintient des statistiques de moments compactes et de petite taille sur l'ensemble de jetons évincés. Ces statistiques servent un double objectif : guider les décisions d'éviction et corriger la sortie de l'attention pendant l'inférence.

1. Statistiques de moments compactes

Au lieu de stocker les jetons évincés, MOMENTKV maintient quatre sommes glissantes par tête d'attention, mises à jour de manière incrémentale lors de chaque éviction :

• Compte ($n_e$) : Le nombre de jetons évincés.
• Somme des clés ($s_k$) : $\sum_{i \in E} k_i$.
• Somme des valeurs ($s_v$) : $\sum_{i \in E} v_i$.
• Somme du produit extérieur ($S$) : $\sum_{i \in E} v_i k_i^\top$.

Ces statistiques permettent de calculer la moyenne des clés ($\bar{k}$), la moyenne des valeurs ($\bar{v}$) et la covariance valeur-clé ($\tilde{S}$) sans revisiter les jetons évincés. La surcharge de stockage totale est de $O(d^2)$ par tête, indépendante de la longueur du contexte.

2. Approximation du premier ordre de la sortie évincée

Pendant l'inférence, MOMENTKV dérive une approximation de premier ordre en forme close de la sortie d'attention évincée, $\hat{f}_E(q)$. En recentrant les logits autour de la moyenne des clés et en appliant une expansion de Taylor au premier ordre à la fonction softmax, les auteurs dérivent :
$$\hat{f}_E(q) = \bar{v} + \frac{1}{n_e \sqrt{d}} \tilde{S} q$$

• Le terme $\bar{v}$ agit comme une base indépendante de la requête, récupérant la direction moyenne de la sous-sortie évincée.
• Le terme $\tilde{S} q$ agit comme une correction adaptative à la requête, capturant la variation perpendiculaire (contenu directionnel) que la base seule ne peut capturer.

Cette approximation permet au modèle de reconstruire une estimation de la contribution évincée perdue, restaurant efficacement l'information directionnelle que l'éviction standard oblitère.

3. Éviction informée par les moments

Les statistiques de moments guident également le processus d'éviction lui-même. Pour chaque jeton conservé $j$, l'algorithme calcule un résidu de moment :
$$r_j = v_j - \bar{v} - \frac{1}{n_e \sqrt{d}} \tilde{S} k_j$$
Ce résidu mesure la quantité d'information directionnelle nouvelle que le jeton $j$ apporte au-delà de ce que le résumé accumulé capture déjà. Le score d'éviction est défini par le produit du poids d'attention et de la norme du résidu : $\text{score}(j) = \alpha_j \cdot \|r_j\|$.

• Les jetons avec de faibles résidus sont bien prédits par le modèle de moment et sont prioritaires pour l'éviction.
• Les jetons avec des résidus élevés (portant une information directionnelle unique) sont conservés même si leurs poids d'attention sont plus faibles.

Cela crée une boucle de renforcement mutuel : l'éviction informée par les moments maintient l'ensemble évincé « géométriquement régulier » (faible dispersion des logits $\sigma$), ce qui resserre la borne d'approximation pour la correction d'inférence. En retour, la correction précise fournit de meilleurs résidus pour les futures décisions d'éviction.

4. Inférence corrigée par la normalisation

La sortie d'attention finale est reconstruite comme une combinaison convexe de la sortie conservée ($f_R$) et de la sortie évincée approximée ($\hat{f}_E$) :
$$\hat{f}(q) = \hat{w}_R f_R(q) + (1 - \hat{w}_R) \hat{f}_E(q)$$
Le poids de mélange $\hat{w}_R$ est estimé à l'aide d'une borne inférieure sur la partition de l'ensemble évincé dérivée de l'inégalité de Jensen. Cela garantit que le modèle prend en compte la magnitude de la masse évincée tout en corrigeant sa direction.

Contributions clés

1. Identification de l'écart directionnel : L'article établit que la source d'erreur dominante dans l'éviction du cache KV est la quasi-orthogonalité entre les sous-sorties conservées et évincées, et non seulement la magnitude de la masse d'attention évincée.
2. Statistiques de moments compactes : Introduction d'une méthode efficace en mémoire pour suivre les propriétés géométriques (moyenne et covariance) des jetons évincés en utilisant seulement $O(d^2)$ de stockage par tête.
3. Correction en forme close : Dérivation d'une approximation de premier ordre pour la sortie d'attention évincée qui récupère à la fois la direction moyenne et les variations dépendantes de la requête.
4. Mécanisme à double rôle : Un cadre unifié où les statistiques de moments optimisent simultanément la sélection d'éviction (en identifiant les jetons redondants) et corrigent les sorties d'inférence (en restaurant l'information directionnelle perdue).

Résultats expérimentaux

Les auteurs ont évalué MOMENTKV sur les benchmarks LongBench et RULER en utilisant LLaMA-3.1-8B-Instruct et Qwen3-4B-Instruct.

• Performance : MOMENTKV surpasse systématiquement toutes les lignes de base (H2O, SnapKV, PyramidKV, Ada-KV) sur tous les budgets de cache ($L \in \{128, 256, 512, 1024\}$).
• Compression agressive : Les gains les plus importants ont été observés sous une compression agressive. Sur LongBench avec $L=128$ et LLaMA-3.1-8B, MOMENTKV a obtenu une amélioration de score moyen de +1,35 par rapport à la meilleure ligne de base (Ada-KV), conservant 94,1 % de la performance du cache complet.
• Tâches de récupération : Sur RULER, qui teste la récupération précise de jetons, MOMENTKV montre des gains proportionnels de +3,4 et +3,5 par rapport à Ada-KV pour LLaMA et Qwen respectivement, validant son efficacité dans les contextes intensifs en récupération.
• Études d'ablation :
  • L'Inférence corrigée par la normalisation (NC) a contribué de manière plus significative (+2,41) que l'Éviction informée par les moments (MI) seule (+0,57), confirmant que la récupération de la direction évincée est le principal moteur de la performance.
  • La combinaison de ces deux mécanismes produit un effet synergique (+3,17 de gain total), dépassant la somme des gains individuels.
• Efficacité : La méthode ajoute seulement 1,6 ms de latence par jeton et nécessite 4,1 Mo de mémoire auxiliaire (6,4 % du cache conservé), tout en réduisant la latence par jeton de plus de 65 % par rapport à un cache complet.

Signification

L'article soutient que la voie vers l'amélioration de l'inférence à contexte long ne réside pas uniquement dans une meilleure sélection de jetons, mais dans la récupération de l'information directionnelle que les stratégies d'éviction actuelles oblitèrent. MOMENTKV démontre qu'en maintenant des statistiques compactes sur l'ensemble évincé, il est possible de combler l'écart directionnel, transformant le « coût irréductible » de l'éviction en une erreur d'approximation correctible. Cette approche offre une solution pratique et scalable pour déployer des modèles de langage à contexte long sans sacrifier l'efficacité des caches KV à budget fixe.