Compressing Sequences in the Latent Embedding Space: $K$-Token Merging for Large Language Models

1. Problemstellung

Große Sprachmodelle (LLMs) stoßen bei der Verarbeitung langer Eingabekontexte an erhebliche Grenzen hinsichtlich Rechenleistung und Speicherverbrauch. Der Hauptgrund liegt im Self-Attention-Mechanismus (basierend auf Softmax), dessen Komplexität quadratisch mit der Eingabelänge skaliert ($O(n^2)$).

Bestehende Ansätze zur Token-Komprimierung lassen sich in zwei Kategorien einteilen, die jedoch beide Nachteile haben:

• Harte Komprimierung (Hard Prompt Compression): Entfernt oder fasst Token zusammen (z. B. durch Selektion oder Zusammenfassung). Dies führt oft zu Informationsverlusten, besonders bei informationsdichten Aufgaben wie Code-Generierung oder mathematischem Schlussfolgern, wo jedes Token essenziell sein kann.
• Weiche Komprimierung (Soft Prompt Compression): Lernt neue Token oder passt Parameter an, um kompaktere Repräsentationen zu erzeugen. Diese Methoden operieren jedoch fast ausschließlich auf der Token-Ebene und ignorieren die Ineffizienzen im latenten Embedding-Raum.

Ein zentrales Problem ist die Diskrepanz zwischen der theoretischen Informationsdichte und der tatsächlichen Bit-Nutzung: Ein Token aus einem Vokabular von ~152.000 Wörtern benötigt theoretisch nur ~18 Bits zur Identifikation, verbraucht in modernen Embeddings jedoch oft über 28.000 Bits. Dieser enorme Overhead deutet auf signifikante Redundanz im Embedding-Space hin, die bisher ungenutzt blieb.

2. Methodik: K-Token Merging

Das Paper stellt K-Token Merging vor, ein neues Framework zur Komprimierung im latenten Embedding-Raum.

• Kernidee: Statt Token zu löschen, werden aufeinanderfolgende Blöcke von $K$ Token-Embeddings durch einen leichten Encoder ($f$) zu einem einzigen komprimierten Embedding zusammengeführt. Dies geschieht ausschließlich auf der Eingabeseite.
• Architektur:
  • Prefill-Phase: Der Encoder nimmt $K$ originale Token-Embeddings ($T_{iK+1}, \dots, T_{iK+K}$) entgegen und erzeugt ein einziges komprimiertes Embedding ($C_i$).
  • Generierungsphase: Das LLM generiert weiterhin originale, unkomprimierte Token ($G_j$). Diese werden an den gemischten Präfix (bestehend aus komprimierten und generierten Token) angehängt.
  • Anpassung: Das LLM wird mittels LoRA (Low-Rank Adaptation) gemeinsam mit dem Encoder feinabgestimmt, um die komprimierten Eingaben zu verstehen, ohne das Vokabular zu erweitern.
• Initialisierung: Ein entscheidender technischer Aspekt ist die Initialisierung des Encoders. Der Encoder wird so konstruiert, dass er zunächst als Durchschnitt (Mean Pooling) der $K$ Token-Embeddings agiert. Dies wird durch eine kleine MLP-Schicht erreicht, deren Gewichte nahe Null initialisiert werden. Dies sorgt für eine stabile und semantisch fundierte Startposition, was die Konvergenzgeschwindigkeit im Vergleich zur zufälligen Initialisierung deutlich verbessert.
• Unterschied zur Vokabular-Erweiterung: Im Gegensatz zu Ansätzen, die $K$-Gramme als neue Token ins Vokabular aufnehmen (was zu einer Explosion seltener „Tail-Token" führt), bleibt das Vokabular unverändert. Der Encoder generalisiert auf ungesehene Token-Kombinationen.

3. Wichtige Beiträge

1. Neues Framework: Einführung von K-Token Merging als erste Methode, die Redundanz direkt im latenten Embedding-Raum ausnutzt, anstatt Token zu löschen oder das Vokabular zu erweitern.
2. Trainings-Rezept: Vorstellung einer effizienten Trainingsstrategie unter Verwendung von LoRA und einer speziellen Initialisierung des Encoders (Average-Initialized Encoder), die eine schnelle Konvergenz ermöglicht.
3. Experimentelle Validierung: Umfassende Evaluation auf drei unterschiedlichen Aufgaben, die zeigen, dass die Methode die Pareto-Grenze zwischen Kompressionsrate und Leistung erreicht.

4. Ergebnisse

Die Methode wurde auf drei Datensätzen evaluiert:

1. Textualized Tree (Strukturelles Schlussfolgern): Eine synthetische Aufgabe zur Erkennung von Eltern-Kind-Beziehungen in Bäumen.
   • Ergebnis: Mit $K=4$ (75% Reduktion der Eingabelänge) wurde eine Genauigkeit von 98,38% erreicht (nur 1,59% Abfall gegenüber dem unkomprimierten Baseline von 99,97%).
   • P-L F1 Score: 0,851 (bester Wert aller Methoden).
2. Amazon Reviews (Sentiment-Analyse): Klassifikation von Produktbewertungen.
   • Ergebnis: Bei 75% Kompression ($K=4$) wurde eine Genauigkeit von 91,05% erreicht. Der P-L F1 Score betrug 0,822 (25,5% besser als der beste Baseline).
3. CommitPackFT (Code-Editing): Generierung von Code-Updates basierend auf Anweisungen.
   • Ergebnis: Bei 75% Kompression wurde der beste P-L F1 Score von 0,830 erzielt, trotz einer leichten Erhöhung der Perplexität.

Zusammenfassung der Leistung:

• Kompressionsrate: Bis zu 75% Reduktion der Eingabelänge (bei $K=4$).
• Recheneffizienz: Da der Prefill-Teil ($O(n^2)$) auf 25% der Länge reduziert wird, sinkt der Rechenaufwand für den Prefill auf ca. 6,25% des Originals. Bei kurzen Ausgaben führt dies zu einer geschätzten 94%igen Reduktion des Gesamtrechenaufwands.
• Vergleich: Die Methode übertrifft sowohl harte Komprimierungsmethoden (die bei hohen Raten stark an Genauigkeit verlieren) als auch weiche Methoden (die oft geringere Kompressionsraten erreichen).

5. Bedeutung und Ausblick

Die Arbeit zeigt, dass die Ineffizienz im Embedding-Raum eine bisher unterschätzte Quelle für Optimierung darstellt. K-Token Merging bietet einen vielversprechenden Weg, um LLMs für Long-Context-Anwendungen effizienter zu machen, ohne die generative Qualität oder das Vokabular zu beeinträchtigen.

Einschränkungen und zukünftige Richtungen:

• Die Komprimierung erfolgt derzeit nur auf der Eingabeseite, nicht während der Generierung.
• Die Kompressionsrate ist fest ($K$ ist konstant); adaptive Strategien (basierend auf der Häufigkeit von K-Grammen) könnten die Effizienz weiter steigern.
• Die Evaluation beschränkte sich auf kleinere Modelle (Qwen 2.5 0.5B); die Generalisierbarkeit auf größere Modelle muss noch untersucht werden.

Insgesamt etabliert das Paper einen neuen Paradigmenwechsel weg von der Token-Auswahl hin zur Embedding-Optimierung, um die Skalierbarkeit von LLMs zu verbessern.