🤖 machine learning

Looped Latent Attention: Cross-Loop KV Compression for Looped Transformers

Dit artikel introduceert Looped Latent Attention (LLA), een post-training codec die de low-rank structuur van loop-geïndexeerde KV-caches in weight-tied Transformers exploiteert om extreme compressie (tot 32x) te bereiken en de batchcapaciteit aanzienlijk te vergroten, terwijl bijna verliesloze prestaties worden behouden door middel van SVD-initialisatie en distillatie.

Oorspronkelijke auteurs: James O' Neill, Fergal Reid

Gepubliceerd 2026-07-20

📖 7 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: James O' Neill, Fergal Reid

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Het Geheugenprobleem in AI-hersenen

Stel je voor dat je een robot probeert te leren om een verhaal te schrijven. De robot heeft een brein gemaakt van miljoenen kleine schakelaars (parameters) die weten hoe woorden bij elkaar passen. Maar hier is de crux: om een lang verhaal te schrijven, kan de robot niet alleen naar zijn brein kijken; hij heeft een kladblok nodig. Elke keer dat hij een nieuw woord schrijft, moet hij zich alles herinneren wat hij net heeft geschreven, zodat hij zichzelf niet herhaalt of het plot niet kwijtraakt. In de wereld van Kunstmatige Intelligentie wordt dit kladblok de KV-cache genoemd. Het is een privé-notitieblok voor elk afzonderlijk gesprek, en het wordt groter met elk woord dat de robot spreekt.

Stel je nu een speciaal soort robot voor die ontworpen is om super efficiënt te zijn. In plaats van een enorm brein met veel verschillende kamers (lagen) voor verschillende taken, heeft deze robot één kleine, slimme kamer die hij steeds opnieuw bezoekt. Hij gebruikt dezelfde set instructies opnieuw door erdoorheen te loopen om dieper na te denken. Dit wordt een Looped Transformer genoemd. Het bespaart ruimte op het brein zelf, wat geweldig is. Maar er is een verraderlijk probleem: ook al hergebruikt de robot dezelfde kamer, hij schrijft nog steeds een nieuw pagina in zijn kladblok elke keer dat hij een loop uitvoert. Als de robot vier keer door een loop gaat om een wiskundeprobleem op te lossen, eindigt hij met vier aparte pagina's met aantekeningen voor hetzelfde moment in het verhaal. Dit betekent dat de "geloopte" robot nog steeds een enorme hoeveelheid geheugen nodig heeft om te draaien, wat het doel van efficiëntie tenietdoet. De grote vraag voor wetenschappers is: kunnen we dat kladblok verkleinen zonder dat de robot vergeet hoe hij moet denken?

De Ontdekking van het Papier: Een Geheime Afkorting in de Notities

Dit papier introduceert een slimme truc genaamd Looped Latent Attention (LLA) om dat geheugenprobleem op te lossen. De onderzoekers ontdekten dat de aantekeningen van de robot niet zo rommelig zijn als ze lijken. Wanneer de robot door zijn denkproces loopt, veranderen de aantekeningen die hij voor hetzelfde woord maakt niet wild; ze volgen een vloeiend, voorspelbaar pad, zoals een knikker die een zachte heuvel afrolt. In plaats van elke enkele pagina van het kladblok van de robot op te slaan, slaat de nieuwe methode slechts een kleine "samenvatting" van dat pad op en een set instructies over hoe de specifieke pagina's weer opgebouwd kunnen worden wanneer de robot ze daadwerkelijk nodig heeft om te lezen.

Denk er zo over na: Stel je voor dat je een film kijkt waarin het hoofdpersonage door een gang loopt. In een normale computer zou je bij elke stap die ze zetten een volledige, high-definition foto van het personage opslaan. Dat neemt een enorme hoeveelheid opslag in beslag. Maar met LLA besef je dat het personage gewoon in een rechte lijn loopt. Dus in plaats van duizenden foto's op te slaan, sla je één foto van het personage op en een simpele notitie die zegt: "Ga stap 2 één inch naar voren, 2 inches voor stap 3." Wanneer de computer het personage bij stap 3 nodig heeft, tekent hij die afbeelding snel met behulp van de notitie. Dit is precies wat het papier doet voor AI: het comprimeert het "loop"-gedeelte van het geheugen, waardoor een stapel van vier pagina's wordt omgezet in één enkele, kleine samenvatting die direct kan worden uitgebreid naar de volledige pagina's.

De onderzoekers testten dit op verschillende AI-modellen, waaronder een genaamd Ouro-1.4B. Ze ontdekten dat deze methode ongelooflijk goed werkt. Sterker nog, op een krachtige computerchip genaamd H200 waren ze in staat om 768 verschillende gesprekken tegelijkertijd in het geheugen te passen, terwijl ze voorheen slechts 32 konden passen. Dat is een enorme sprong in hoe meer mensen de AI tegelijkertijd kunnen gebruiken. Ze lieten ook zien dat deze truc werkt, zelfs wanneer de AI zeer lange, moeilijke wiskundeproblemen oplost, hoewel het een beetje extra training nodig heeft om ervoor te zorgen dat het niet in de war raakt wanneer het zijn eigen verhaal schrijft in plaats van alleen maar een leraar te kopiëren.

Waar dit Papier "Nee" tegen zegt

Het is belangrijk om te weten wat deze methode niet is. De onderzoekers testten eerst een heel simpel idee: "Wat als we gewoon de eerste drie pagina's met aantekeningen weggooien en alleen de allerlaatste pagina bewaren?" Ze dachten misschien dat de robot na een paar loops tot een definitief antwoord komt, dus dat de eerdere aantekeningen er niet toe deden. Ze probeerden dit, en het was een ramp. De AI vergat volledig hoe hij wiskunde moest doen, en scoorde nul op tests. Dit bewijst dat het "pad" dat de aantekeningen afleggen belangrijk is; je kunt niet alleen het eindpunt pakken. Je moet de reis vastleggen.

Ze vergeleken hun methode ook met andere manieren om geheugen te verkleinen, zoals het samendrukken van de aantekeningen (kwantisatie) of het delen van aantekeningen tussen verschillende delen van het brein (head-axis compressie). Hoewel die methoden een beetje helpen, laat het papier zien dat het comprimeren van de "loop"-as het geheime ingrediënt is. Als je probeert de andere delen in plaats van de loops te comprimeren, verliest de AI zijn vermogen om te redeneren. Het papier is heel duidelijk: de loop is het meest redundante deel van het geheugen, en dat is de enige plek waar je moet knijpen.

Hoe zeker zijn ze?

De auteurs zijn vrij zelfverzekerd over deze resultaten omdat ze ze direct hebben gemeten. Ze hebben niet alleen gegokt; ze hebben de AI getest op echte taken zoals het oplossen van wiskundeproblemen (GSM8K) en het schrijven van code (HumanEval). Ze lieten zien dat met een 4x compressie (het geheugen 4 keer kleiner maken) de AI nog steeds bijna net zo goed presteert als de originele, ongecomprimeerde versie. Wanneer ze de extreme compressie van 21,3x opzochten, kon de AI nog steeds veel taken aan, hoewel het begon te worstelen met de zeer lange, complexe redeneerstappen.

Ze testten dit ook op een ander type AI-model genaamd Huginn-3.5B, dat een andere hersenstructuur heeft en tot 32 keer loopt. Zelfs daar werkte de methode, wat bewijst dat dit geen toevalstreffer is voor één specifieke robot. Het papier suggereert dat deze "low-rank trajectory" (het vloeiende pad van de aantekeningen) een fundamentele eigenschap is van hoe deze loepende robots denken. Hoewel ze bewezen hebben dat het werkt voor het besparen van geheugen en het vergroten van de capaciteit, merken ze op dat voor de absolute beste prestaties op zeer lange taken, de AI een tweede ronde training nodig heeft waarbij het oefent met zijn eigen gecomprimeerde aantekeningen, in plaats van alleen maar een leraar te kopiëren. Deze "on-policy" training herstelt de kleine fouten die optreden wanneer de robot zijn eigen verhaal begint te schrijven.

Kortom, dit papier laat zien dat we loepende AI-modellen veel efficiënter kunnen maken door te beseffen dat hun geheugen op een specifieke manier repetitief is. Door een kleine samenvatting van die herhaling op te slaan in plaats van de volledige aantekeningen, kunnen we honderden keren meer gesprekken op een enkele computerchip passen, waardoor krachtige AI beschikbaar komt voor meer mensen zonder dat er superdure hardware nodig is.

Technische Samenvatting: Looped Latent Attention (LLA)

Probleemstelling

In Looped, weight-tied Transformers wordt een enkele blok toegepast herhaaldelijk (recurrentie) om het aantal parameters te verminderen. Hoewel dit de gewichten van het model effectief comprimeert, blijft de geheugenbottleneck tijdens inferentie de Key/Value (KV) cache. In standaard looped inferentie berekent het model de aandacht bij elke recurrentiestap opnieuw, waarbij voor elke stap, laag en kop een aparte KV-cache wordt opgeslagen. Bij een model met $T$ loops over $D$ lagen resulteert dit in een KV-cache die gelijk is aan een decoder met $T \times D$ lagen, wat de voordelen van weight tying tenietdoet.

Het artikel stelt dat deze loop-geïndexeerde cache sterk gestructureerd is. Voor een vaste token, laag en kop volgen de K/V-vectoren over de recurrentiestappen heen een kort, laag-dimensionaal traject (low-rank trajectory), in plaats van een hoog-dimensionale ruimte te verkennen. De redundantie is specifief voor de loop-as; de spectra langs de kop-as en de laag-as blijven aanzienlijk vlakker (hogere rang). De uitdaging bestaat erin deze recurrentie-as te comprimeren zonder het traject te laten instorten tot één enkel vast punt, wat de prestaties zou verslechteren.

Methodologie: Looped Latent Attention (LLA)

De auteurs introduceren Looped Latent Attention (LLA), een post-training cache-codec die de $T$ per-loop K/V-vectoren vervangt door een compacte latente representatie en loop-specifieke reconstructiekaarten.

Kernmechanisme

Latente Opslag: In plaats van $T$ $T$ vectoren ( $x_1, \dots, x_T$ $x_{1}, \dots, x_{T}$ ) op te slaan voor een specifieke token, laag en kop, slaat LLA een enkele gecentreerde laag-dimensionale latente vector $c \in \mathbb{R}^r$ $c \in R^{r}$ op.
- $c = (x - \mu)W_{\text{down}}$
- Hierbij is $\mu$ het gemiddelde van een kalibratieset, en $W_{\text{down}}$ een down-projectie matrix.
Loop-Specifieke Reconstructie: De specifieke K/V-vector voor recurrentiestap $t$ $t$ wordt gereconstrueerd met behulp van een loop-specifieke up-projectie:
- $\hat{x}_t = c W_{\text{up},t}^\top + \mu_t$
Asymmetrische Ranks: De methode gebruikt aparte ranks voor Keys ( $r_k$ ) en Values ( $r_v$ ). Empirische analyse toont aan dat Value-vectoren langzamer stabiliseren over de loops dan Key-vectoren, wat motiveert dat $r_v > r_k$ .
LLA-2D Variant: Voor extreme compressie vouwt LLA de kop-as in de latente ruimte, waarbij één latent per token en laag wordt opgeslagen om alle $H$ koppen te reconstrueren. Dit vermindert de cachegrootte met een factor $H$ , maar koppelt de koppen aan elkaar.

Training en Initialisatie

SVD Initialisatie: De codec wordt geïnitialiseerd via Singular Value Decomposition (SVD) van de teacher-activaties op een kalibratieset. De down-projectie wordt ingesteld op de top- $r$ rechtse singuliere vectoren. Deze initialisatie is cruciaal; willekeurige initialisatie resulteert in aanzienlijk hogere KL-divergentie.
Distillatie: De codec wordt verfijnd door de teacher-modellen te bevriezen en de forward KL-divergentie tussen de distributie van de teacher en de geswapped model-distributie te minimaliseren, aangevuld met een attention-output matching term.
On-Policy Verfijning: Om exposure bias in lange autoregressieve rollouts (waarbij het model conditioneert op zijn eigen gegenereerde prefixes in plaats van die van de teacher) aan te pakken, wordt een tweede fase van on-policy distillatie toegepast. De codec wordt bijgewerkt met prefixes die gesampled zijn uit de student-modellen zelf, wat de prestaties bij lange generaties stabiliseert.

Belangrijkste Bijdragen

Definitie van LLA: Een cross-loop latente K/V-codec voor weight-tied Transformers die een compressieratio $\rho = \frac{2Td_{\text{head}}}{r_k + r_v}$ bereikt. Het slaat $(r_k + r_v)H$ scalairen per token per laag op in plaats van $2Td_{\text{head}}$ .
As-Specifieke Compressie: Het artikel demonstreert dat de recurrentie (loop)-as de meest comprimeerbare as is in looped modellen (Ouro-1.4B, Ouro-2.6B-Thinking, en Huginn-3.5B). Het comprimeren van de loop-as presteert beter dan het comprimeren van de kop-as (MLA), de laag-as (cross-layer sharing), of het toepassen van kwantisatie bij gelijke cache-budgetten.
Afwijzing van Fixed-Point Collapse: De auteurs tonen aan dat het loop-traject niet kan worden vervangen door een enkele staat (de laatste loop). Een zero-parameter controle die alleen de K/V van de laatste loop hergebruikt, laat de generatieprestaties instorten (GSM8K nauwkeurigheid daalt naar 0.000), wat bewijst dat het "korte laag-dimensionale traject" noodzakelijk is.
Impact op Deployment: De exacte scalaire reductie vertaalt zich direct in een grotere serving-capaciteit. Op een enkele H200 GPU verhoogt LLA de maximale batch size voor Ouro-1.4B bij 4k context van 32 naar 768 sequenties (bij 21.3× compressie).

Experimentele Resultaten

Compressie vs. Nauwkeurigheid

GSM8K Prestaties: Op Ouro-1.4B evenaart lichte compressie (1.33× tot 2×) de nauwkeurigheid van de teacher (ongeveer 0.79–0.80). Bij 4× compressie behoudt LLA een nauwkeurigheid van 0.59, wat aanzienlijk beter is dan kop-as compressie (0.52) en hergebruik van de laatste loop (0.00).
Schaalbaarheid: Op Ouro-2.6B-Thinking blijft LLA bijna lossless (0.81–0.83) bij 4× tot 10.67× compressie, terwijl kop-as baselines aanzienlijk aan nauwkeurigheid verliezen.
Transfer naar Huginn: De methode is overdraagbaar naar Huginn-3.5B (32 recurrentiestappen). Een training-vrije SVD-codec bereikt bijna lossless fideliteit tot 32× in decoder-onafhankelijke metrieken.

Stabiliteit bij Lange Generaties

MATH-500: Lange wiskundige rollouts zijn gevoelig voor prefix drift. Off-policy training (met gebruik van teacher-prefixes) leidt tot degradatie van de nauwkeurigheid en hoge "no-answer" ratio's naarmate de generatielengte toeneemt.
On-Policy Winst: Het toepassen van on-policy verfijning verhoogt de MATH-500 nauwkeurigheid van 0.43 naar 0.66 bij 4× compressie en vermindert het aantal "no-answer" generaties aanzienlijk.

Serving Capaciteit

Geheugenreductie: De reductie van de cache is exact. Bij 21.3× compressie daalt de KV-cache voor Ouro-1.4B bij 1M context van 768 GiB naar 36 GiB.
Doorvoersnelheid (Throughput): Het "latent-store capacity path" (alleen latents opslaan) is compute-bound maar maakt batches en contexten mogelijk die voorheen onmogelijk waren door OOM (Out of Memory). Het "fast path" (reconstrueren en volledige K/V cachen) behoudt de doorvoersnelheid maar realiseert de geheugenbesparing niet.

Betekenis en Claims

Het artikel claimt dat recurrentie de meest comprimeerbare cache-as is in weight-tied looped Transformers. De betekenis ligt in:

Mechanistisch Inzicht: Het identificeert dat de loop-as spectraal onderscheidend is (lagere rang) van de kop- en laag-assen, waardoor efficiënte compressie mogelijk is zonder het dynamische gedrag van het model te laten instorten tot een statisch vast punt.
Praktische Implementatie: Het biedt een post-training oplossing die de batch size en de mogelijkheden voor contextlengte op bestaande hardware direct vergroot zonder de basisarchitectuur van het model opnieuw te trainen.
Toekomstige Richting: De resultaten suggereren dat native recurrente architecturen kunnen profiteren van ontkoppelde latente caches, waarbij de recurrentie gedeelde inhoud plus goedkope per-loop delta's berekent, in plaats van volledige K/V-states opnieuw te materialiseren.

De auteurs benadrukken dat de methode niet leunt op absorbed attention implementaties voor hun claims over nauwkeurigheid; de kern van het resultaat is de geldigheid van de lage-rang traject aanname en de effectiviteit van de codec in het behouden van het gedrag van de teacher onder strikte geheugenbeperkingen.