Large Language Model as Token Compressor and Decompressor

Probleemstelling

Grote Taalmodellen (LLM's) hebben een fundamentele beperking in het verwerken van lange contexten. De aandachtmechanismen (attention mechanisms) die ten grondslag liggen aan deze modellen hebben een kwadratische complexiteit ($O(N^2)$) ten opzichte van de sequentielengte. Dit maakt het rekenkundig en geheugenintensief om sequenties van tienduizenden tot honderdduizenden tokens te verwerken, wat schaalbare inferentie, documentbegrip en multi-hop retrieval onpraktisch duur maakt.

Bestaande oplossingen vallen in twee categorieën, die beide tekortschieten:

1. Heuristische methoden: Token-merging of -pruning op basis van gelijkenis of aandachtsscores. Deze kunnen discursieve of pragmatische betekenis verliezen die moeilijk te herstellen is.
2. Lerende/auto-encoder methoden: Methoden zoals AutoCompressor of ICAE die context comprimeren naar vaste lengte, continue vectoren (soft embeddings). Deze zijn vaak niet-lexicaal, moeilijk interpreteerbaar, en vereisen vaak externe netwerken voor decompressie. Ze missen vaak de mogelijkheid tot exacte reconstructie en passen zich niet dynamisch aan aan de informatiedichtheid.

Methodologie

De auteurs introduceren een nieuw inzicht: een bestaande LLM kan fungeren als zowel compressor als decompressor door een "eigen, intern taal" te leren. Het voorgestelde framework bestaat uit drie componenten: een Compressor, een Inferencer (optioneel), en een Decompressor.

1. Het Z-token Concept:
In plaats van continue vectoren, leert het model om tekst te vertalen naar een compacte reeks discrete, variabele lengte latent codes, genaamd Z-tokens.

• Compressor: Mapt een invoersequentie $X$ (lengte $N$) naar een kortere sequentie $Z$ (lengte $K \ll N$). Dit gebeurt autoregressief.
• Variabele Lengte: Semantisch dichte segmenten krijgen meer Z-tokens, terwijl redundante of voorspelbare gebieden agressief worden gecomprimeerd.
• Training: Het model wordt gefine-tuned met LoRA-adapters. Het doel is "self-expressive autoencoding": het model leert tekst te coderen in zijn eigen abstracte taal en deze exact te reconstrueren.
• Discretisatie: Om discrete tokens te genereren tijdens training, wordt een Gumbel-Softmax schatter gebruikt in combinatie met een "straight-through" estimator, zodat gradients kunnen stromen terwijl de output discreet blijft.

2. Twee Gebruiksparadigma's:

• Decompressor als Inferencer: De compressor fungeert als een "learned prompt compressor". De decompressor voert downstream taken (zoals samenvatting of QA) direct uit op basis van de Z-tokens, zonder de originele tekst te reconstrueren.
• Decompressor als Reconstructor: De Z-tokens worden eerst gebruikt voor redenering in de latent space (door een aparte inferencer LLM), waarna de decompressor de originele tekst of het antwoord reconstrueert. Dit koppel compressie en redenering los van elkaar.

3. Sliding Window:
Voor zeer lange teksten wordt de invoer opgesplitst in overlappende vensters. Er wordt een consistentie-beperking (overlap loss) toegepast op de overlappende Z-token sequenties om semantische coherentie te garanderen.

4. Interpretatie van Polysemy:
Een belangrijke observatie is dat Z-tokens polysemisch zijn: dezelfde Z-token-index kan in verschillende contexten verschillende betekenissen hebben. Dit is geen chaos, maar een gestructureerde vorm van abstractie waarbij oppervlaktevariabiliteit wordt weggelaten ten gunste van conceptuele kernbetekenis. De decompressor lost deze ambiguïteit op door gebruik te maken van de volledige Z-sequentie, autoregressieve prefixen en beperkingen tot de basisvocabulary.

Belangrijkste Bijdragen

• Nieuw Paradigma: Het internaliseren van compressie in het redeneringsproces van de LLM zelf, in plaats van het behandelen als een extern, hulpstuk.
• Discrete & Interpretabele Codes: Het introduceren van Z-tokens als discrete, herbruikbare semantische eenheden (vergelijkbaar met morfemen) in plaats van continue vectoren.
• Adaptieve Compressie: Een methode die dynamisch de compressieverhouding aanpast op basis van semantische dichtheid, wat leidt tot een hogere effectieve compressie dan vaste-ratio methoden.
• Exacte Reconstructie: Het vermogen om de originele tekst exact te reconstrueren, wat essentieel is voor verifieerbaarheid en toegang tot broninhoud.

Resultaten

De methode is getest op zes grote datasets (Wikipedia, CNN/DailyMail, HotpotQA, QuALITY, NarrativeQA, QASPER) met verschillende taken (reconstructie, voortzetting, samenvatting, QA).

• Compressie: Bereikt tot 18x tokenreductie op sommige datasets, met een gemiddelde effectieve compressie van 5x tot 9x bij nominale instellingen van 4x of 8x.
• Reconstructie: Bereikt state-of-the-art prestaties op Wikipedia (BLEU-4 van 99.31 bij 4x compressie), wat aantoont dat de Z-tokens voldoende semantische informatie bevatten.
• Downstream Taken:
  • Samenvatting (CNN/DailyMail): Presteert vergelijkbaar of beter dan baselines, met een 2x versnelling in inferentietijd.
  • Lange Context QA: Overtreft bestaande methoden (zoals ICAE, Gist-Token) op alle vier QA-datasets, zelfs bij hoge compressie.
• Efficiëntie:
  • Snelheid: Tot 4.6x snellere inferentie (afhankelijk van de taak en batchgrootte).
  • Geheugen: Significante reductie in GPU-geheugenverbruik (bijv. van 2.81 GB naar 2.04 GB bij een specifieke configuratie).
  • Kwadratische reductie: Door de lengte te reduceren van $N$ naar $K$, daalt de complexiteit van de aandacht van $O(N^2)$ naar $O(K^2)$.

Betekenis en Toekomstperspectief

Dit paper biedt een veelbelovende route naar token-efficiënt redeneren in lange contexten. Door LLM's in staat te stellen hun eigen "korte taal" te spreken, wordt de computatiedruk drastisch verlaagd zonder in te leveren op semantische nauwkeurigheid.

De implicaties zijn groot voor:

• Schalbaarheid: Het mogelijk maken van inferentie op documenten die momenteel te groot zijn voor bestaande modellen.
• Interpreteerbaarheid: Het bieden van een nieuw venster in hoe LLM's semantische kennis intern representeren en manipuleren.
• Toekomstige Werk: De auteurs wijzen op de noodzaak van een mechanisme voor globale coördinatie van Z-tokens over venstergrenzen heen en het maken van compressie "taak-bewust" om specifieke structuren (zoals feiten of argumenten) te behouden.

Kortom, de paper bewijst dat compressie niet hoeft te gaan ten koste van kwaliteit, maar juist kan dienen als een krachtig instrument voor efficiëntere en diepere semantische verwerking.