Technische Samenvatting: End-to-End Contextcompressie op Schaal

Probleemstelling

Inference met lange context in Large Language Models (LLMs) wordt steeds vaker beperkt door geheugen en latentie vanwege de lineaire groei van de Key-Value (KV) cache met de contextlengte. Hoewel recente technieken voor KV-cachecompressie proberen dit te verzachten door cache-items te verwijderen of compacte caches te trainen, kampen zij met significante beperkingen:

1. Kwaliteit versus Efficiëntie-trade-off: Veel methoden verslechteren de prestaties van het model aanzienlijk of vereisen domeinspecifieke training die niet generaliseert.
2. Compatibiliteit met Inference Engines: Technieken die items niet-uniform over attention-heads of lagen verwijderen, zijn moeilijk te integreren met moderne productie-inference engines (bijv. vLLM, SGLang) die uitgaan van gedeelde sequentielengtes.
3. Prefill Overhead: Veel benaderingen vereisen dat de volledige context eerst is geprefilled voordat compressie plaatsvindt, wat potentiële winst in latentie tenietdoet, of vereisen aanzienlijke extra rekenkracht.
4. Beperkingen van Soft-Tokens: Bestaande encoder-decoder "soft-token" methoden vertrouwen vaak op offline preprocessing, falen in het behouden van het brede in-context gedrag van het basismodel, of worden alleen geëvalueerd op smalle, domeinspecifieke taken in plaats van informatie-dichte long-context benchmarks.

De centrale vraag die wordt behandeld is of een algemene, taak-agnostische compressor getraind kan worden om de oorspronkelijke capaciteiten van een model te behouden terwijl het geheugengebruik en de inferentietijd aanzienlijk worden verminderd.

Methodologie

De auteurs introduceren Latent Context Language Models (LCLMs), een familie van end-to-end getrainde encoder-decoder compressors. De kernarchitectuur brengt een sequentie van input-tokens $x_{1:T}$ in naar een kortere sequentie van latente embeddings $z_{1:M}$ via een encoder, die vervolgens door een adapter worden geprojecteerd en door een decoder worden geconsumeerd in plaats van de oorspronkelijke context.

Architectuurontwerp

Door middel van een grootschalige architectuurzoektocht vanaf nul, hebben de auteurs optimale ontwerpkeuzes geïdentificeerd:

• Pooling: Mean pooling bleek superieur aan token-gebaseerde (EOS/CLS) pooling en concurrerend met concatenatie, vooral bij hoge compressieratio's (16×).
• Encoding Granulariteit: Een encoder window size ($W$) van 1024 tokens werd geadopteerd, waardoor de encoder over een grotere lokale context kan attentionen dan de compressie blokgrootte ($N$) zonder excessieve geheugenoverhead.
• Attention Masking: Causal masking binnen de encoder presteerde consequent beter dan bidirectionele masking.
• Adapter: Een lichtgewicht MLP adapter (zonder self-attention) bleek effectiever en efficiënter dan attention-gebaseerde adapters voor het projecteren van encoder latents naar de decoder dimensie.
• Initialisatie: Het initialiseren van de encoder vanuit een vooraf getraind embedding model (Qwen3-Embedding-0.6B) in plaats van een standaard LLM leverde betere initiële representaties op, hoewel het verschil kleiner werd bij uitgebreide training.

Trainingsrecept

De auteurs hebben een multi-stage trainingsrecept ontwikkeld om catastrofale vergetelheid te voorkomen en stabiliteit te garanderen:

1. Stage 0 (Adapter Warmup): Bevries de encoder en decoder; train alleen de adapter.
2. Stage 1 (Encoder Training): Ontvries de encoder; houd de decoder bevroren.
3. Stage 2 (Continual Pre-training): Ontvries de decoder met een kleine learning rate; train end-to-end op een mix van next-token prediction en reconstructie data.
4. Stage 3 (Supervised Fine-Tuning): Train op SFT data om instructie-opvolging en redeneren met gecomprimeerde inputs te verbeteren.

De trainingsdata bevat een zorgvuldig samengestelde mix van code, tekst, redenering en long-context documenten. Cruciaal is dat de pre-training data een interleaved format gebruikt waarbij gecomprimeerde segmenten (verpakt in speciale tokens) afwisselen met ongecomprimeerde segmenten, wat het model dwingt om te conditioneren op latente context op meerdere posities in plaats van alleen aan het begin.

Agent Integratie

De auteurs stellen een agentisch systeem voor waarbij de LCLM agent de volledige gecomprimeerde context ontvangt. De agent kan selectief een EXPAND tool aanroepen om de originele ongecomprimeerde tekst voor specifieke segmenten op te halen, wat een "skim-and-zoom" strategie mogelijk maakt voor langdurige taken.

Belangrijkste Resultaten

De LCLM-familie (specifiek een 0.6B encoder en 4B decoder) werd geëvalueerd bij compressieratio's van 4×, 8× en 16× op benchmarks inclusief RULER, LongBench, LongHealth en GSM8K.

• Pareto Frontier: LCLMs vestigen een nieuwe Pareto frontier op de nauwkeurigheid-efficiëntie trade-off. Ze bereiken een aanzienlijk snellere Time-to-First-Token (TTFT) en een lager piek GPU-geheugengebruik vergeleken met KV cache compressie baselines (bijv. SnapKV, KVzip, Attention Matching), vooral bij hoge compressieratio's.
• Behoud van Nauwkeurigheid: Bij 16× compressie behouden LCLMs een hoge nauwkeurigheid op long-context benchmarks (bijv. ~65% op RULER 16k vs. ~93% voor volledige KV, wat aanzienlijk beter is dan baselines die vaak onder de 20-40% zakken).
• Geheugenefficiëntie: In tegenstelling tot KV cache methoden waarbij het geheugen schaalt met de volledige contextlengte, vlakt het LCLM geheugengebruik af bij langere contexten omdat de encoder vaste window-groottes verwerkt. Dit stelt LCLMs in staat om contexten tot 1M tokens aan te kunnen op een enkele H200 GPU, terwijl baselines out-of-memory gaan bij 512K of 1M tokens.
• Fijnmazige Prestaties: Op GSM8K (korte, dichte contexten) bereiken LCLMs de hoogste nauwkeurigheid over alle compressieratio's, wat aantoont dat ze niet beperkt zijn tot long-context QA maar ook informatie-dichte inputs effectief kunnen comprimeren.
• Agent Prestaties: Het agentische systeem met adaptieve expansie verbetert de exact string-match nauwkeurigheid op "needle-in-a-haystack" taken aanzienlijk, en matcht vaak de prestaties van het ongecomprimeerde model terwijl de efficiëntie van de gecomprimeerde context behouden blijft.

Betekenis en Claims

Het artikel claimt dat LCLMs een praktische, schaalbare oplossing bieden voor long-context inference die compatibel is met standaard productie-inference engines. Door de compressie-last te verschuiven naar een kleinere, paralleliseerbare encoder, ontkoppelen LCLMs de geheugen- en latentiekosten van de contextlengte van het generatieproces van de decoder.

De auteurs benadrukken dat hun aanpak:

1. Algemene Capaciteiten Behoudt: In tegen tegenover de meeste eerdere soft-token methoden die prestaties verslechteren of specifieke tuning vereisen, behouden LCLMs sterke prestaties over algemene taken, redeneren en instructie-opvolging.
2. Adaptieve Retrieval Mogelijk Maakt: De integratie met agentische frameworks maakt dynamische expansie van relevante context mogelijk, wat de kloof overbrugt tussen globale corpus zichtbaarheid en fijnmazige informatie-retrieval.
3. Schaalbaar Is: De methode schaalt effectief naar miljoen-token contexten op commodity hardware (H200 GPU's), waarmee de geheugenbottlenecks worden overwonnen die momenteel de implementatie van long-horizon LLM systemen beperken.

Het werk suggereert dat geleerde, end-to-end compressie een levensvatbaar alternatief is voor heuristische KV cache verwijdering, en een robuuste basis biedt voor toekomstige long-horizon agents met persistente werkgeheugens.