Per-Token Fixed-Point Convergence in Depth-Recurrent Transformers

Technische Samenvatting: Per-Token Fixed-Point Convergentie in Depth-Recurrent Transformers

Probleemstelling
Standaard fixed-depth transformers koppelen de modelkwaliteit aan één enkel compute-werkpunt. Het bedienen van meerdere kwaliteit/kosten-tiers vereist het trainen en opslaan van aparte modellen, terwijl het aanpassen van de compute aan de moeilijkheidsgraad van de input tijdens inferentie de prompt cache tenietdoet. Depth-recurrent transformers (die een gewichts-gebonden kern een variabel aantal keren toepassen) bieden een theoretische oplossing, waardoor één enkele parameterset kan opereren over een bereik van dieptes. Bestaande benaderingen voor per-token allocatie lijden aan significante beperkingen: early-exit methoden op niet-gedeelde stacks falen door inconsistenties in de key-value cache; mixture-of-recursions methoden fixeren het compute-budget in de gewichten tijdens de pretraining, wat post-hoc retargeting verhindert; en gezamenlijk getrainde halting-policies leiden vaak tot degeneratieve dieptes of vereisen delicate afstemming. Cruciaal is dat geen enkel voorgaand werk direct de per-token computation van een getraind elastisch model heeft gemeten om te bepalen of een eenvoudige, training-vrije policy beter zou kunnen presteren dan geleerde policies.

Methodologie
De auteurs stellen een tweestapsbenadering voor die diepte-robuustheid scheidt van diepte-allocatie:

1. Diepte-robuustheid (Pretraining): Een randomized-depth backbone wordt getraind met een minimaal recept. Bij elke optimizer-stap wordt een recursie-aantal $r$ gesampled uit een log-normale verdeling en toegepast op de hele stap. Het model gebruikt standaard cross-entropy loss zonder halting units, ponder loss of FLOP-straffen. De architectuur bestaat uit een unieke prelude, een gewichts-gebonden kern (met input re-injectie en sandwich normalisatie om state-collapse te voorkomen), en een unieke coda.
2. Diepte-allocatie (Post-hoc Policy): Zodra de backbone bevroren is, worden drie policies geëvalueerd op hetzelfde model om de per-token diepte te bepalen:
   • A0 (Uniforme Diepte): Past een vast loop-aantal $r$ toe op alle tokens (baseline).
   • A1 (Convergence Exit): Een training-vrije regel die een token stopt zodra de KL-divergentie tussen opeenvolgende output-distributies ($KL(p_i || p_{i-1})$) onder een drempelwaarde $\epsilon$ valt.
   • A2 (Oracle-Labeled Router): Een lineaire probe getraind op de eigen hidden states van het model om te voorspellen of de argmax-voorspelling van een token in toekomstige loops zal veranderen.

Belangrijkste Bijdragen

1. Directe Per-Token Meting: De studie biedt de eerste karakterisering van per-token convergentie in een depth-recurrent taalmodel. Het demonstreert dat de recurrente staat convergeert naar een per-token fixed point. Terwijl de mediane token convergeert bij loop zes, blijft een staart van ongeveer 10% van de tokens updaten bij de training-gemiddelde diepte van acht. Convergentiediepte is geordend naar token-type (whitespace convergeert het vroegst, content words het diepst).
2. Lezen versus Leren van Convergentie: Het paper demonstreert dat het direct lezen van de model-state dynamiek beter presteert dan het leren voorspellen ervan. De training-vrije convergence exit (A1) bereikt de kwaliteit van uniforme diepte-8 bij een gemiddelde van 4,94 loops (een reductie van 38% in gemiddelde diepte). In contrast hiermee vereist de geleerde lineaire router (A2), getraind op convergentie-labels afkomstig van hetzelfde model, bijna de volledige diepte en levert geen reductie op.
3. Training-Mean Saturation Hypothese: Ondersteund door interne data en externe consistentiecontroles, hypothetiseren de auteurs dat de nuttige diepte van een recursief model verzadigt bij de training-gemiddelde loop-distributie. Modellen getraind met een gemiddelde van 4 verzadigen bij diepte 4; modellen getraind met een gemiddelde van 8 verzadigen bij diepte 8.

Resultaten

• Elasticiteit: De randomized-depth backbone vertoont een monotone kwaliteitsverbetering van diepte 1 naar 8, waarna de validatie loss stabiliseert (3,20 nats) zelfs tot 32 loops. Dit bevestigt dat het model diepte-robuust is zonder specifieke afstemming voor elke diepte nodig te hebben.
• Convergentie Dynamiek: De gemiddelde opeenvolgende output KL-divergentie daalt van $3,9 \times 10^{-1}$ bij loop 2 naar $8,5 \times 10^{-6}$ bij loop 16. De verandering in de state neemt af per stap, wat bevestigt dat het systeem een fixed point bereikt in plaats van nutteloos rond te draaien.
• Allocatie-efficiëntie: De training-vrije A1-policy evenaart de kwaliteit van een uniforme diepte-8 model bij 4,94 gemiddelde loops. Het presteert beter dan de uniforme baseline bij ondiepe dieptes en komt overeen met de baseline over het gemiddelde dieptebereik. De geleerde router (A2) faalt in het leveren van besparingen en vereist bijna volledige diepte om de baseline te evenaren.
• Token-niveau Gedrag: Structurele tokens (whitespace) verlaten het proces het vroegst, terwijl content tokens langer rekenen. De relatieve ordening van cijfers en woorden keert om tussen datasets (TinyStories vs. FineWeb-Edu), toegeschreven aan distributieverschuivingen in plaats van architecturale fouten.
• Efficiëntie Grenzen: Hoewel de studie een reductie van 38% in gemiddelde diepte rapporteert (een FLOP-proxy), merkt het expliciet op dat de gerealiseerde wall-clock snelheid afhangt van diepte-gewijze continuous batching, wat niet is geïmplementeerd. De best-case wall-clock versnelling wordt geschat op 1,48–1,51x, terwijl de huidige harness 0% gerealiseerde versnelling laat zien door masking in plaats van het overslaan van werk.

Betekenis en Claims
Het paper claimt dat de "elasticiteit" van depth-recurrent transformers een stabiele eigenschap is die kan worden benut via eenvoudige, training-vrije policies. De centrale bevinding is dat de per-token variatie in convergentiediepte direct leesbaar en exploiteerbaar is, waardoor complexe geleerde halting-policies overbodig en minder effectief zijn. De auteurs positioneren dit werk als een metings- en allocatiestudie in plaats van een bijdrage aan de onderliggende elastische backbone (die zij toeschrijven aan eerder werk zoals Huginn en LoopFormer).

De studie is bescheiden in haar reikwijdte:

• Resultaten zijn vastgesteld op een enkele schaal (135M-klasse parameters) en een enkele seed (voor het hoofdmodel).
• Er wordt geen FLOP-matched parity geclaimd tegen standaard transformers; vergelijkingen worden gemaakt met een plafondreferentie (SmolLM2-135M) die een 150x groter aantal tokens heeft.
• De gerealiseerde wall-clock versnelling wordt niet gerapporteerd als een definitieve metriek, maar als een marge tussen FLOP-reductie en de best-case batching potentie.
• De hypothese dat de verzadigingsdiepte gelijk is aan de training-gemiddelde diepte wordt gepresenteerd als een open hypothese ondersteund door beperkte datapunten, niet als een bewezen wet.

De volledige studie was reproduceerbaar op een enkele RTX 4090 in ongeveer 100 GPU-uren.