Technische Samenvatting: Dense Supervisie is Niet Genoeg: De Readout Blind Spot in Looped Language Models

1. Probleemstelling

Looped language models (LLM's) hergebruiken dezelfde transformer-block over recurrente diepte, waardoor de diepte tijdens de testfase kan fungeren als een dynamische rekenknop (bijv. vroegtijdig stoppen voor eenvoudige voorspellingen). In deze architectuur dient een verborgen toestand $H_k$ een dubbelrol: het wordt gedecodeerd voor voorspelling (de "readout" interface) en gevoed als input voor toekomstige berekeningen (de "recurrente" toestand).

Het centrale probleem dat wordt geïdentificeerd, is een zichtbaarheid–activiteitsmismatch. Hoewel dense supervisie (het toepassen van cross-entropy loss bij elke loop) ervoor zorgt dat intermediaire toestanden bruikbaar zijn voor voorspelling, controleert dit niet noodzakelijkerwijs alle variabelen die actief zijn binnen de recurrente transitie. Specifiek onderzoekt het paper de schaal van de verborgen toestand (de radiale magnitude van de verborgen vector).

In standaard pre-norm residual loops met schaal-invariante readouts (bijv. RMSNorm of LayerNorm), normaliseert de readout de verborgen toestand, waardoor de directe cross-entropy loss ongevoelig is voor de radiale schaal van de input. Echter, het recurrente pad (de skip-verbinding en de residual update) blijft deze schaal dragen en bijwerken. Bij consequent gevolg kan het model bij elke loop getraind worden terwijl de normen van de verborgen toestandstoestanden ongecontroleerd naar duizenden of tienduizenden driften, wat de "readout blind spot" tot een faalmodus maakt waarbij dense supervisie faalt om actieve recurrente variabelen te beheersen.

2. Methodologie en Mechanisme

Theoretisch Mechanisme

De auteurs formaliseren de faalmodus via twee lemma's:

1. Visibility Lemma: Als een readout $r$ schaal-invariant is (d.w.z. $r(\alpha H) = r(H)$ voor $\alpha > 0$), is de directe cross-entropy loss $L$ een radiale afgeleide van ongeveer nul ($\langle \nabla_H L, H \rangle \approx 0$). De loss kan de schaal van de verborgen toestand niet "zien" of corrigeren.
2. Activity Lemma: In pre-norm residual updates is de schaal actief. De updatefunctie $F(su) = su + b(u)$ (waarbij $s$ de schaal is en $u$ de richting) staat toe dat de residual branch $b(u)$ de schaal bij de leidende orde bijwerkt. Zelfs als de readout de schaal verbergt, behoudt de recurrentie deze en kan deze zelfs versterken.

Dit creëert een lokale mismatch: het loss-signaal kan de geleerde radiale snelheid van de recurrente update niet tegenwerken, wat leidt tot schaaldrift.

Experimentele Opzet

De auteurs trainden autoregressieve looped transformers op de WikiText-103 dataset met $K=4$ recurrente toepassingen van een gedeelde decoder stack. Ze gebruikten twee modelgroottes: 44M en 129M parameters.

• Core Ablation: Een $2 \times 2$ design waarbij Loss Placement (Terminal-only CE versus Per-loop CE) en Readout Type (genormaliseerde RMSNorm versus Raw/Unnormalized) worden gekruist.
• Controles:
  • Norm Penalty: Het toevoegen van een expliciete hulp-loss term $\lambda \sum \|H_k\|^2$ om de schaal zichtbaar te maken zonder de readout te veranderen.
  • Final-only Norm: Het gebruik van raw readouts voor intermediaire loops maar een genormaliseerde readout voor de finale output.
  • Recurrent Scale Clamp: Het handmatig herschalen van de verborgen toestanden naar hun initiële loop-1 magnitude om te testen of de geaccumuleerde schaalgroei voorspellend is.
• Evaluatie: Gemeten final-loop verborgen toestand-normen, perplexiteit (PPL) en de "compute-quality frontier" (PPL versus throughput) bij variabele inferentie-dieptes ($K=1$ tot $K=4$).

3. Belangrijkste Bijdragen

1. Identificatie van de Readout Blind Spot: Het paper definieert een specifieke faalmodus waarbij dense cross-entropy supervisie wel de voorspellingsinterfaces traint, maar faalt in het controleren van actieve recurrente variabelen (specifiek schaal) als die variabelen worden verborgen door schaal-invariante readouts.
2. Validatie van het Mechanisme: Door middel van ablaties en diagnostiek tonen de auteurs aan dat:
   • Per-loop CE met genormaliseerde readouts leidt tot enorme normdrift (bijv. normen die 39.000+ bereiken in het 44M model).
   • Raw readouts of expliciete norm-penalties deze normen reduceren tot de tientallen, zelfs met per-loop CE.
   • De drift op token-niveau plaatsvindt en heavy-tailed is, en niet slechts een verschuiving van het gemiddelde.
3. Ontwerpregel: De auteurs stellen een eenvoudige ontwerpregel voor: recurrent schaalbeheer vereist ofwel het zichtbaar maken van de schaal van de verborgen toestand voor een loss (via raw readouts of norm penalties), of het verwijderen ervan uit het recurrente pad (via inter-loop normalisatie). Dense CE alleen is onvoldoende.
4. Variabele-Diepte Prestaties: Schaal-gecontroleerde modellen (gebruikmakend van raw readouts of norm penalties) behouden bruikbare vroege exits terwijl ze een lagere perplexiteit bereiken bij gelijke inferentie-dieptes vergeleken met schaal-invariante baselines. De baselines vertonen vaak "K-invariance" (prestaties verbeteren niet met meer loops) omdat de schaaldrift de hoekverbeteringen verkleint ($1/s_k$), wat de genormaliseerde toestand effectief bevriest.

4. Resultaten

• Norm Drift: In 44M en 129M modellen met per-loop CE en RMSNorm readouts, driften de finale verborgen toestand-normen naar respectievelijk 39.0 de 56.051. In contrast hiermee hielden modellen met raw readouts of norm penalties de normen in de range van 17–83.
• Gradient Diagnostics: Radiale gradiëntfracties voor genormaliseerde readouts waren bijna nul ($\sim 10^{-9}$), terwijl raw readouts significante radiale sensitiviteit vertoonden ($\sim 10^{-2}$).
• Perplexity en Diepte:
  • RMSNorm Baseline: Modellen waren bijna K-invariant; stoppen bij $K=1$ leverde een vergelijkbare perplexiteit op als $K=4$. Dit resulteerde in een hoge schijnbare versnelling, maar duidde op een gebrek aan nuttige diepte-benutting.
  • Scale-Controlled Variants: Deze modellen vertoonden een significante verbetering in perplexiteit naarmate de diepte toenam ($K=1$ tot $K=4$). Ze bezetten een superieure PPL-throughput frontier, waarbij ze een lagere perplexiteit boden bij dezelfde throughput vergeleken met de RMSNorm baseline.
• Scale Clamp: Het kunstmatig vastzetten (clamping) van de schaal van de verborgen toestanden in getrainde modellen resulteerde in verwaarloosbare veranderingen in cross-entropy, wat suggereert dat de geaccumuleerde schaalgroei zelf niet de primaire bron van voorspellende informatie is, maar eerder een bijproduct van de ongecontroleerde recurrentie.

5. Betekenis en Claims

Het paper claimt dat de "readout blind spot" een fundamentele architecturale overweging is voor looped en adaptieve-diepte modellen. De betekenis ligt in de verschuiving van de focus van enkel het toepassen van dense supervisie naar het waarborgen dat het supervisiesignaal zichtbaar is voor de specifieke variabelen die in de recurrentie worden bijgewerkt.

• Praktische Implicatie: Voor practitioners die genormaliseerde readouts gebruiken, suggereert het paper het toevoegen van een expliciete norm penalty als een "drop-in" baseline om de schaal te controleren zonder de output-interface te veranderen.
• Theoretisch Inzicht: Het werk verheldert dat dense supervisie de interface (de readout) traint, maar niet automatisch de toestand-dynamiek (de recurrentie) reguleert als de interface de magnitude van de toestand maskeert.
• Scope Beperking: De auteurs merken bescheiden op dat hun conclusies gebaseerd zijn op specifieke instellingen (WikiText-103, 44M/129M modellen, teacher-forced evaluatie) en dat het mechanisme lokaal is. Ze beweren niet dat ze alle stabiliteitsproblemen in looped modellen hebben opgelost, maar adresseren specifief de onder-geconstrueerde natuur van de schaal in schaal-invariante readout loops.

Het paper concludeert dat adaptive-depth looped modellen zowel getrainde exits (via per-loop loss) als een mechanisme vereisen om de actieve recurrente schaal te controleren (via zichtbaarheid of verwijdering), en dat dit orthogonale ontwerpvereisten zijn.