Technische Samenvatting: Latent Recurrent Transformer (LRT)

Probleemstelling

Autoregressieve transformers zijn de standaard voor taalmodellering, maar ze vertrouwen op feedforward-berekeningen met een vaste diepte voor elke gegenereerde token. Hoewel recurrentie theoretisch de rekenkracht kan vergroten voor iteratieve verfijning, introduceren bestaande benaderingen aanzienlijke inferentiekosten. Methoden die recurrentie toevoegen in diepte (herhaaldelijk blokken toepassen op dezelfde token) of in tijd (pauze- of denk-tokens invoegen), verhogen de inferentielatentie door extra decodestappen of bloktoe passingen.

De auteurs stellen vast dat autoregressieve decoding al "gratis" hoogwaardige recurrente signalen berekent: de verborgen toestanden van de vorige token. Specifiek bevatten toestanden van hogere lagen, die zijn getraind voor voorspelling van de volgende token, bruikbare latente context. De kernvraag van het onderzoek is of deze al berekende hoogwaardige representaties van de vorige token kunnen worden hergebruikt als recurrente geheugen voor de volgende token zonder extra decodestappen toe te voegen.

Methodologie

Architectuur: Latent Recurrent Transformer (LRT)

LRT is een lichtgewicht uitbreiding van standaard autoregressieve transformers die een cross-layer, cross-token latente pad introduceert.

• Mechanisme: Op tokenpositie $t$ hergebruikt LRT een verborgen toestand van een bronlaag van de vorige positie, $m_{t-1} = h^{\ell_{src}}_{t-1}$, als recurrente geheugen. Hierbij is $\ell_{src}$ een geselecteerde bronlaag (doorgaans een boven-middenlaag).
• Injectie: Het geheugen $m_{t-1}$ wordt via twee complementaire mechanismen geïnjecteerd in de transformerlagen van de huidige token $t$:
  1. KV-projectie: Projecteert het geheugen naar de attention key-value ruimtes ($\tilde{k}_{rec}, \tilde{v}_{rec}$) en combineert deze met lokale keys/values met behulp van input-afhankelijke poorten. Dit biedt een directe route voor de representatie van de vorige token naar het attention-mechanisme.
  2. Residuale injectie: Voegt het geheugen direct toe aan de residustroom ($\bar{x}^{(\ell)}_t = \alpha_\ell x^{(\ell)}_t + \gamma_\ell m_{t-1}$), waardoor het gehele blok blootgesteld wordt aan het recurrente signaal.
• Inferentie-efficiëntie: Cruciaal behoudt LRT de standaard decoder-only backbone, het attention-mechanisme en de KV-cache-interface. Het voegt geen geheugentokens, pauzetokens of dieptelussen in. Elke gegenereerde token gebruikt nog steeds precies één normale transformer forward pass, waarbij de recurrente toestand eenvoudig wordt bijgewerkt en doorgegeven aan de volgende stap. De standaardvariant voegt slechts 0,3% parameters toe.

Trainingsstrategie: Verweven Parallelle Training

Een grote uitdaging is het pretrainen van deze recurrentie op schaal. Exacte training vereist het sequentieel ontrollen van de transformer over de volledige sequentie (Backpropagation Through Time), wat de parallelisme vernietigt die essentieel is voor efficiënt transformer pretraining. Geblokkeerde training (het sequentieel verwerken van aaneengesloten blokken) behoudt enig parallelisme, maar propageert geheugen alleen over blokgrenzen heen, wat een grove benadering biedt.

De auteurs introduceren Verweven Parallelle Training om token-level recurrentie te benaderen terwijl parallelle berekening behouden blijft:

1. Initialisatiepass: Een enkele forward pass over de volledige sequentie bouwt een gedeelde buffer op die keys, values en recurrente toestanden van de bronlaag bevat voor alle posities.
2. Verfijningpasses: De tokenposities worden gepartitioneerd in $S$ disjuncte subsets (bijvoorbeeld gestructureerde partities). Deze subsets worden parallel verfijnd.
3. Schrijven terug: Bijgewerkte toestanden van een subset worden teruggeschreven naar de gedeelde buffer. Aansluitende subsets kunnen het verfijnde geheugen van eerdere subsets consumeren.
4. Doel: De loss is het gemiddelde van de initialisatie-loss en de verfijning-losses voor elke subset.

Deze strategie zorgt ervoor dat elke token een supervisiestap ontvangt die rekening houdt met recurrente geheugen. In termen van rekenkracht kost het ongeveer 2× een standaard transformer-update (één initialisatie + één effectieve verfijning), maar het bootst token-level recurrentie nauwkeuriger na dan geblokkeerde training.

Belangrijkste Bijdragen

1. LRT-architectuur: Een lichtgewicht recurrente uitbreiding die hoogwaardige verborgen toestanden van een bronlaag van de vorige token hergebruikt als geheugen, waardoor cross-layer recurrentie mogelijk wordt zonder extra decodestappen of architecturale wijzigingen aan de KV-cache-interface.
2. Verweven Parallelle Training: Een parallelle trainingsbenadering die disjuncte token-subsets verfijnt om token-level recurrentie na te bootsen binnen een vast trainingsbudget, en zo de sequentiële ontrolling vermijdt die vereist is voor standaard RNN-achtige training.
3. Empirische Validatie: Aantonen dat LRT de afweging tussen rekenkracht en kwaliteit verbetert over verschillende modelgroottes en tokens-per-parameter-budgetten, waarbij de standaard gedeelde-projectievariant minimale parameter-overhead toevoegt (0,3%).

Resultaten

Experimenten werden uitgevoerd op nanochat-stijl GPT-backbones (20L en 24L lagen) die vooraf waren getraind op FineWeb-Edu.

• Schalingsgedrag: Onder trainingsrekenkracht die gelijkwaardig is aan de baseline (rekening houdend met de 2× kosten van verweven training), verschuift LRT de schalingscurves naar lagere Bits Per Byte (BPB) en hogere CORE (een metriek voor in-context leren over 22 taken).
• Prestatiewinst:
  • Op de 24L-backbone bij rekenkracht 80 gelijkwaardig aan de baseline, bereikte LRT-shared 0,695 BPB tegenover 0,699 voor de baseline.
  • In-context leren (CORE) verbeterde eveneens; op de 20L-backbone bij rekenkracht 80 bereikte LRT-shared 0,274 versus 0,271 voor de baseline.
• Ablatiebevindingen:
  • Bronlaag: Een boven-middenlaag (bijvoorbeeld laag 12 in een 20L-model) presteert beter dan de laatste laag, wat suggereert dat bruikbaar geheugen hoogwaardig moet zijn maar niet te gespecialiseerd voor logits.
  • Temporele bron: De onmiddellijk voorafgaande toestand ($m_{t-1}$) is de sterkste enkele geheugenbron die de beperking van één forward pass behoudt.
  • Injectie: KV-projectie en residuale injectie zijn complementair; het gebruik van beide levert de beste resultaten op.
  • Training: Verweven parallelle training presteert aanzienlijk beter dan geblokkeerde training, die een grovere benadering van recurrentie biedt.

Betekenis en Aanspraken

Het artikel beweert dat LRT succesvol de afweging tussen rekenkracht en kwaliteit voor autoregressieve transformers verschuift. Door "gratis" hoogwaardige signalen van de vorige token te hergebruiken en deze te trainen via een parallelle benadering, maakt LRT iteratieve verfijning en beter contextgebruik mogelijk zonder de inferentielatentieboetes van diepterecurrente methoden of methoden met pauzetokens.

De auteurs benadrukken dat LRT een lichtgewicht uitbreiding is die de standaard transformer-interface behoudt. De winst wordt toegeschreven aan het recurrente geheugenmechanisme zelf en niet simpelweg aan het toevoegen van extra parameters of poorten, aangezien generieke architecturale toevoegingen (zoals gated attention of laagschaalvergroting) op dezelfde sterke baseline verwaarloosbare verbeteringen opleverden. Het werk suggereert dat bruikbaar recurrente geheugen hoogwaardig moet zijn, maar niet de laatste logit-gespecialiseerde laag, en dat token-level recurrentie tijdens pretraining effectief parallel kan worden benaderd.