Technische Samenvatting: T2MLR: Transformer met Temporale Middelste-Laag Recurrente Structuren

Probleemstelling

De auteurs identificeren een fundamentele beperking in standaard Transformer-architecturen met betrekking tot redeneervermogen. Hoewel Transformers uitblinken in in-context retrieval, creëert hun autoregressieve decoderingmechanisme een informatie-bottleneck. Bij elke decoderingsstap worden rijke, hoogdimensionale latente representaties herhaaldelijk gecomprimeerd naar een ijle tokenruimte (one-hot vectoren) voordat ze als input voor de volgende stap dienen. Dit proces voorkomt dat abstracte intermediaire redeneertoestanden over de tijd heen voortbestaan, waardoor het model gedwongen wordt om noodzakelijke informatie indirect via aandacht of via sterk gecomprimeerde input-embeddings te reconstrueren.

Bestaande benaderingen voor latent redeneren proberen dit te verzachten door ofwel volledig in de continue ruimte te redeneren (vaak in de laatste laag), ofwel door onzekerheid te propageren via lineaire combinaties van token-embeddings. Echter, deze methoden voeden recurrente signalen doorgaans terug via de input-embedding, waardoor de middelste lagen – waar abstract redeneren bekend is te plaatsvinden – worden omzeild. Alternatieve "geloopte" (looped) Transformer-architecturen vergroten de effectieve diepte door over alle lagen te itereren, maar dit brengt inferentiekosten met zich mee die lineair schalen met het aantal loops, wat onpraktisch is voor implementatie.

Methodologie

Het artikel introduceert T2MLR (Transformers met Temporale Middelste-Laag Recurrente Structuren), een nieuwe architectuur die is ontworpt om abstracte intermediaire representaties te laten voortbestaan en evolueren over decoderingsstappen heen, zonder de inferentie-tijd diepte te verhogen.

Architecturaal Ontwerp

T2MLR wijzigt een standaard $L$-laagse decoder-only Transformer door een gegateerde recurrente pathway te introduceren die een diepere laag van het vorige token verbindt met een eerdere laag van het huidige token.

1. Recurrente Grenzen: De architectuur wordt geparametriseerd door twee laag-indices, $\ell_{start}$ en $\ell_{end}$, die een "middelste" blok van lagen definiëren.
2. Recurrente Cache ($R_t$): Voor elke decoderingsstap $t$ wordt een cache-vector van constante grootte bijgehouden. Deze slaat de representatie op na laag $\ell_{end}$ van de vorige stap.
3. Fusie-module ($\Phi$): Voordat de huidige token wordt verwerkt bij laag $\ell_{start}$, fuseert het model de representatie van de huidige token ($h^{(\ell_{start}-1)}_t$) met de recurrente cache ($R_{t-1}$) van de vorige token. De fusie wordt beheerst door een leerbaar gegateerd mechanisme:
   $$\Phi(h, R) = h + \tanh(\gamma_{cur}) \sigma(f_{cur}([h, R])) \odot h + \tanh(\gamma_{rec}) \sigma(f_{rec}([h, R])) \odot W_{rec}R$$
   Dit ontwerp stelt het model in staat om te leren wanneer het moet vertrouwen op de stroom van de huidige token versus de temporeel geaccumuleerde latente toestand.
4. Cache Update: Na passage door laag $\ell_{end}$ wordt de recurrente cache bijgewerkt met behulp van de nieuwe representatie en de vorige cache via een residuele verbinding en RMSNorm.

Trainingsstrategie: Temporale Parallelisme

Een belangrijke uitdaging voor latente redeneerarchitecturen is dat standaard sequentie-parallel training (token-parallel) wordt doorbroken wanneer er afhankelijkheden bestaan tussen tijdstappen. Om dit aan te pakken, stellen de auteurs een benaderde temporale-parallelle trainingsmethode voor met behulp van Jacobi fixed-point iteraties:

• In plaats van sequentiële tokenverwerking, voert het model $d_{forward}$ iteraties van verfijning parallel uit.
• Het begint met een initiële cache afgeleid van een standaard forward pass, iteratieve fusie van deze cache met de huidige token-representatie, verfijnt het middelste blok, en werkt de cache bij.
• Gradiënten worden door deze iteraties teruggepropageerd met een afgekapte diepte ($d_{backward}$), analoog aan Truncated Backpropagation Through Time (TBPTT).
• Deze aanpak maakt het mogelijk om de modellen te trainen met dichte teacher forcing en sequentie-parallelisme, een capaciteit die de auteurs uniek achten onder de varianten van continue chain-of-thought.

Belangrijkste Bijdragen

1. Architecturale Innovatie: T2MLR introduceert een gerichte recurrente pathway specifiek in de middelste lagen van de Transformer, in plaats van bij de input-embedding of de finale output. Dit behoudt de standaard autoregressieve interface terwijl het latente toestand-persistentie mogelijk maakt.
2. Schaalbare Training: De auteurs demonstreren een methode om deze recurrente afhankelijkheden te trainen met benaderde temporale parallelisme, wat pretraining en fine-tuning op grote datasets mogelijk maakt zonder de schaalbaarheid op te offeren.
3. Retrofitting Capaciteit: Het artikel toont aan dat T2MLR geen training vanaf nul vereist. De recurrente pathway kan in een bestaand pretrained model (bijv. SmolLM2-1.7B) worden geïnjecteerd en wordt gefinetuned om significante verbeteringen in redeneren te leveren.
4. Efficiëntie: In tegen tegenstelling tot geloopte Transformers voegt T2MLR verwaarloosbare inferentie-overhead toe (maximaal ~8% per token) omdat de recurrentie wordt afgehandeld via een cache van constante grootte en een enkele forward pass tijdens inferentie.

Experimentele Resultaten

De auteurs evalueren T2MLR over synthetische benchmarks, natuurlijke taal pretraining en downstream redeneertaken.

• Synthetisch Redeneren (S5-Retrieval): T2MLR presteert aanzienlijk beter dan zowel standaard Transformers als LSTM's op een taak die gelijktijdige toestand-tracking en in-context retrieval vereist. Terwijl Transformers falen bij lange sequentielengtes door dieptebeperkingen en LSTM's falen bij retrieval door vaste toestand-compressie, behoudt T2MLR een hoge nauwkeurigheid.
• Pretraining (FineWeb-Edu): Op 135M parameter modellen presteren T2MLR-varianten consequent gelijk aan of beter dan parameter-gematchte Transformer-baselines op zero-shot NLP-benchmarks (ARC, HellaSwag, etc.). Opvallend genoeg presteert middelste-laag recurrentie (het loopen over ~20% van de lagen, bijv. lagen 13–18) vaak beter dan volledige-laag recurrentie.
• Downstream Redeneren: Op taken zoals GSM-Aug (wiskunde), HotPotQA (multi-hop QA) en Variable Assignment laat T2MLR consistente winst zien ten opzichte van baselines. De varianten met middelste-laag recurrentie presteren opnieuw beter dan volledige-laag recurrentie, wat suggereert dat gerichte recurrentie effectiever is dan globale looping.
• Schaling: De prestatiewinsten blijven aanwezig en groeien bij opschaling naar 361M en 1B parameters en het uitbreiden van pretraining naar 50B tokens.
• Retrofitting: Het toepassen van T2MLR op een pretrained SmolLM2-1.7B-Instruct model en het finetunen op wiskundige data verbeterde de GSM8K nauwkeurigheid van 35.8% naar 39.9% en MATH500 van 12.8% naar 18.0%, wat de praktische toepasbaarheid zonder volledige pretraining aantoont.

Betekenis en Claims

Het artikel stelt dat effectief latent redeneren in Transformers niet vereist om over alle lagen te loopen (wat de inferentiekosten verhoogt) noch om de recurrentie te beperken tot de input-embedding (wat de abstracte computatie omzeilt). In plaats daarvan claimen de auteurs dat gerichte middelste-laag recurrentie voldoende is om significante redeneercapaciteiten te ontsluiten.

De betekenis van T2MLR ligt in het vermogen om:

1. Redeneerdiepte te ontkoppelen van Inferentiekosten: Het introduceert een temporale latente pathway die over de tijd heen persisteert zonder de lineaire inferentie-straf die gepaard gaat met geloopte architecturen.
2. De Drempel voor Adoptie te Verlagen: Door aan te tonen dat de architectuur kan worden geretrofit op bestaande pretrained modellen, beargumenteren de auteurs dat geavanceerde redeneercapaciteiten kunnen worden toegevoegd aan huidige modellen tegen een fractie van de kosten van pretraining vanaf nul.
3. Inductieve Biases te Valideren: De resultaten suggeren dat het "sweet spot" voor recurrentie in Transformers de middelste lagen zijn, waar abstracte computatie plaatsvindt, in plaats van de input- of output-interfaces.

De auteurs erkennen dat T2MLR een training-time overhead met zich meebrengt (2–4x wandkloktijd afhankelijk van de grootte van het recurrente blok) vanwege de Jacobi-benaderingsstappen. Zij kaderen dit als een trade-off: T2MLR offert training-compute op om een laag-overhead, hoog-presterende inferentie-architectuur te verkrijgen voor redeneertaken.