Tensor Cache: Eviction-conditioned Associative Memory for Transformers

Technische Samenvatting: Tensor Cache

Probleemstelling
Autoregressieve Transformer-inferentie is afhankelijk van het cachen van Key-Value (KV)-paren om het herberekenen van prefixes te vermijden. De behouden KV-toestand groeit echter lineair met de contextlengte, diepte en concurrentie, wat een memory-bottleneck creëert. Schuifvenster-caching begrenst het geheugengebruik door alleen de meest recente $W$ tokens te behouden, maar deze aanpak lijdt aan "hard vergeten": zodra een token uit het venster wordt verwijderd, wordt zijn informatie volledig weggegooid, waardoor relevante bewijsvoering buiten het venster ontoegankelijk wordt. Bestaande methoden met een begrenste toestand gooien verwijderde tokens volledig weg of proberen elke token te comprimeren tot een lopend gemiddelde, wat ruis kan introduceren of er niet in slaagt om recent en ver verleden context effectief te onderscheiden.

Methodologie: Tensor Cache (TC)
Tensor Cache introduceert een tweeledige architectuur met een begrenste toestand die KV-verwijdering transformeert van een verwijderingsgebeurtenis naar een schrijfbewerking voor associatief geheugen.

1. Tweeledige Architectuur:

   • L1 (Eerste-niveau cache): Een ringbuffer van vaste grootte voor KV-paren die de meest recente $W$ tokens bevat. Dit ondersteunt exacte softmax-attention voor lokale context.
   • L2 (Tweede-niveau cache): Een matrix $A$ met een vaste grootte voor uitwendige producten (snel-gewicht geheugen) die uitsluitend informatie opslaat van tokens die uit het L1-venster zijn verwijderd.
   • Schrijfmechanisme: Wanneer de ringbuffer een vermelding overschrijft, wordt het verplaatste $(k, v)$-paar niet weggegooid. In plaats daarvan wordt het in de L2-matrix $A$ geschreven via een update van uitwendige producten: $A \leftarrow \lambda A + \eta (k \otimes v)$. Hierbij zijn $\lambda$ (afname) en $\eta$ (schrijfsnelheid) per-head geleerde scalairen.
   • Leesmechanisme: Toekomstige queries $q_t$ krijgen toegang tot het gecomprimeerde verleden door $r_t = q_t A$ te berekenen. Door gebruik te maken van de linear-attention-identiteit $q(k \otimes v) = \langle q, k \rangle v$, haalt deze bewerking effectief informatie op uit alle verwijderde tokens zonder deze individueel op te slaan.
   • Fusie: Een geleerde scalair-gate $g$ fuseert de uitvoer van de lokale attention (L1) en de geheugenuitlezing (L2): $y_t = y_{local} + \sigma(g) m_t$.

2. Trainingsoptimalisatie (Parallelle Scan):
   Het artikel identificeert een kritiek probleem met standaard verkorte methoden voor chunked training. Een naïeve aanpak vat een chunk samen via zijn gemiddelde ($\bar{k} \otimes \bar{v}$), wat $C^2 - C$ spurious cross-token uitwendige producten introduceert ($k_i \otimes v_j$ waar $i \neq j$) die niet bestaan tijdens streaming-inferentie.

   • Oplossing: De auteurs implementeren een parallelle gewogen-som-scan die wiskundig equivalent is aan sequentiële schrijfbewerkingen per token. Dit dicht de kloof tussen training en inferentie binnen float32-epsilon ($< 10^{-7}$ relatieve fout), zodat het model de juiste associatieve dynamiek leert zonder de ruis van cross-token-interferentie.

Belangrijkste Bijdragen

• Verwijderings-geconditioneerd Associatief Geheugen: In tegenstelling tot eerdere uitwendige-productgeheugens (zoals Infini-attention, mLSTM) die elke token opnemen, gebruikt TC de uitwendige-productmatrix strikt als een L2-cache die uitsluitend wordt gevoed door schuifvenster-verwijderingen. Dit partitioneert tokens op een schone manier in "resident in L1" en "gecomprimeerd in L2".
• Correctie van Trainingartefacten: Het artikel identificeert en lost een eerder over het hoofd gezien trainingssnelkoppeling (chunked-mean) op die spurious cross-producten introduceert. De voorgestelde parallelle gewogen-som-scan elimineert deze mismatch tegen dezelfde computatiekosten.
• Efficiëntie met Begrensde Toestand: TC biedt een mechanisme waarbij oudere context voorspellingen beïnvloedt zonder individuele verwijderde vermeldingen te behouden, en biedt zo een praktisch alternatief voor volledige KV-behoud.

Experimentele Resultaten
De auteurs hebben TC geëvalueerd op schaal van systemen, synthetisch associatief herinneren en taalmodellering met lange context op echte tekst (OpenWebText en Shakespeare).

• Synthetisch Associatief Herinneren: In taken die het ophalen van informatie buiten het schuifvenster vereisten, behaalde TC 100% nauwkeurigheid, gelijk aan de prestaties van Full KV. Het verminderde de behouden inferentiestatus met 72–84% vergeleken met Full KV, terwijl het aanzienlijk beter presteerde dan Infini-attention (wat degradeerde naarmate de gaten toenamen) en schuifvenster-baselines (die op willekeurig niveau presteerden).
• Taalmodellering met Lange Context: Op OpenWebText behaalde TC de laagste gemiddelde Negative Log-Likelihood (NLL) bij elke geëvalueerde contextlengte van 1.024 tot 32.768 tokens, met als enige uitzondering $L=2.048$ waar het statistisch niet te onderscheiden was van Infini-attention. Bij $L=32.768$ bereikte TC een NLL van 5,14, beter dan Full KV (6,00) en Infini-attention (5,42), terwijl het 2,4$\times$ minder piek-GPU-geheugen gebruikte dan Full KV.
• Throughput: TC veroorzaakte een overhead van ongeveer 30% in throughput vergeleken met pure schuifvenster-methoden (vanwege L2-lezen/bijwerken), maar bleef consistent sneller dan Infini-attention over alle geteste contextlengtes.
• Capaciteitsdiagnostiek: Experimenten met ruwe capaciteit van associatief geheugen toonden aan dat de prestaties bij het ophalen verslechteren naarmate meer ongerelateerde associaties in de vaste toestand worden gecomprimeerd, wat de begrende aard van het geheugen bevestigt.

Betekenis en Claims
Het artikel beweert dat Tensor Cache succesvol de kloof overbrugt tussen de geheugenefficiëntie van schuifvenster-attention en de herinneringscapaciteiten van full-context attention. Door KV-verwijdering te behandelen als een gestructureerde geheugenbewerking in plaats van een destructieve verwijdering, stelt TC Transformers in staat om relevante bewijsvoering buiten het directe venster te benaderen zonder de lineaire geheugenkosten van volledige KV-caches. De auteurs benadrukken dat hun primaire bijdrage niet de uitvinding van het uitwendige-productgeheugen zelf is, maar de specifieke toepassing ervan als een verwijderings-geconditioneerde L2-cache en de rigoureuze correctie van trainingartefacten die dergelijke methoden eerder belemmerden. TC wordt gepresenteerd als een praktisch, alternatief met een begrenste toestand dat de grens van geheugenkwaliteit verbetert voor inferentie met lange context.