Context Memorization for Efficient Long Context Generation

Technische Samenvatting: Attention-State Memory voor Efficiënte Generatie met Lange Context

Probleemstelling
Moderne toepassingen van Large Language Models (LLM) vertrouwen steeds meer op lange conditionering-prefixen (bijvoorbeeld contextvoorbeelden, opgehaalde documenten of agentische instructies) om inferentie te sturen. Prefix-verrijkte inferentie stuit echter op twee primaire structurele beperkingen:

1. Prefix Decay: Naarmate de generatie vordert, verdeelt de aandacht van het model zich over nieuw gegenereerde tokens, waardoor de invloed van de prefix afneemt, met name in scenario's met lange context.
2. Inefficiëntie bij Inferentie: De berekening van aandacht over de prefix schaalt lineair met de lengte ervan ($O(L)$), wat aanzienlijke latentie en geheugenoverhead oplevert tijdens zowel de prefill-fase als elke decode-stap. Hoewel prefix-caching de prefill-kosten amortiseert, blijft het aanzienlijke geheugengebruik met zich meebrengen.

Bestaande oplossingen vallen in twee categorieën, beide met nadelen:

• Prefix Internalisatie: Methoden die modelparameters fine-tunen of hypernetwerken gebruiken om prefix-gedrag te internaliseren. Deze vereisen op gradiënten gebaseerde training, waardoor ze resource-intensief zijn, traag bij te werken zijn en ongeschikt voor dynamische prefix-wijzigingen.
• Prefix Compressie: Methoden die de prefix in het aandachtsmechanisme houden maar deze comprimeren (bijvoorbeeld via token-pruning of het verwijderen van KV-cache). Deze vereisen nog steeds aandacht voor de gecomprimeerde prefix tijdens inferentie, wat betekent dat de latentie schaalt met de gecomprimeerde lengte en de prefix-invloed onderhevig blijft aan decay.

Methodologie: Attention-State Memory (ASM)
De auteurs stellen Attention-State Memory voor, een trainingsvrije aanpak die de prefix externaliseert naar een lichtgewicht, op opzoeken gebaseerd geheugen van vooraf berekende aandachtstoestanden. Deze methode maakt gebruik van de online-softmax-identiteit om aandacht over een samengevoegd sleutelblok te decomponeren in aandacht over sub-blokken.

Het kerninzicht is dat voor een gegeven query de aandachtstoestand $(a, Z)$ – bestaande uit de aandachtsoutput en de normalisatieterm – voldoende is om de bijdrage van de prefix aan de aandacht te reconstrueren zonder opnieuw aandacht te besteden aan de oorspronkelijke sleutels en waarden.

Het framework werkt in twee fasen:

1. Offline Kalibratiefase (Geheugenconstructie):

   • Verzameling: Het model voert forward passes uit over een set representatieve query-response traces die zijn samengevoegd met de prefix. Voor elke query-token in de respons registreert het systeem de query-vector en de bijbehorende aandachtstoestand $(a, Z)$ over de prefix.
   • Clustering: De verzamelde query-vectoren worden geclusterd (bijvoorbeeld via K-means) tot $K$ centroiden.
   • Aggregatie: Voor elke cluster wordt een enkele representatieve invoer gemaakt. De centroid-query is het gemiddelde van de queries van de cluster. De aandachtstoestand $(\bar{a}, \bar{Z})$ wordt berekend via aandachtsbewuste aggregatie, waarbij de toestanden van alle leden in de cluster worden samengevoegd met behulp van de online-softmax-updateregel. Dit zorgt ervoor dat de centroid een verliesloos gemiddelde vertegenwoordigt van de prefix-aandacht van de cluster.
   • Efficiëntie: Om zeer lange prefixes (bijvoorbeeld 16K+ tokens) te verwerken zonder het GPU-geheugen te overschrijden, kan de prefix worden opgedeeld in chunks, onafhankelijk worden gecodeerd en binnen het geheugen worden samengevoegd met behulp van het samenstellingskarakter van de online-softmax-identiteit.

2. Online Inferentiefase:

   • Retrieval: Voor een inkomende query identificeert het systeem de dichtstbijzijnde geheugencentroïde met behulp van cosinus-ähnelijkheid (optioneel met whitening-transformaties en RoPE-handling).
   • Samenvoegen: De opgehaalde vooraf berekende aandachtstoestand $(\bar{a}, \bar{Z})$ wordt samengevoegd met de self-attention-output van de query (berekend over niet-prefix tokens) met behulp van de samenvoegoperator die wordt gedefinieerd door de online-softmax-identiteit. Dit herstelt de exacte aandachtsoutput die zou zijn geproduceerd als de prefix aanwezig was, zonder aandacht te besteden aan de prefix-tokens.
   • Complexiteit: De opzoekkosten schalen logaritmisch ($O(\log K)$) met de geheugengrootte via hiërarchische indexering, onafhankelijk van de oorspronkelijke prefix-lengte.

Belangrijkste Bijdragen

1. Trainingsvrije Externalisatie: Het artikel stelt een methode voor om lange prefixes te externaliseren naar een compact woordenboek van geheugen zonder gradiënt-gebaseerde training of auxiliaire modellen.
2. Uitbreiding van Online-Softmax: De auteurs breiden de online-softmax-identiteit uit van efficiënte aandachtsberekening naar het domein van cross-query prefix-hergebruik, waardoor verliesloze reconstructie van prefix-aandacht via retrieval mogelijk wordt.
3. Prestaties en Efficiëntie: Experimenten tonen aan dat ASM de prestaties van volledige aandacht evenaart of overtreft, terwijl de rekenkosten aanzienlijk worden verlaagd.

Experimentele Resultaten
De methode werd geëvalueerd op ManyICLBench (In-Context Learning) en RuleArena/NBA (Retrieval-Augmented Generation) met LLaMA-3.1-8B.

• In-Context Learning (ICL): Op ManyICLBench verbeterde ASM de nauwkeurigheid ten opzichte van standaard in-context learning over geheugenbudgetten van 1K tot 8K invoeren. Bij een geheugenbudget van 8K verminderde het de aandachtslatentie met 1,36× ten opzichte van volledige aandacht. ASM presteerde beter dan KV-cache compressie-baselines (KVzip), die vaak kritieke informatie zoals label-tokens verloren lieten gaan.
• Retrieval-Augmented Generation (RAG): Op de NBA-benchmark (met ~20K tokens aan regels) overtrof ASM de prestaties van volledige aandacht-RAG met slechts 20% van het geheugenvoetafdruk. De optimale prestaties werden waargenomen bij een intermediaire codeboekgrootte (4K invoeren), wat suggereert dat het aantal invoeren een taakspecifieke hyperparameter is in plaats van een die simpelweg gemaximaliseerd moet worden.
• Latentie: Hoewel hiërarchische opzoeken bij zeer kleine geheugenbudgetten hogere overhead met zich meebrengen, is de groei van de latentie sub-lineair. ASM wordt sneller dan volledige aandacht bij ongeveer 4K geheugenvoeren en bereikt een 1,8× snelheidswinst bij 16K invoeren.

Betekenis en Claims
Het artikel claimt dat Attention-State Memory de structurele kosten van generatie met lange context aanpakt door prefixkennis te ontkoppelen van per-query aandachtsberekening. Door de prefix te externaliseren naar een vooraf berekend, op opzoeken gebaseerd geheugen, elimineert de methode de lineaire schaling van inferentie-latentie met prefix-lengte en verwijdert het het prefix-decay-probleem dat geassocieerd wordt met standaard aandachtsmechanismen.

De auteurs positioneren dit werk als een stap in de richting van een bredere onderzoeksrichting: het externaliseren van LLM-kennis naar compacte, herbruikbare representaties die bestaan buiten tekstprompts en modelparameters. Ze benadrukken dat hun aanpak een gunstige afweging biedt tussen nauwkeurigheid en efficiëntie, en een trainingsvrij alternatief biedt voor parameter-internalisatie en een efficiënter alternatief voor prefix-compressie.