Key-Value Means

Het artikel introduceert Key-Value Means (KVM), een nieuw blokgewijs recurrentie-mechanisme voor attention dat de voordelen van transformers en lineaire RNN's verenigt door efficiënt, chunk-paralleliseerbaar trainen met flexibele statiegroei en subkwadratische prefill-tijd mogelijk te maken, alles met standaardoperaties en minimale extra parameters.

De kern

Stel je voor dat je probeert een heel lang boek te lezen, maar je hersenen hebben een beperkte hoeveelheid "werkgeheugen" om het verhaal in je hoofd vast te houden terwijl je leest.

Het probleem met huidige AI
Huidige AI-modellen (Transformers) gedragen zich als een student die probeert om elk enkel woord dat ze ooit in het boek hebben gelezen, te onthouden.

• Het goede: Ze zijn ongelooflijk nauwkeurig omdat ze het hele verhaal voor zich hebben.
• Het slechte: Naarmate het boek langer wordt, groeit hun "werkgeheugen" enorm. Het lezen van een boek van 100 pagina's kost een klein beetje moeite, maar het lezen van een boek van 1.000 pagina's kost een enorme hoeveelheid tijd en energie. Het is alsof je probeert een rugzak te dragen die met elke stap die je zet zwaarder wordt.

Het probleem met Recurrente (RNN-stijl) Modellen
RNN-stijl modellen nemen een andere aanpak: ze houden een kleine, vaste samenvatting van wat ze tot nu toe hebben gelezen bij en werken deze bij naarmate ze verder lezen.

• Het goede: Ze zijn supersnel en licht. Hun "rugzak" wordt nooit zwaarder, ongeacht hoe lang het boek is.
• Het slechte: Ze vergeten het begin van het verhaal. Als je ze vraagt naar een plotpunt van pagina 10, herinneren ze het zich misschien niet omdat ze alleen de laatste paar pagina's vasthouden.

De nieuwe oplossing: Key-Value Means (KVM)
De auteurs van dit paper introduceren een nieuwe methode genaamd Key-Value Means (KVM). Denk aan KVM als een slim, magisch notitieboek dat het beste van twee werelden combineert.

Zo werkt het met een eenvoudige analogie:

1. Het "Schuifvenster" (De directe context)

Stel je voor dat je een boek leest en je hebt een vergrotingsglas dat alleen de laatste paar pagina's duidelijk laat zien. Dit is het "Schuifvenster". KVM besteedt perfecte aandacht aan de meest recente woorden, net als een standaard-AI. Dit zorgt ervoor dat het de directe context niet mist.

2. De "Gecomprimeerde Samenvatting" (Het langetermijngeheugen)

Terwijl je voorbij die paar pagina's leest, glijden de oude pagina's uit je vergrotingsglas. In plaats van ze weg te gooien (zoals RNN-stijl modellen doen) of te proberen het hele boek te dragen (zoals de huidige AI), doet KVM iets slims:

• Het kijkt naar de pagina's die net zijn weggelopen.
• Het vraagt zich af: "Welke van deze pagina's zijn het belangrijkst of uniekst?"
• Het schrijft een korte, gecomprimeerde samenvatting van die belangrijke pagina's in een speciaal notitieboek.
• Als er een nieuwe pagina langskomt die erg lijkt op wat er al in het notitieboek staat, werkt het alleen de bestaande notitie bij. Als het iets heel nieuws en verrassends is, voegt het een nieuwe regel toe aan het notitieboek.

3. De "Slimme Samenvoeging" (De magische truc)

Het paper beschrijft een specifieke manier van informatie samenvoegen die een "Winner-Take-All" regel wordt genoemd.

• Stel je voor dat je een emmer water hebt (de nieuwe informatie) en een spons (het notitieboek).
• In plaats van het water er gewoon in te gieten, vindt KVM de exacte plek in de spons die het beste bij het water past en absorbeert het daar.
• Het gebruikt ook een "Just-in-Time" normalisatie. Stel je voor dat je de notities in het boek eerst schrijft als ruwe sommen en aantallen, zonder ze direct te delen of te middelen. KVM slaat deze "onbewerkte" gegevens op terwijl je schrijft. Pas op het moment dat je echt naar het notitieboek kijkt (bij het lezen van een vraag), doet het de deling om het juiste gemiddelde te berekenen. Door deze berekening uit te stellen tot het exacte moment van gebruik ("just-in-time"), hoeft het niet elke keer opnieuw te rekenen als er nieuwe informatie wordt toegevoegd.

Waarom dit belangrijk is

• Flexibele grootte: Je kunt KVM vertellen om een klein notitieboek (vaste grootte) te houden voor snelheid, of het notitieboek te laten groeien naarmate het boek langer wordt (uitbreidbare grootte).
• Snelheid versus geheugen: Het stelt je in staat om een middenweg te kiezen. Je hoeft niet te kiezen tussen "supersnel maar vergeetachtig" of "superslim maar traag". Je kunt het zo afstellen dat het snel genoeg is voor realtime gebruik, maar slim genoeg om het hele verhaal te onthouden.
• Geen aangepaste hardware: In tegenstelling tot sommige andere nieuwe methoden die speciale, dure computerchips nodig hebben om te draaien, kan KVM draaien op standaardcomputers met normale softwarebewerkingen.

De resultaten

De auteurs hebben dit getest op taalmodellen (AI die tekst leest en schrijft).

• Korte contexten: Het presteerde net zo goed als de beste standaard-AI-modellen.
• Lange contexten: Toen de invoer erg lang werd (duizenden tokens), onthield KVM met een groeiend notitieboek details veel beter dan de "vergeetachtige" RNN-stijl modellen en was het veel sneller dan de "het-hele-boek-dragende" modellen.
• De "Naald in een hooiberg": Bij tests waarbij de AI een specifieke zin moest vinden die verborgen zat in een enorme tekst, presteerde de versie van KVM die zijn notitieboek liet groeien zeer goed, wat bewees dat het daadwerkelijk informatie uit het verre verleden kon terugroepen.

Kortom, KVM is een nieuwe manier voor AI om lange boeken te lezen zonder moe te worden, zonder het begin te vergeten en zonder een rugzak te hoeven dragen die oneindig zwaar wordt. Dit doet het door een helder zicht op het heden te houden terwijl het tegelijkertijd een slimme, gecomprimeerde samenvatting van het verleden behoudt.

Technische samenvatting

Technische Samenvatting: Key-Value Means (KVM)

Probleemstelling

Transformers bieden efficiënt trainen op moderne hardware, maar lijden onder lineaire schaling in geheugen en tijd per output-token ten opzichte van de contextlengte ($O(N)$ geheugen, $O(N)$ decodeertijd). Omgekeerd bereiken moderne Lineaire RNN's (LRNN's) constant geheugen en tijd per token ($O(1)$), maar hebben ze doorgaans moeite met beperkte recall op lange contexten. Bestaande architecturen die proberen deze kloof te overbruggen, vertrouwen vaak op staten met een vaste grootte (wat recall beperkt) of complexe training tijdens het testen met runtime-optimizers (wat de snelheid beïnvloedt). Er is behoefte aan een architectuur die geheugenefficiëntie, snelheid en recall op lange contexten in evenwicht brengt, zonder dat er aangepaste kernels of complexe hyperparameter-tuning voor training tijdens het testen vereist zijn.

Methodologie: Key-Value Means (KVM)

KVM is een nieuw blokrecurrerend attentiemechanisme dat een blokglijdend raamattentie (BSWA) integreert met een dynamisch uitbreidbare, gecomprimeerde staat. Het werkt binnen een enkele softmax-attentielayer en verenigt de voordelen van traditionele transformers (uitbreidbare context, parallelisme per chunk) en lineaire RNN's.

Kernmechanismen

1. Blokglijdend Raam met Gecomprimeerde Staat:
   KVM verwerkt invoer in chunks. Het handhaaft een raam met vaste grootte (BSWA) voor recente tokens en een aparte, periodiek bijgewerkte staat voor oudere tokens. Wanneer een blok tokens het BSWA-raam overstroomt, wordt het verwerkt om de staat bij te werken in plaats van te worden weggegooid.

2. Staatcompressie en Samenvoeging:
   Overstroomde tokens worden gecomprimeerd in de staat met behulp van een "winner-take-all" samenvoegregel die lijkt op cosinus-ähnelijkheid.

   • Similariteitsmetriek: In plaats van standaard softmax gebruikt KVM een maximaal schaars update-matrix (geïnspireerd door Online Vector Quantization), waarbij elke overstroomde key wordt toegewezen aan de enige meest gecorreleerde state-key.
   • Just-In-Time (JIT) Hernormalisatie: Om te voorkomen dat de norm van state-vectoren in de loop van de tijd krimpt door het middelen van orthogonale of tegenstrijdige vectoren, past KVM JIT-normalisatie toe. State-keys worden genormaliseerd met LayerNorm voordat de attentie wordt toegepast. State-waarden worden genormaliseerd naar een vaste "leesstraal" ($\rho_i$) die wordt bepaald bij het aanmaken van de sleuf, waardoor waardenmagnitudes behouden blijven terwijl richtingsveranderingen mogelijk zijn.
   • Merge-gate: Een data-afhankelijke scalair gate moduleert de hoeveelheid inkomende overstroomde key/waarde die door de staat wordt geabsorbeerd.

3. Staatuitbreidingsstrategie:
   In tegenstelling tot RNN's met een vaste grootte, ondersteunt KVM een uitbreidbare staat. De meest "verrassende" (minst redundante) overstroomde tokens worden direct aan de staat toegevoegd, terwijl de rest wordt samengevoegd. Dit zorgt voor sublineaire geheugengroei (bijv. $O(\sqrt{N})$) terwijl recall op vroege contexten behouden blijft.

4. Verwerking van Positie-encoding:
   Om compatibiliteit te behouden met Rotary Positional Embeddings (RoPE) in het BSWA-raam, terwijl RoPE in de gecomprimeerde staat wordt vermeden (die tokens van sterk variërende posities aggregeert), hanteert KVM gedeeltelijk RoPE-nulstellen. De rotatie-onderruimte van state-keys wordt op nul gezet, terwijl het BSWA-raam volledige RoPE behoudt. Dit stelt het model in staat om niet-gerotereerde queries voor de staat en gerotereerde queries voor het raam te gebruiken binnen dezelfde attentiepass.

5. Sink-tokens:
   Een beschermde set van initiële staatrijen (sinks) wordt behouden om de degradatie van kritieke informatie uit de vroege context te voorkomen, waarmee het probleem wordt aangepakt dat sink-tokens verschillende waardenmagnitudes hebben.

Belangrijkste Bijdragen

Het artikel presenteert de volgende specifieke bijdragen:

• Nieuwe Blokrecurrerende Formulering: Een methode om overstroomde tokens te comprimeren in een dynamisch hernormaliseerde staat met behulp van een "winner-take-all" samenvoegregel, waardoor de noodzaak van aparte compressielagen wordt geëlimineerd.
• Staatuitbreidingsstrategie: Een mechanisme om nieuwe overstroomde tokens aan de staat toe te voegen, waardoor sublineaire geheugengroei mogelijk is zonder in te leveren op recall.
• JIT Hernormalisatie: Een schema om state-keys en -waarden just-in-time te normaliseren om vectornormen te behouden en destructieve interferentie tijdens het middelen te voorkomen.
• Gedeelde RoPE: Een methode om positie-encoding te delen over gecomprimeerde en niet-gecomprimeerde gebieden door de RoPE-dimensie in state-keys op nul te zetten, waardoor de noodzaak van extra parameters of complexe her-samenvoegmechanismen wordt vermeden.
• Gefuseerde Architectuur: Een enkele attentielayer die interpoleert tussen RNN's met een vaste staat en full-attention Transformers, en een continue afweging biedt tussen geheugenefficiëntie en recall.

Experimentele Resultaten

De auteurs trainden modellen (120M en 350M parameters) op de Prolong-dataset met contextlengtes van 8k.

• Prestaties op Lange Contexten:
  • Fixed-State KVM (256 tokens): Presteerde beter dan grotere OVQ/SWA-modellen op sequentiepositieverlies en benchmarks voor korte contexten. Het had echter moeite in "Needle In A Haystack" (NIAH)-tests met nieuwe afleidingsfactoren (NIAH-S2/S3) bij extreme lengtes, waarbij de statencapaciteit een bottleneck werd.
  • Growable KVM (Power-law/Saturating schema's): De "KVM sqrt"-variant (staatgrootte $\propto \sqrt{N}$) behaalde concurrerende resultaten op benchmarks voor lange contexten (RULER, LongBench, NIAH), waarbij het niet-hybride GPTAlpha-modellen in extrapolatiezones buiten de 8k trainingscontext evenaarde of versloeg. Het presteerde aanzienlijk beter dan fixed-state KVM en pure LRNN's (RWKV-7) op taken die het ophalen van nieuwe informatie over lange afstanden vereisten.
• Prestaties op Korte Contexten: KVM-varianten presteerden consistent met standaard Transformers op benchmarks voor korte contexten (LAMBADA, ARC, HellaSwag, etc.), wat bevestigt dat het BSWA-raam standaardattentiecapaciteiten behoudt.
• Ablatiestudies: Het verwijderen van normalisatie van waardenlengte veroorzaakte de meest significante prestatiedegradatie. Het verwijderen van sink-bescherming en de merge-gate verzwakte ook de recall op lange contexten aanzienlijk.

Betekenis en Claims

Het artikel beweert dat KVM de kloof tussen RNN's met een vaste staat en full-attention Transformers succesvol overbrugt.

• Efficiëntie versus Recall: Het biedt een flexibele keuze voor staatgrootte, waardoor gebruikers de afweging tussen geheugenefficiëntie en recall kunnen afstemmen. Met een vaste staat biedt het $O(N)$ chunked recurrent gedrag; met een uitbreidbare staat bereikt het sublineaire geheugengroei met sterke recall op lange contexten.
• Implementatie-eenvoud: KVM is implementeerbaar met standaardoperaties zonder aangepaste kernels en ondersteunt chunk-wise paralleliseerbaar trainen en prefill.
• Hybride Potentieel: De architectuur kan worden gebruikt in hybride oplossingen naast LRNN-lagen om deze aan te vullen met verbeterde sublineaire geheugengroei en mogelijkheden voor decoding op lange contexten.
• Geen Runtime-optimizers: In tegenstelling tot Test-Time Training (TTT)-benaderingen, vertrouwt KVM op een eenvoudige staatupdate-regel in plaats van runtime-optimizers zoals SGD of Adam, waardoor de bijbehorende hyperparameter-uitdagingen worden vermeden.

De auteurs concluderen dat KVM aantoont dat het mogelijk is om op een eenvoudige en effectieve manier vloeiend te interpoleer tussen RNN's met een vaste staat en full attention, en zo een gefuseerde oplossing biedt voor modellering op lange contexten.