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

Probleemstelling

Grote Taalmodellen (LLMs) ondervinden aanzienlijke rekenkundige en geheugenkosten bij het verwerken van lange prompts. Dit komt doordat de zelf-attention-mechanisme (self-attention) kwadratisch schaalt met de invoerlengte ($O(n^2)$). Bestaande oplossingen voor tokencompressie richten zich voornamelijk op de token-ruimte (het verwijderen of samenvatten van tokens), maar negeren inefficiënties in de latente embedding-ruimte.

De auteurs wijzen op een fundamenteel probleem: tokens verbruiken veel meer bits dan theoretisch nodig is om ze te identificeren. Bijvoorbeeld, in het QWEN 2.5-model (vocabulaire ~152k) kost een token ~28.672 bits, terwijl de theoretische minimumwaarde slechts ~18 bits is. Er is dus een enorme hoeveelheid redundantie in de huidige token-embeddings die niet wordt benut.

Methodologie: K-Token Merging

Het paper introduceert K-Token Merging, een nieuw raamwerk voor compressie in de latente ruimte. In plaats van tokens te verwijderen, worden groepen van $K$ opeenvolgende tokens samengevoegd tot één enkele embedding.

Kerncomponenten:

1. Lightweight Encoder ($f$): Een klein encoder-model (een MLP van ongeveer 50 MB) dat elke opeenvolgende blok van $K$ token-embeddings in één enkele "compressed embedding" ($C_i$) omzet.
   • Initialisatie: De encoder wordt zo ontworpen dat hij bij start van het trainen de gemiddelde (mean pooling) van de $K$ originele embeddings output. Dit versnelt de convergentie aanzienlijk ten opzichte van willekeurige initialisatie.
2. LoRA-Adaptatie: De LLM wordt fijngefineerd (fine-tuned) met behulp van LoRA (Low-Rank Adaptation) om deze gecomprimeerde embeddings te begrijpen.
3. Workflow:
   • Prefill-fase: De invoer wordt opgesplitst in blokken van $K$ tokens. De encoder verwerkt deze blokken tot één embedding per blok. Dit vermindert de lengte van de sequentie die de attention-laag moet verwerken.
   • Generatie-fase: Het model genereert altijd originele, niet-gecomprimeerde tokens. Deze nieuwe tokens worden toegevoegd aan de mix van gecomprimeerde en niet-gecomprimeerde tokens, waarna de autoregressieve generatie doorgaat.
4. Trainingsdoel: De verliesfunctie (loss) wordt alleen berekend op de posities van de gegenereerde (niet-gecomprimeerde) tokens. Dit zorgt ervoor dat het model leert om de gecomprimeerde invoer te interpreteren zonder de kwaliteit van de output te verliezen.

Verschil met bestaande methoden:

• Hard Prompt Compression: Verwijdert tokens (risico op informatieverlies).
• Soft Prompt Compression (bestaand): Voegt nieuwe meta-tokens toe of leert nieuwe token-sequenties, maar werkt vaak nog steeds op token-niveau.
• K-Token Merging: Werkt direct in de embedding-ruimte, behoudt alle informatie (geen tokens worden verwijderd) en generaliseert goed naar onbekende token-combinaties zonder de vocabulairegrootte explosief te laten toenemen.

Belangrijkste Bijdragen

1. Nieuw Framework: Introductie van K-Token Merging als een methode voor compressie in de latente embedding-ruimte.
2. Trainingsrecept: Een specifieke trainingsschema met LoRA en een "Average-Initialized Encoder" voor snelle convergentie.
3. Pareto-Optimaliteit: Experimenteel bewijs dat de methode een superieure balans biedt tussen compressie en prestaties vergeleken met state-of-the-art baselines.
4. Efficiëntie: Potentieel tot 94% reductie in rekenkosten (FLOPs) in scenario's waar de outputlengte veel korter is dan de invoerlengte, dankzij de kwadratische schaal van attention.

Resultaten

De methode is getest op drie verschillende taken en datasets:

1. Structuurredening (Textualized Tree): Een synthetische dataset om te testen of structurele relaties behouden blijven.
   • Resultaat: Met $K=4$ (75% lengtereductie) werd een nauwkeurigheid van 98,38% behaald (t.o.v. 99,97% voor de oncompressed baseline). Dit is een verbetering van 28,2% op de P-L F1-score ten opzichte van de beste baseline.
2. Sentimentanalyse (Amazon Reviews): Natuurlijke taalverwerking.
   • Resultaat: De 4-Token variant bereikte een lengtereductie van 75% met een nauwkeurigheid van 91,05%. De P-L F1-score was 25,5% hoger dan de sterkste baseline.
3. Code-editing (CommitPackFT): Code-generatie en -aanpassing.
   • Resultaat: De methode behaalde de beste P-L F1-score (0,830) met 75% compressie, terwijl de perplexiteit (verwarring) slechts marginaal toenam.

Conclusie van de resultaten: K-Token Merging ligt consistent op de Pareto-grens (de optimale balans tussen prestatie en compressie) voor alle geteste datasets. Het behoudt bijna alle informatie terwijl de invoerlengte drastisch wordt verkort.

Betekenis en Toekomstperspectief

Dit onderzoek toont aan dat er aanzienlijke inefficiënties bestaan in de huidige token-embeddings van LLMs en dat deze kunnen worden opgelost door compressie in de latente ruimte in plaats van op token-niveau.

• Impact: De methode biedt een praktische oplossing voor het verwerken van lange contexten zonder de kwaliteit van het antwoord te verlagen, wat essentieel is voor toepassingen zoals documentanalyse, code-generatie en complexe redenering.
• Beperkingen & Toekomst:
  • Momenteel wordt alleen de invoer gecomprimeerd, niet de gegenereerde output (wat de totale efficiëntie beperkt bij lange antwoorden).
  • De compressie is uniform ($K$ is constant); adaptieve strategieën (variabele $K$ gebaseerd op frequentie) worden voorgesteld voor toekomstig onderzoek.
  • De tests zijn uitgevoerd op relatief kleine modellen (Qwen 2.5 0.5B); validatie op grotere modellen is nodig.

Samenvattend biedt K-Token Merging een veelbelovende richting om de rekenkosten van LLMs te verlagen door de redundantie in de embedding-ruimte te benutten, zonder de semantische integriteit van de invoer te offeren.