Maglev: Sliding Recurrent Memory

Technische Samenvatting: Maglev – Sliding Recurrent Memory

Probleemstelling
Huidige Transformer-architecturen worden geconfronteerd met een fundamentele afruil tussen geheugencapaciteit en computationele efficiëntie. Volledige causale zelf-attentie maakt krachtige nietlineaire retrieval en in-context computatie mogelijk, maar brengt kosten met zich mee die lineair groeien met de contextlengte. Daarentegen beperkt sliding-window aandacht de aandacht en KV-cache kosten, maar verwerpt het verre informatie, waardoor het faalt om persistent langetermijngeheugen te bieden. Klassieke recurrente modellen (bijv. LSTM's) comprimeren de historie, maar propageren kleine toestanden sequentieel, wat de capaciteit en sequentie-parallellisme beperkt. Hoewel Memory Transformers en lineaire aandachtmodellen proberen deze kloof te overbruggen, leggen zij vaak kunstmatige segmentgrenzen op, vertrouwen zij op gestructureerde lineaire updates, of vereisen zij sequentiële unrolling tijdens de training. De kernuitdaging is het ontwikkelen van een praktische, token-gewijze nietlineaire recurrente geheugenmechanisme voor Transformers dat parallelle pretraining mogelijk maakt terwijl de inferentiekosten begrensd blijven.

Methodologie: Maglev Architectuur
Maglev adresseert dit door de constructie van het geheugen tijdens de training te ontkoppelen van de propagatie van het geheugen tijdens de inferentie. De architectuur bestaat uit twee gekoppelde modellen:

1. Prefiller ($Q$): Een causale Transformer die de geobserveerde sequentie parallel verwerkt met behulp van een groter receptief veld (bijv. geïnterleavde volledige en sliding-window aandacht). Het genereert een traject van doel-geheugenvectoren, $m'_t$, voor elke tokenpositie.
2. Decoder ($P$): Een sliding-window Transformer die het volgende token voorspelt en zijn eigen geheugenvectoren genereert, $m_t$. Cruciaal is dat $P$ geen toegang heeft tot de volledige historie; het ziet alleen een begrensde window van recente tokens en een verschoven window van geheugenvectoren.

Trainingsmechanisme: Lifted Parallel Training
In plaats van de recurrente toestand sequentieel te unrollen (wat parallellisme verhindert), gebruikt Maglev een twee-pass trainingstrategie:

• Pass 1: De prefiller $Q$ genereert het volledige doel-geheugentraject $m'_{1:T}$ in parallel.
• Pass 2: De decoder $P$ consumeert het verschoven doel-traject $m'_{0:T-1}$ samen met de input-tokens om het volgende token te voorspellen en zijn eigen geheugen $m_{1:T}$ te produceren.
• Doelstelling: Het model wordt geoptimaliseerd met een samengestelde verliesfunctie:
  • Next-Token Prediction Loss ($L_{CE}$): Standaard cross-entropy verlies voor taalmodellering.
  • Memory Consistency Loss ($L_{cons}$): Een term die de output-geheugen van de decoder $m_t$ afstemt op de doelstelling van de prefiller $m'_t$. Dit leert de decoder om de noodzakelijke geheugentoestand voor de volgende stap te reproduceren zonder sequentiële unrolling.

Inferentiemechanisme
Tijdens de inferentie wordt de prefiller $Q$ weggegooid. De decoder $P$ draait recurrent en voert zijn eigen eerder gegenereerde geheugens $m_{t-1}$ terug in het systeem.

• Recurrent K/V Injectie: In plaats van specifieke geheugentokens te gebruiken, injecteert Maglev de verschoven geheugentoestand in de Key/Value (K/V) kenmerken van het aandachtmechanisme via geleerde projecties en gating-mechanismen.
• Begrensde Complexiteit: De decoder behoudt een vaste grootte van een sliding window van tokens en geheugens. De recurrente toestand wordt voortgevoerd door de K/V-cache, waardoor de inferentiekosten (compute en cache-grootte) begrensd blijven en onafhankelijk zijn van de totale sequentielengte, passend bij het profiel van standaard sliding-window aandacht.

Belangrijkste Bijdragen

1. Token-Gewijze Nietlineaire Recurrentie: Maglev introduceert een mechanisme waarbij de decoder-toestand van elk token dient als geheugen, waarbij het volledige nietlineaire diepte van de Transformer doorloopt zonder extra geheugentokens of kunstmatige segmentgrenzen te vereisen.
2. Parallel Training voor Recurrente Modellen: Door een prefiller te gebruiken om doelstellingen in parallel te generen, maakt Maglev parallelle pretraining mogelijk voor nietlineaire recurrente architecturen, waarmee de sequentiële bottlenecks van traditionele RNN-training worden vermeden.
3. Parameter Sharing: De auteurs demonstreren dat de prefiller $Q$ en de decoder $P$ het merendeel van hun parameters (specifiek de Transformer-blokken) kunnen delen zonder de prestaties significant te verslechteren, wat de parametergeheugenkosten van de trainingsopstelling vermindert.
4. Begrensde Inferentie: Het ingezette model behoudt de efficiëntie van sliding-window aandacht, terwijl het persistente geheugencapaciteiten verkrijgt door de recurrente K/V-injectie.

Experimentele Resultaten
Geëvalueerd op de nanochat d20 architectuur (435M parameters, getraind op 43.52B tokens), werd Maglev vergeleken met geïnterleavde volledige/sliding-window Transformers, pure sliding-window modellen en Latent Recurrent Transformer (LRT) baselines.

• Prestaties: De beste Maglev-variant (gescheiden parameters, $\lambda=1$) verbeterde de FineWeb-Edu validatie bits-per-byte (BPB) van 0.7413 (sliding-window baseline) naar 0.7251.
• Downstream Taken: Het bereikte een gemiddelde downstream nauwkeurigheid van 56,4%, waarmee het zowel de gematchte sliding-window Transformer (54,1%) als de gematchte LRT-baseline versloeg.
• Consistency Weight: De studie merkte op dat de optimale consistentie gewicht ($\lambda$) afhangt van de parameter sharing strategie; een hogere $\lambda$ was gunstig voor het model met gescheiden parameters, maar kon het model met gedeelde parameters beperken.

Betekenis en Claims
Het artikel claimt dat Maglev een praktisch pad biedt om nietlineaire Transformers uit te rusten met persistent, token-gewijs geheugen dat na elke token herschreven kan worden. Het slaagt erin de kloof te overbruggen tussen de expressieve kracht van volledige aandacht en de efficiëntie van modellen met een begrensde context. Door de generatie van het geheugen te scheiden van het gebruik van het geheugen tijdens de training, bereikt Maglev parallelle pretraining voor recurrente architecturen, een prestatie die voorheen moeilijk was voor nietlineaire recurrentie. De auteurs positioneren Maglev als een voorlopig onderzoek in deze ontwerpplek, waarbij zij opmerken dat hoewel het de sliding-window en latent recurrent baselines in de geteste setting verbetert, verder werk nodig is om de aanpak te schalen, de kernel-efficiëntie voor implementatie te optimaliseren en alternatieve geheugeninjectiemechanismen (bijv. residual-stream injectie of cross-attention) te verkennen.