Technische Samenvatting: RecurrentGPT

Probleemstelling

Het schalen van transformer-taalmodellen creëert een fundamentele spanning tussen expressiviteit en geheugenefficiëntie. Standaardarchitecturen koppelen diepte en parameters rigide aan elkaar; elke extra laag introduceert een unieke set gewichten, wat een geheugenbottleneck creëert die effectieve diepte onder een vaste hardwarebudget beperkt. Omgekeerd dwingen standaard strategieën voor gewichtsdeling (bijv. het delen van alle lagen over de diepte) uniforme transformaties af op evoluerende representaties. Deze statische aanpak reduceert de functionele diversiteit die diepte waardevol maakt, aangezien de verborgen toestand na de eerste recursie fundamenteel verschilt van de toestand na de achtste, terwijl beide worden gedwongen door dezelfde statische transformatie.

Het artikel stelt dat de kracht van een redeneersysteem niet ligt in de breedte van de parameters, maar in de diepte van de iteratie (verwijzend naar Turing [2] en moderne test-time compute paradigma's [3–5]). De uitdaging is om diepe, iteratieve verfijning van inputrepresentaties te bereiken zonder de geheugenkosten van unieke gewichten voor elke laag, terwijl de representatieve instorting wordt vermeden die wordt veroorzaakt door eenvoudige gewichtsdeling.

Methodologie: RecurrentGPT Architectuur

De auteurs stellen RecurrentGPT voor, een recurrente diepte-transformer die de modeldiepte ontkoppelt van het aantal parameters door een gedeelde kern te itereren terwijl deze wordt omlijst door vaste prelude- en coda-blokken.

1. Architectuur Lay-out

Het model volgt een $n_{pre} + n_{rec} \times R + n_{coda}$ configuratie:

• Prelude ($n_{pre}$ blokken): Wordt eenmalig toegepast om een stabiel, contextbewust conditioneringssignaal ($h^{(pre)}$) te genereren vanuit de input.
• Gedeelde Kern ($n_{rec}$ blokken): Een set transformer-blokken die $R$ keer recursief wordt toegepast. In tegenstelling tot standaard gewichtsdeling wordt deze kern niet uniform toegepast; deze wordt bij elke stap gemoduleerd.
• Coda ($n_{coda}$ blokken): Verwerkt de uiteindelijke verfijnde toestand om logits te produceren.

Het totale aantal unieke parameters schaalt als $\Theta((n_{pre} + n_{rec} + n_{coda}) \cdot d^2)$, onafhankelijk van het aantal recursies $R$.

2. Adaptieve Gating en Recurrente Projectie

Om representatieve instorting te voorkomen en een enkele gewichtstensor te laten fungeren als meerdere gespecialiseerde lagen, introduceert RecurrentGPT een per-element geleerde gate en een recurrente projectie.

Bij elke recursiestap $r$:

1. Input Voorbereiding: De vorige verborgen toestand $h^{(r-1)}$ wordt geprojecteerd samen met de vaste prelude-output $h^{(pre)}$ en geïnjecteerde stochastische ruis $\epsilon_x$.
   $$\tilde{h}^{(r)} = W_{proj} [h^{(r-1)} + \epsilon_x, h^{(pre)}]$$
2. Gedeelde Transformatie: De gedeelde blokken ($B_{shared}$) verwerken deze input om een voorstel te produceren $o^{(r)}$.
   $$o^{(r)} = B_{shared}(\tilde{h}^{(r)})$$
3. Elementwise Gating: Een geleerde gate $g^{(r)}$ bepaalt hoeveel van het voorstel de huidige toestand vervangt. De gate is geconditioneerd op de genormaliseerde huidige toestand, de genormaliseerde prelude en geïnjecteerde ruis $\epsilon_g$.
   $$g^{(r)} = \sigma(f_{gate}([LN(h^{(r-1)}), LN(h^{(pre)})]/\tau + \epsilon_g))$$
   $$h^{(r)} = g^{(r)} \odot h^{(r-1)} + (1 - g^{(r)}) \odot o^{(r)}$$

Belangrijke Ontwerpkeuzes:

• Initialisatie: De gate-bias is geïnitialiseerd op $+4$, wat resulteert in $g^{(r)} \approx 0.98$ aan het begin van de training. Dit zorgt ervoor dat de residuele stroom aanvankelijk bijna onveranderd blijft, zodat het model kan leren welke elementen geleidelijk overschreven moeten worden.
• Ruis: Additieve ruis ($\epsilon_x$) en gate-ruis ($\epsilon_g$) voorkomen breekbare exacte-match patronen en gate-instorting.
• Diepte Sampling: Tijdens de training wordt de recursiediepte $r$ uniform gesampled uit $\{1, \dots, R\}$. Dit traint impliciet elk exitpunt, wat emergent early exit mogelijk maakt tijdens inferentie zonder hulp-losses.

Belangrijkste Bijdragen

1. Architectuur: RecurrentGPT maakt gebruik van per-element gating en stochastische perturbaties om representatieve instorting te voorkomen, waardoor een gedeelde kern zich verschillend kan gedragen over de recursiestappen heen.
2. Uniforme Evaluatie: Het model wordt geëvalueerd onder zowel isoFLOPS (overeenkomstige compute, minder parameters) als isoPARAMS (overeenkomstige parameters, meer compute) regimes.
   • isoFLOPS: RecurrentGPT leidt Mixture-of-Recursions (MoR) en heavy-tail diepte-sampling in alle negen scale-by-budget cellen. Op kleine schaal evenaart het de dense GPT-2 Small baseline met 64% minder parameters. Op medium en grote schaal benadert het de dense kwaliteit naarmate het tokenbudget groeit.
   • isoPARAMS: Diepere recursie levert consistente winsten in validatieverlies (bijv. 2.76 vs 2.84 voor een niet-recurrente tegenhanger bij gelijk aantal parameters).
3. Emergent Early Exit: Zonder hulp-losses ondersteunt het model van nature vroege exit tijdens inferentie, waarbij het $\sim92\%$ van zijn nauwkeurigheid behoudt wanneer het slechts de helft van zijn recursiestappen uitvoert.

Experimentele Resultaten

• Prestaties: Onder isoFLOPS-beperkingen evenaart een 3-laagse RecurrentGPT de nauwkeurigheid van een 12-laagse GPT-2 Small baseline. Op medium schaal (equivalent aan 24 lagen) bereikt het een validatieverlies van 2.89 vergeleken met de dense baseline van 2.84, terwijl het slechts 124M parameters gebruikt versus 354M.
• Efficiëntie: Onder isoFLOPS slaat het model slechts 36–37% van de unieke parameters van de dense baseline op.
• Downstream Taken: Op grootschalige benchmarks presteert de isoPARAMS-variant beter dan de dense baseline op 8 van de 9 taken, met een gemiddelde verbetering van +2.10 punten. De isoFLOPS-variant evenaart de gemiddelde prestaties van de dense baseline (42.08 vs. 42.05) met slechts 37% van de parameters.
• Latentie en Geheugen: Onder torch.compile veroorzaakt RecurrentGPT-Large slechts een toename van 10% in generatielatentie terwijl het piek GPU-geheugen met 59% en het aantal parameters met 63% vermindert.
• Mechanistische Inzichten: Analyse onthult een twee-fasen computationele strategie:
  • Vroege Stappen (Schrijf-intensief): De gate staat meer open en het model voert snelle herstructurering uit, waarbij ~77% van de verliesreductie in de eerste twee stappen wordt gerealiseerd.
  • Late Stappen (Kopieer-intensief): De gate wordt nauwer en het model voert gerichte correcties uit.
  • Token Routing: Het model wijst impliciet meer berekening toe aan moeilijke tokens (die een hoge initiële loss hebben) zonder expliciete routing-signalen.

Betekenis en Claims

Het artikel claimt dat adaptieve diepte-hergebruik een fundamentele strategie is voor het ruilen van parameters voor kwaliteit. Door de recursie te conditioneren op de huidige verborgen toestand, de vaste prelude en ruis, bereikt het model functionele diversiteit zonder unieke gewichten.

De auteurs benadrukken dat RecurrentGPT de spanning tussen geheugenefficiëntie en expressiviteit oplost door een enkele gewichtstensor te laten fungeren als meerdere gespecialiseerde lagen. De resultaten tonen aan dat:

1. Parameter-efficiëntie: Significante parameterreducties (tot 63%) kunnen worden bereikt met minimale latentie-overhead.
2. Compute-Quality Tradeoff: De architectuur maakt een continue afweging mogelijk tussen inferentiesnelheid en kwaliteit via early exiting, waarbij één enkel checkpoint meerdere operationele punten dient.
3. Scaling Laws: De aanpak verbetert de scaling laws voor loss versus aantal parameters en loss versus training FLOPs, waarbij consistente winsten zichtbaar zijn op kleine, medium en grote schaal.

Het artikel concludeert dat hoewel de huidige implementatie een vaste recursiediepte gebruikt tijdens inferentie, het geleerde gating-mechanisme een robuuste basis biedt voor toekomstig werk op het gebied van dynamische per-token halting en kennisdistillatie.