KVSculpt: KV Cache Compression as Distillation

1. Problemstellung

Bei der Inferenz von autoregressiven Large Language Models (LLMs) mit langen Kontexten wird der Key-Value (KV) Cache zum dominanten Speicher-Engpass. Um diesen zu komprimieren, existieren zwei Hauptansätze:

1. Reduktion der Paargröße: Quantisierung oder Low-Rank-Zerlegung.
2. Reduktion der Sequenzlänge: Verringerung der Anzahl der KV-Paare.

Bestehende Methoden zur Sequenzlängenreduktion fallen in zwei Kategorien, die beide an den originalen Cache-Einträgen festhalten:

• Eviction (Verwerfen): Es wird eine Teilmenge der originalen KV-Paare basierend auf Metriken (z. B. Aufmerksamkeitsscore, Recency) ausgewählt und der Rest verworfen.
• Merging (Zusammenführen): Ähnliche Paare werden gewichtet gemittelt, bleiben aber als Kombination der Originalwerte gebunden.

Die Autoren argumentieren, dass dieser diskrete Auswahlrahmen unnötig restriktiv ist. Da KV-Paare nach der RoPE-Positionscodierung Vektoren in einem kontinuierlichen Embedding-Raum sind, gibt es keine zwingende Notwendigkeit, dass der komprimierte Cache eine Teilmenge der Original-Einträge sein muss.

2. Methodik: KVSCULPT

KVSCULPT reformuliert das Problem der KV-Cache-Kompression von einer diskreten Auswahl in einen Distillations-Prozess. Ziel ist es, eine kleinere Menge von $k$ unbeschränkten KV-Paaren zu optimieren, die das Aufmerksamkeitsverhalten der ursprünglichen $N$ Paare exakt nachahmen.

Kernkomponenten:

• Unbeschränkte Optimierung: Im Gegensatz zu Eviction oder Merging sind die Keys ($K_c$) und Values ($V_c$) im komprimierten Cache frei im $\mathbb{R}^d$-Raum optimierbar. Sie müssen nicht mit originalen Token-Positionen übereinstimmen.
• Alternierender Optimierer:
  • Keys ($K_c$): Werden mit L-BFGS (ein Quasi-Newton-Verfahren) optimiert. Dies ist entscheidend, da der durch Softmax erzeugte Loss-Landscape nicht-konvex ist und Gradientenabstiegsverfahren (wie Adam) in lokalen Minima stecken bleiben.
  • Values ($V_c$): Werden analytisch in geschlossener Form durch Ridge-Regression (Least Squares) gelöst, gegeben die aktuellen Attention-Weights.
  • Die Schritte werden abwechselnd durchgeführt (z. B. 5 L-BFGS-Schritte für Keys, dann eine LS-Lösung für Values).
• Trainings-Query-Konstruktion: Da zukünftige Queries beim Komprimieren unbekannt sind, werden diese durch eine Kombination aus:
  1. Den realen "Retain"-Queries (die jüngsten Token).
  2. Synthetischen zukünftigen Queries: Diese werden durch Ent-RoPE-faktorierung (Rückgewinnung positionsunabhängiger Content-Vektoren) und erneute Anwendung von RoPE auf zukünftige Positionen generiert. Ein gleichmäßiges Sampling über den Kontext hat sich als robusteste Strategie erwiesen.
• Loss-Funktion: Ein zweigliedriger Loss, der sowohl den Output-MSE (Fehler der Attention-Ausgabe) als auch die Übereinstimmung des Log-Sum-Exp (LSE) der Attention-Scores minimiert, um die korrekte Gewichtung des Kontexts sicherzustellen.

Adaptive Budget-Allokation

Da die Schwierigkeit der Kompression zwischen verschiedenen Schichten (Layers) und KV-Heads stark variiert, führt KVSCULPT eine adaptive Budget-Verteilung ein:

1. Pilot-Lauf: Ein kurzer, kostengünstiger Kompressionslauf (ohne vollständige Optimierung) ermittelt den Mean Squared Error (MSE) pro Layer und Head.
2. Verteilung: Das feste Gesamtbudget wird basierend auf diesem Pilot-MSE neu verteilt. Schichten/Heads mit höherem Kompressionsfehler erhalten mehr KV-Paare, einfache weniger.
3. Dämpfung: Eine Quadratwurzel-Dämpfung ($MSE^{0.5}$) verhindert, dass Ausreißer-Schichten das gesamte Budget monopolisieren.

3. Wichtige Beiträge

1. Paradigmenwechsel: Umstellung von diskreter Token-Auswahl auf kontinuierliche Distillation, was die kombinatorische Suche eliminiert.
2. Leistungsfähiger Optimierer: Ein L-BFGS + Least-Squares-Verfahren, das signifikant bessere Ergebnisse liefert als reine Eviction-Methoden.
3. Granulare Budget-Steuerung: Einführung einer Pilot-basierten, adaptiven Allokation auf Layer- und Head-Ebene, die ohne zusätzliche Inferenzkosten die Qualität steigert.
4. Analyse der Komplexität: Nachweis, dass die Kompressions-Schwierigkeit extrem ungleich verteilt ist (Faktoren bis zu 467x innerhalb einer Schicht zwischen Heads).

4. Ergebnisse

Die Evaluation erfolgte auf Qwen2.5-1.5B-Instruct mit einem Kontext von 2048 Token.

• Hauptergebnisse: KVSCULPT reduziert die KL-Divergenz (Kullback-Leibler-Divergenz) im Vergleich zum besten Eviction-Baseline ("Select+Fit") um den Faktor 3,5 bis 4,1 über verschiedene Kompressionsraten ($r \in \{0.3, 0.5, 0.7\}$).
• Vergleich mit Joint Optimization: Ein Ansatz, der nur die Values optimiert, aber die Keys auf Originalpositionen beschränkt, verbessert die Ergebnisse kaum. Der Hauptgewinn stammt aus der Freiheit der Keys im Embedding-Raum.
• Adaptive Allokation: Die Pilot-basierte Budget-Verteilung reduziert die KL-Divergenz zusätzlich um den Faktor 1,3 bei null zusätzlichen Inferenzkosten.
• Optimierbarkeit: Der L-BFGS-Optimierer erreicht bereits mit einem Startpunkt (1 Restart) Ergebnisse, die nur 2–8% vom theoretischen Optimum (Oracle mit 100 Restarts) entfernt sind. Dies zeigt, dass das Problem der einzelnen Schicht im Wesentlichen gelöst ist.

5. Signifikanz und Limitationen

Bedeutung:
KVSCULPT beweist, dass KV-Cache-Kompression nicht auf das "Aussortieren" oder "Mischen" existierender Token beschränkt sein muss. Durch die Nutzung des kontinuierlichen Raums können komprimierte KV-Paare Attention-Muster "zusammenfassen", die von keinem einzelnen Original-Token repräsentiert werden. Dies ermöglicht nahezu verlustfreie Kompression auf einfachen Sequenzen, wo Eviction-Methoden um eine Größenordnung schlechter abschneiden.

Limitationen:

• Offline-Nutzung: Der Kompressionsprozess dauert ca. 170 Sekunden pro Kontext (auf einer A100 GPU), was KVSCULPT für Offline-Szenarien (z. B. Dokumentenverarbeitung) geeignet macht, aber nicht für Online-Echtzeit-Inferenz.
• Schichtübergreifende Fehlerakkumulation: Obwohl die Optimierung pro Schicht fast optimal ist, akkumulieren sich Fehler über die Schichten hinweg (um den Faktor 200–5800). Dies ist der aktuelle Engpass für "harte" Sequenzen.
• Evaluation: Die Ergebnisse basieren nur auf einem Modell (Qwen2.5-1.5B) und einer Kontextlänge (2048). Es fehlen Vergleiche mit Merge-Methoden und Quantisierung.

Zukunftsausblick:
Die Autoren schlagen "Cascade-aware Optimization" vor, bei der die Optimierung einer Schicht die Fehler der vorherigen Schichten berücksichtigt, um die verbleibende Lücke bei komplexen Sequenzen zu schließen.