🤖 machine learning

Training Transformers for KV Cache Compressibility

Dieser Beitrag stellt KV-Compression Aware Training (KV-CAT) vor, eine Methode zur fortgesetzten Vorverarbeitung, die während des Trainings KV-Slots maskiert, um das Erlernen inhärent komprimierbarer Repräsentationen zu fördern und dadurch die Wirksamkeit der nachträglichen KV-Cache-Komprimierung für die Modellierung langer Kontexte in Sprachmodellen erheblich verbessert.

Ursprüngliche Autoren: Yoav Gelberg, Yam Eitan, Michael Bronstein, Yarin Gal, Haggai Maron

Veröffentlicht 2026-05-08

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Yoav Gelberg, Yam Eitan, Michael Bronstein, Yarin Gal, Haggai Maron

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich vor, Sie lesen ein sehr langes Buch und möchten sich die wichtigsten Teile merken, um später Fragen dazu beantworten zu können. In der Welt der KI wird dieses „Gedächtnis" als KV-Cache (Key-Value Cache) bezeichnet.

Hier ist das Problem: Je länger das Buch wird, desto mehr Notizen muss die KI für jedes einzelne Wort, das sie liest, führen. Irgendwann wird diese Liste so riesig, dass der Speicherplatz ausgeht oder das Durchlesen zu langsam wird.

Um dies zu beheben, haben Wissenschaftler versucht, diese Notizen „zu zusammenfassen", nachdem die KI bereits trainiert war. Sie versuchten, die langweiligen Notizen wegzuwerfen und nur die wichtigen zu behalten. Doch die Autoren dieses Papers entdeckten einen Fehler in diesem Ansatz: Es ist, als würde man versuchen, ein Buch zusammenzufassen, das in einer unordentlichen, chaotischen Handschrift geschrieben wurde. Egal wie sehr man sich auch bemüht, es zusammenzufassen, das ursprüngliche Chaos macht es unmöglich, die wichtigen Details zu behalten, ohne die Bedeutung zu verlieren.

Die große Idee: Von Anfang an sauber schreiben

Das Paper argumentiert, dass wir die KI nicht danach lehren sollten, die unordentlichen Notizen nach dem Training zu bereinigen, sondern dass wir die KI von Anfang an lehren sollten, saubere, komprimierbare Notizen zu schreiben.

Sie nennen diese neue Trainingsmethode KV-CAT (KV-Compression Aware Training).

Wie es funktioniert: Das „Versteckspiel"

Um die KI beizubringen, saubere Notizen zu schreiben, spielten die Forscher während des Trainings ein Spiel namens „Versteckspiel" (oder technisch präziser: KV-Sparsifizierung).

Das Setup: Stellen Sie sich vor, die KI liest einen Satz. Normalerweise betrachtet sie jedes einzelne Wort, um das nächste zu verstehen.
Die Wendung: Während des Trainings „verstecken" (maskieren) die Forscher zufällig etwa die Hälfte der Wörter. Die KI ist gezwungen, das nächste Wort zu erraten, ohne alle vorherigen Notizen zu sehen.
Die Lektion: Da sich die KI nicht darauf verlassen kann, jede einzelne Notiz zu sehen, lernt sie, ihr Gedächtnis anders zu organisieren. Sie lernt, die wichtigsten Informationen in die geringstmögliche Anzahl von Notizen zu packen, genau wie ein Schüler, der lernt, eine perfekte Zusammenfassung zu schreiben, weil er weiß, dass er an einer kurzen Version des Textes geprüft wird.
Das Sicherheitsnetz: Um sicherzustellen, dass die KI nicht vergisst, wie man normal liest, lassen sie sie gelegentlich auch den vollständigen, unversteckten Text lesen. Dies stellt sicher, dass sie auch dann intelligent und genau bleibt, wenn sie nicht komprimieren muss.

Das Ergebnis: Ein intelligenteres Gedächtnis

Als sie diese neue Methode testeten, waren die Ergebnisse beeindruckend:

Bessere Zusammenfassungen: Als sie versuchten, das Gedächtnis der KI später zu komprimieren (unter Verwendung Standard-Tools), behielt die mit KV-CAT trainierte KI im Vergleich zu einer Standard-KI viel mehr der ursprünglichen Bedeutung bei.
Schnellere Verarbeitung: Das Komprimieren des Gedächtnisses der KV-CAT-KI benötigte weniger Zeit und Aufwand.
Kein Verlust an Intelligenz: Entscheidend ist, dass die KI bei normalen Aufgaben nicht „dümmer" wurde. Sie konnte Fragen beantworten und Texte schreiben, genau so gut wie das ursprüngliche Modell, wenn sie nicht komprimiert wurde.

Das Fazit

Stellen Sie es sich wie das Packen für eine Reise vor.

Der alte Weg: Sie packen einen riesigen Koffer mit all Ihren Besitzungen, versuchen dann aber, ihn am Flughafen in eine winzige Tasche zu zwängen. Am Ende verlieren Sie Ihre Zahnbürste oder Ihr Lieblingshemd.
Der KV-CAT-Weg: Ihnen wird beigebracht, von Anfang an effizient zu packen. Sie lernen genau, was in die kleine Tasche passt, sodass Sie am Flughafen alles perfekt zuknöpfen können, ohne etwas Wichtiges zu verlieren.

Das Paper beweist, dass wir durch das Trainieren von KI, damit sie „kompressionsbewusst" ist, diese viel effizienter im Umgang mit langen Kontexten machen können, ohne die Architektur der KI ändern oder ihre Intelligenz opfern zu müssen.

Technischer Überblick: Training von Transformern für die Komprimierbarkeit des KV-Caches

1. Problemstellung

Die Modellierung von langen Kontexten wird zunehmend durch den Key–Value (KV)-Cache zu einem Engpass, dessen Speicherbedarf und Zugriffskosten zur Dekodierungszeit linear mit der Länge des Präfixes skalieren. Während neuere Arbeiten verschiedene Methoden zur Kontextkomprimierung vorgeschlagen haben – von der Zusammenfassung auf Token-Ebene bis hin zu optimierungsbasierten Komprimierungen des KV-Caches (z. B. gradientenbasiertes Matching oder Attention-Matching) – operieren diese Ansätze als nachgelagerte Eingriffe in fest vortrainierte Modelle.

Die Kernbeschränkung dieser nachgelagerten Methoden besteht darin, dass ihre Wirksamkeit fundamental durch die Komprimierbarkeit der internen Repräsentationen des Modells begrenzt ist. Die Arbeit argumentiert, dass Komprimierbarkeit nicht allein eine Eigenschaft der Eingabesequenz oder der Aufgabe ist, sondern eine Eigenschaft der gelernten Repräsentationen innerhalb des Transformers. Konkret können verschiedene Transformer-Implementierungen für dieselbe Eingabesequenz identische Next-Token-Verteilungen erzeugen, während sie völlig unterschiedliche interne Strukturen nutzen, von denen einige inhärent widerstandsfähig gegen Komprimierung sind.

2. Theoretische Motivation

Die Autoren formalisieren den Begriff der KV-Komprimierbarkeit: Ein Transformer ist KV-komprimierbar, wenn eine Komprimierungsstrategie existiert, die den KV-Cache langer Eingaben auf einen kürzeren Cache abbildet, wobei die Next-Token-Verteilung des Modells innerhalb eines begrenzten Fehlers erhalten bleibt.

Wichtige theoretische Erkenntnisse:

Existenz nicht-komprimierbarer Implementierungen: Die Arbeit beweist, dass für fast jede Sequenz-zu-Vektor-Funktion Transformer-Implementierungen existieren, die inhärent nicht komprimierbar sind (wobei jede nicht-triviale Komprimierung einen konstanten Fehler verursacht), neben Implementierungen, die hochgradig komprimierbar sind (bis hin zu einem einzigen KV-Paar komprimierbar).
Motivierendes Beispiel (Histogrammberechnung): Die Autoren zeigen, dass eine „natürliche" Transformer-Implementierung einer Histogrammfunktion (Berechnung von Token-Häufigkeiten) nicht komprimierbar sein kann, da sie auf der Mittelung von Token-Embeddings beruht, ohne die für die Korrektur von Skalierungsfehlern während der Komprimierung notwendigen Positionsinformationen zu bewahren. Im Gegensatz dazu ermöglicht eine leicht modifizierte, strukturierte Implementierung, die Positionsinformationen und unnormalisierte Zählungen explizit kodiert, eine perfekte Komprimierung auf ein einziges KV-Paar.
Implikation: Standard-Trainingsziele (Next-Token-Vorhersage) fördern das Entstehen komprimierbarer Repräsentationen nicht inhärent; sie können einfachere, nicht komprimierbare Lösungen begünstigen.

3. Methodik: KV-Kompressionsbewusstes Training (KV-CAT)

Um die theoretische Lücke zu schließen, schlagen die Autoren KV-Kompressionsbewusstes Training (KV-CAT) vor, ein Verfahren zum fortgesetzten Vortraining (CPT), das darauf ausgelegt ist, das Entstehen komprimierbarer interner Repräsentationen zu incentivieren.

Trainingsmechanismus

KV-CAT modifiziert die Standard-Trainings Schleife durch die Einführung einer Trainingszeit-KV-Sparsifizierungsstrategie:

Dualer Forward-Pass: Für jede Trainingssequenz führt das Modell zwei Durchläufe durch:
- Dichter Pass: Der Standard-Forward-Pass, bei dem alle KV-Slots aktiv sind.
- Maskierter Pass: Ein Forward-Pass, bei dem ein gelernter Router einen Anteil der KV-Slots maskiert (z. B. nur 50% beibehält).
Router-Mechanismus: Leichte gelernte Router (implementiert als lineare Attention-Module) werden zwischen Layer-Gruppen eingefügt. Sie geben einen skalaren Score für jeden Token aus, der durch einen Schwellenwert in eine binäre Maske umgewandelt wird. Diese Router werden gemeinsam mit den Modellparametern trainiert.
Trainingsziel: Die Verlustfunktion kombiniert drei Terme:
- Self-Distillation-Verlust ( $L_{mask}$ ): Minimiert die KL-Divergenz zwischen der Verteilung des maskierten Modells und der Verteilung des dichten Modells. Dies zwingt das Modell, die Leistung auch dann aufrechtzuerhalten, wenn KV-Slots maskiert sind.
- Anchor-Verlust ( $L_{anchor}$ ): Standard-Next-Token-Vorhersage (NTP)-Verlust, angewendet auf den dichten Forward-Pass, um das ursprüngliche Verhalten des Modells zu bewahren und katastrophales Vergessen zu verhindern.
- Budget-Verlust ( $L_{budget}$ ): Ein Regularisierungsterm, der den Router dazu anregt, eine Zielbeibehaltungsrate (z. B. 50%) einzuhalten und verhindert, dass das Modell trivialerweise alle Tokens oder keine maskiert.

Inferenz

Zum Evaluierungszeitpunkt werden die Router deaktiviert, und das Modell arbeitet als standardmäßiger dichter Transformer. Nachgelagerte Komprimierungsmethoden (z. B. Attention-Matching oder gradientenbasierte Optimierung) werden anschließend auf den KV-Cache dieses trainierten Modells angewendet.

4. Empirische Evaluation

Die Autoren evaluierten KV-CAT an QWEN2.5-Modellen (0,5 Mrd. und 1,5 Mrd. Parameter) unter Verwendung von fortgesetztem Vortraining auf dem FineWeb-Edu-Datensatz. Sie verglichen die KV-CAT-Checkpoints mit den Basismodellen unter Verwendung modernster nachgelagerter Komprimierungsmethoden (Attention-Matching und gradientenbasierte Optimierung) entlang dreier Achsen:

Leistung ohne Komprimierung (Q1):
- KV-CAT-Modelle behielten auf Standard-Multiple-Choice-Benchmarks (HellaSwag, WinoGrande, PIQA usw.) eine Leistung bei, die mit den Basismodellen vergleichbar war, mit vernachlässigbaren Unterschieden in der Genauigkeit. Dies bestätigt, dass das Trainingsverfahren die Standardfähigkeit nicht zugunsten der Komprimierbarkeit opfert.
Komprimierungsqualität bei festem Budget (Q2):
- Suffix-Perplexität: Unter festen Optimierungsbudgets erzielten KV-CAT-Modelle konsistent eine bessere Erhaltung der Suffix-Perplexität als Basismodelle. Die Verbesserungen reichten bis zu 3,21× in der Perplexitätserhaltung.
- Optimierungsgeschwindigkeit: Bei Verwendung gradientenbasierter Komprimierung benötigten KV-CAT-Modelle bis zu 5× weniger Optimierungsschritte, um dieselbe Komprimierungsqualität wie das Basismodell zu erreichen.
Downstream-Aufgaben (Q3 & Q4):
- Wiederauffinden (Needle-in-a-Haystack): KV-CAT verbesserte die Genauigkeit des exakten Treffens beim Wiederauffinden aus komprimierten Kontexten, mit Gewinnen von bis zu 68% bei bestimmten Beibehaltungsraten (z. B. 30–50% Beibehaltungsraten).
- Fragenbeantwortung in langen Kontexten: Bei LongBench v2-Aufgaben zeigten KV-CAT-Modelle signifikante Verbesserungen der Genauigkeit (bis zu 39%) bei komprimierten Kontexten und übertrafen Basismodelle in verschiedenen Teilbereichen (Recht, Wissenschaft, Agentenverlauf usw.).

5. Wichtige Beiträge

Die Arbeit beansprucht folgende Beiträge:

Charakterisierung der Komprimierbarkeit: Definition der KV-Komprimierbarkeit als Eigenschaft der gelernten Transformer-Repräsentation und nicht nur der Eingabesequenz.
Theoretischer Beweis: Nachweis, dass fast alle Sequenz-zu-Vektor-Funktionen sowohl komprimierbare als auch nicht komprimierbare Transformer-Implementierungen zulassen.
KV-CAT-Rahmenwerk: Vorschlag eines Trainingsverfahrens, das durch Sparsifizierung zur Trainingszeit und Self-Distillation explizit komprimierbare Repräsentationen fördert.
Empirische Validierung: Nachweis, dass KV-CAT das Qualitäts-Budget-Verhältnis nachgelagerter Komprimierungsmethoden bei Wiederauffinden, Perplexität und Fragenbeantwortung in langen Kontexten verbessert, ohne die Leistung ohne Komprimierung zu verschlechtern.

6. Bedeutung und Behauptungen

Die Arbeit positioniert KV-CAT nicht als Ersatz für nachgelagerte Komprimierungsalgorithmen, sondern als ergänzenden Ansatz, der Modelle empfänglicher dafür macht. Die Bedeutung liegt in der Verschiebung des Optimierungsziels vom Komprimierungsalgorithmus selbst hin zu den Trainingsdynamiken des Modells. Indem das Modell während des Trainings Komprimierungsbeschränkungen ausgesetzt wird, argumentieren die Autoren, dass Modelle Repräsentationen lernen, die strukturell robust gegenüber Informationsverlust sind, wodurch eine effektivere Inferenz in langen Kontexten mit reduzierten Speicher- und Rechenkosten ermöglicht wird.

Die Autoren bleiben hinsichtlich der Einschränkungen bescheiden und weisen darauf hin, dass die Methode Trainings-Overhead (zusätzliche Router und duale Durchläufe) und Implementierungskomplexität mit sich bringt und dass die Generalisierung auf Modelle größeren Maßstabs oder andere Domänen weiterer Untersuchungen bedarf.