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

1. Problemstellung

Large Language Models (LLMs) stoßen bei der Inferenz mit langem Kontext auf einen kritischen Speicherengpass. Während die autoregressive Generierung auf den Key-Value (KV) Cache angewiesen ist, um die Neuberechnung vergangener Tokens zu vermeiden, wächst der Speicherverbrauch dieses Caches linear mit der Sequenzlänge ($O(L \cdot N \cdot d)$, wobei $L$ die Anzahl der Schichten, $N$ die Sequenzlänge und $d$ die verborgene Dimension ist).

Bestehende Lösungen nach dem Training nutzen primär KV-Cache-Pruning, bei dem Tokens mit niedrigen Attention-Werten verworfen werden, um innerhalb eines festen Speicherbudgets zu bleiben. Eine grundlegende Einschränkung aktueller Methoden (z. B. H2O, StreamingLLM) ist jedoch die Annahme eines uniformen Prunings: Sie wenden dasselbe Pruning-Verhältnis über alle Transformer-Schichten hinweg an. Die Autoren argumentieren, dass dies suboptimal ist, da implizit davon ausgegangen wird, dass alle Schichten gleichermaßen zur Modellleistung beitragen, wodurch die heterogene Sensitivität verschiedener Schichten gegenüber dem Entfernen von Tokens ignoriert wird.

2. Methodik: DepthKV

Das vorgeschlagene Framework, DepthKV, hinterfragt die Annahme des uniformen Prunings durch die Einführung einer schichtabhängigen KV-Cache-Allokation. Anstatt einen festen Prozentsatz an Tokens pro Schicht zu verwerfen, verteilt DepthKV ein festes globales KV-Budget über die Schichten hinweg basierend auf deren spezifischer Sensitivität gegenüber Pruning.

A. Vorstudie: Schicht-Sensitivitätsanalyse

Bevor die Allokationsstrategie entwickelt wurde, führten die Autoren eine schichtweise Ablationsstudie durch:

• Methode: Sie prunten zu jedem Zeitpunkt Tokens nur in einer spezifischen Schicht, während sie die anderen intakt ließen, und maßen den daraus resultierenden Leistungsabfall (ROUGE-1).
• Ergebnis: Die Leistungsverschlechterung ist hochgradig nicht-uniform. Bestimmte „kritische" Schichten verursachen massive Einbußen beim Pruning, während andere einen vernachlässigbaren Einfluss oder sogar Redundanz zeigen.
• Content Amplification: Sie identifizierten „Content-Amplification-Schichten". Das Pruning dieser Schichten senkt nicht nur die Genauigkeit, sondern unterdrückt signifikant die Ausgabe-Länge und Informationsdichte (gemessen durch YapScore), was darauf hindeutet, dass diese Schichten entscheidend für die Aufrechterhaltung der Generierung sind.
• Statistische Validierung: Permutationstests bestätigten, dass die Variation der Schichtwichtigkeit über verschiedene Modelle (Gemma, LLaMA, Qwen) und Datensätze hinweg statistisch signifikant ist.

B. Repräsentationsmetriken für Wichtigkeit

Um die Allokation ohne exhaustive Ablation zu automatisieren, analysierten die Autoren Repräsentationseigenschaften, um einen Proxy für die Schichtwichtigkeit zu finden.

• Metrikenauswahl: Sie bewerteten sechs Metriken (Spectral Entropy, Effective Rank, Curvature, DiME, LiDAR und InfoNCE).
• Hauptergebnis: InfoNCE (ein kontrastiver Loss, der die Robustheit von Repräsentationen gegenüber Eingabe-Perturbationen misst) zeigte die stärkste negative Korrelation mit der Leistungsverschlechterung.
  • Interpretation: Schichten mit niedrigeren InfoNCE-Werten (robustere Repräsentationen) sind kritischer. Das Pruning dieser Schichten führt zu größeren Leistungseinbußen.
  • Verwendung: DepthKV nutzt InfoNCE-Werte, um Schichten zu rangieren und die Budget-Allokation zu bestimmen.

C. Allokationsstrategien

DepthKV operiert unter einem festen globalen Pruning-Verhältnis ($\rho$), verteilt dieses jedoch nicht-uniform. Drei Strategien werden vorgeschlagen:

1. Middle-Layer Protection (MLP): Ein heuristischer Ansatz, der eine feste Teilmenge mittlerer Schichten (z. B. Schichten $\lfloor L/2 \rfloor$ und $\lfloor L/2 \rfloor + 1$) bewahrt, während andere uniform geprunt werden.
2. Metric-Guided Allocation (MGA): Die primär vorgeschlagene Methode. Sie berechnet einen Wichtigkeitswert $s^{(l)}$ basierend auf InfoNCE. Schichten mit höherer Robustheit (niedrigerer InfoNCE) erhalten ein größeres KV-Budget (niedrigeres Pruning-Verhältnis), während weniger kritische Schichten aggressiver geprunt werden.
3. Middle-Layer Metric Allocation (MLMA): Ein hybrider Ansatz, der mittlere Schichten strukturell schützt, während das verbleibende Budget basierend auf InfoNCE-Werten auf andere Schichten verteilt wird.

3. Hauptbeiträge

1. Theoretische Erkenntnis: Es wurde gezeigt, dass Transformer-Schichten eine statistisch signifikante, nicht-uniforme Sensitivität gegenüber KV-Cache-Pruning aufweisen, was die Annahme des uniformen Prunings widerlegt.
2. Entdeckung der Content Amplification: Identifizierung spezifischer Schichten, deren Pruning die Inhaltsgenerierung unterdrückt, und direkte Verknüpfung der Repräsentationsrobustheit mit Ausgabequalität und -länge.
3. Framework-Vorschlag: Einführung von DepthKV, einem schichtabhängigen Allokationsframework, das ein festes Speicherbudget basierend auf der Schichtwichtigkeit (speziell unter Verwendung von InfoNCE) neu verteilt.
4. Empirische Validierung: Es wurde gezeigt, dass DepthKV über diverse Aufgaben und Modelle hinweg konsistent das uniforme Pruning übertrifft, ohne zusätzliches Training zu erfordern.

4. Experimentelle Ergebnisse

Die Autoren bewerteten DepthKV an drei Modellfamilien (Gemma-7B, LLaMA-3.1-8B, Qwen2.5-7B) über sieben Datensätze:

• Aufgaben: Zusammenfassung langer Dokumente (arXiv, PubMed, GovReport, LegalCase), dokumentenbasierte QA (HotpotQA, Qasper) und mathematisches Schlussfolgern (GSM-$\infty$).
• Setup: Feste globale KV-Reduktion von 60 % (40 % Beibehaltung).
• Leistung:
  • Zusammenfassung: DepthKV (speziell MGA) verbesserte ROUGE-1 und SBERT-Ähnlichkeit signifikant im Vergleich zum uniformen Pruning (z. B. +3,0 ROUGE-1 auf arXiv für Gemma).
  • QA & Schlussfolgern: DepthKV-Varianten erzielten höhere Exact Match (EM) und F1-Werte. Beispielsweise übertraf MLMA-6L auf Qasper alle Baselines.
  • Robustheit: Die Vorteile erstreckten sich auf Schlussfolgerungsaufgaben (GSM-$\infty$), bei denen alle DepthKV-Varianten das uniforme Pruning übertrafen.
  • LLM-as-a-Judge: Evaluierungen mit Prometheus bestätigten, dass DepthKV-Ausgaben korrekter, vollständiger und prägnanter waren als Baselines mit uniformem Pruning.
• Ablation: Die Studie bestätigte, dass wie das Budget verteilt wird (schichtweise Allokation) wichtiger ist als der spezifische Token-Wichtigkeits-Schätzer (z. B. H2O mit V vs. ohne V).

5. Bedeutung und Einschränkungen

Bedeutung:

• Effizienz: DepthKV bietet eine einfache, nach-trainingsbasierte Lösung für den Speicherengpass bei langem Kontext und ermöglicht längere Kontextfenster auf bestehender Hardware ohne Neutrainieren.
• Generalisierbarkeit: Die Methode ist architekturunabhängig und über verschiedene Modellfamilien und Aufgabentypen hinweg effektiv.
• Paradigmenwechsel: Sie verschiebt das Feld vom „Token-Level"-Pruning hin zur „Layer-Level"-Ressourcenzuteilung und anerkennt die funktionale Heterogenität von Transformer-Schichten.

Einschränkungen:

• Nicht-Query-bewusst: DepthKV schätzt die Token-Wichtigkeit während der Prefill-Phase basierend auf historischer Attention. Es passt sich nicht dynamisch an die spezifische Decoding-Query an und könnte Tokens verpassen, die erst für spezifische zukünftige Queries kritisch werden.
• Granularität auf Head-Ebene: Das aktuelle Framework weist Budgets auf der Schicht-Ebene zu, aggregiert die Attention jedoch über Heads. Es berücksichtigt keine Intra-Schicht-Variabilität, bei der bestimmte Attention-Heads kritischer sein könnten als andere.

Fazit

DepthKV stellt fest, dass nicht alle Transformer-Schichten bei der Inferenz mit langem Kontext gleichwertig sind. Durch die Neuverteilung von Speicherbudgets, um die empfindlichsten und robustesten Schichten zu erhalten (identifiziert via InfoNCE), erreicht die Methode überlegene Performance-Speicher-Verhältnisse im Vergleich zum uniformen Pruning und bietet einen praktischen Weg nach vorne für die Skalierung von LLM-Kontextfenstern.