DepthKV: Layer-Dependent KV Cache Pruning for Long-Context LLM Inference

1. Énoncé du problème

Les grands modèles de langage (LLM) font face à un goulot d'étranglement mémoire critique lors de l'inférence à contexte long. Bien que la génération autoregressive repose sur le cache clé-valeur (KV) pour éviter de recomputer les tokens passés, l'empreinte mémoire de ce cache croît linéairement avec la longueur de la séquence ($O(L \cdot N \cdot d)$, où $L$ est le nombre de couches, $N$ la longueur de la séquence et $d$ la dimension cachée).

Les solutions existantes post-entraînement utilisent principalement l'élagage du cache KV, qui élimine les tokens ayant des scores d'attention faibles pour s'adapter à un budget mémoire fixe. Cependant, une limitation fondamentale des méthodes actuelles (par exemple H2O, StreamingLLM) réside dans l'hypothèse d'un élagage uniforme : elles appliquent le même ratio d'élagage à toutes les couches du transformeur. Les auteurs soutiennent que cela est sous-optimal car cela suppose implicitement que toutes les couches contribuent également aux performances du modèle, ignorant la sensibilité hétérogène des différentes couches à la suppression de tokens.

2. Méthodologie : DepthKV

Le cadre proposé, DepthKV, remet en question l'hypothèse d'un élagage uniforme en introduisant une allocation de cache KV dépendante de la couche. Au lieu de rejeter un pourcentage fixe de tokens par couche, DepthKV réalloue un budget KV global fixe entre les couches en fonction de leur sensibilité spécifique à l'élagage.

A. Pré-étude : Analyse de la sensibilité des couches

Avant de concevoir la stratégie d'allocation, les auteurs ont mené une étude d'ablation couche par couche :

• Méthode : Ils ont élagué des tokens dans une seule couche spécifique à la fois tout en maintenant les autres intactes, mesurant ainsi la baisse de performance résultante (ROUGE-1).
• Résultat : La dégradation des performances est hautement non uniforme. Certaines couches « critiques » provoquent des chutes massives lorsqu'elles sont élaguées, tandis que d'autres montrent un impact négligeable ou même une redondance.
• Amplification du contenu : Ils ont identifié des « couches d'amplification du contenu ». L'élagage de ces couches réduit non seulement la précision, mais supprime également de manière significative la longueur et l'informativité de la sortie (mesurées par YapScore), indiquant que ces couches sont cruciales pour soutenir la génération.
• Validation statistique : Des tests de permutation ont confirmé que la variation de l'importance des couches est statistiquement significative à travers les modèles (Gemma, LLaMA, Qwen) et les jeux de données.

B. Métriques de représentation pour l'importance

Pour automatiser l'allocation sans recourir à une ablation exhaustive, les auteurs ont analysé les propriétés de représentation afin de trouver un proxy pour l'importance des couches.

• Sélection de la métrique : Ils ont évalué six métriques (Entropie spectrale, Rang effectif, Courbure, DiME, LiDAR et InfoNCE).
• Résultat clé : InfoNCE (une perte contrastive mesurant la robustesse de la représentation aux perturbations d'entrée) a montré la corrélation négative la plus forte avec la dégradation des performances.
  • Interprétation : Les couches avec des valeurs plus faibles d'InfoNCE (représentations plus robustes) sont plus critiques. Leur élagage provoque des chutes de performance plus importantes.
  • Utilisation : DepthKV utilise les scores InfoNCE pour classer les couches et déterminer l'allocation du budget.

C. Stratégies d'allocation

DepthKV opère sous un ratio d'élagage global fixe ($\rho$) mais le distribue de manière non uniforme. Trois stratégies sont proposées :

1. Protection des couches intermédiaires (MLP) : Une approche heuristique préservant un sous-ensemble fixe de couches intermédiaires (par exemple, les couches $\lfloor L/2 \rfloor$ et $\lfloor L/2 \rfloor + 1$) tout en élaguant uniformément les autres.
2. Allocation guidée par la métrique (MGA) : La méthode principale proposée. Elle calcule un score d'importance $s^{(l)}$ basé sur InfoNCE. Les couches ayant une robustesse plus élevée (InfoNCE plus faible) reçoivent un budget KV plus important (ratio d'élagage plus faible), tandis que les couches moins critiques sont élaguées plus agressivement.
3. Allocation métrique des couches intermédiaires (MLMA) : Une approche hybride qui protège structurellement les couches intermédiaires tout en distribuant le budget restant parmi les autres couches en fonction des scores InfoNCE.

3. Contributions clés

1. Insight théorique : Démonstration que les couches du transformeur présentent une sensibilité statistiquement significative et non uniforme à l'élagage du cache KV, invalidant l'hypothèse d'un élagage uniforme.
2. Découverte de l'amplification du contenu : Identification de couches spécifiques dont l'élagage supprime la génération de contenu, reliant directement la robustesse de la représentation à la qualité et à la longueur de la sortie.
3. Proposition de cadre : Introduction de DepthKV, un cadre d'allocation dépendant de la couche qui redistribue un budget mémoire fixe basé sur l'importance des couches (en utilisant spécifiquement InfoNCE).
4. Validation empirique : Démonstration que DepthKV surpasse systématiquement l'élagage uniforme sur diverses tâches et modèles sans nécessiter d'entraînement supplémentaire.

4. Résultats expérimentaux

Les auteurs ont évalué DepthKV sur trois familles de modèles (Gemma-7B, LLaMA-3.1-8B, Qwen2.5-7B) à travers sept jeux de données :

• Tâches : Résumé de longs documents (arXiv, PubMed, GovReport, LegalCase), QA ancrée dans le document (HotpotQA, Qasper) et raisonnement mathématique (GSM-$\infty$).
• Configuration : Réduction KV globale fixe de 60 % (40 % de rétention).
• Performances :
  • Résumé : DepthKV (spécifiquement MGA) a considérablement amélioré ROUGE-1 et la similarité SBERT par rapport à l'élagage uniforme (par exemple, +3,0 ROUGE-1 sur arXiv pour Gemma).
  • QA et Raisonnement : Les variantes de DepthKV ont obtenu des scores Exact Match (EM) et F1 plus élevés. Par exemple, sur Qasper, MLMA-6L a surpassé toutes les lignes de base.
  • Robustesse : Les avantages se sont étendus aux tâches de raisonnement (GSM-$\infty$), où toutes les variantes de DepthKV ont surpassé l'élagage uniforme.
  • LLM-as-a-Judge : Les évaluations utilisant Prometheus ont confirmé que les sorties de DepthKV étaient plus correctes, complètes et concises que les lignes de base d'élagage uniforme.
• Ablation : L'étude a confirmé que la manière dont le budget est distribué (allocation couche par couche) importe plus que l'estimateur spécifique de l'importance des tokens (par exemple, H2O avec V vs sans V).

5. Importance et Limites

Importance :

• Efficacité : DepthKV offre une solution simple et post-entraînement au goulot d'étranglement mémoire du contexte long, permettant des fenêtres de contexte plus longues sur le matériel existant sans réentraînement.
• Généralisabilité : La méthode est agnostique à l'architecture et efficace à travers différentes familles de modèles et types de tâches.
• Changement de paradigme : Elle fait passer le domaine de l'élagage au niveau « token » à l'allocation de ressources au niveau « couche », reconnaissant l'hétérogénéité fonctionnelle des couches du transformeur.

Limites :

• Non conscient de la requête : DepthKV estime l'importance des tokens lors de l'étape de préremplissage basée sur l'attention historique. Il ne s'ajuste pas dynamiquement à la requête de décodage spécifique, risquant ainsi de manquer des tokens qui ne deviennent critiques que pour des requêtes futures spécifiques.
• Granularité au niveau des têtes : Le cadre actuel alloue des budgets au niveau de la couche mais agrège l'attention à travers les têtes. Il ne prend pas en compte la variabilité intra-couche où des têtes d'attention spécifiques pourraient être plus critiques que d'autres.

Conclusion

DepthKV établit que toutes les couches du transformeur ne se valent pas dans l'inférence à contexte long. En réallouant les budgets mémoire pour préserver les couches les plus sensibles et robustes (identifiées via InfoNCE), la méthode atteint des compromis performance-mémoire supérieurs par rapport à l'élagage uniforme, offrant une voie pratique pour l'extension des fenêtres de contexte des LLM.