Weight Decay Regimes in Grokking Transformers: Cheap Online Diagnostics

Technische Samenvatting: Gewichtsvervalregimes in Grokking Transformers

Probleemstelling
Grokking verwijst naar het fenomeen waarbij transformers die zijn getraind op algoritmische taken (specifiek modulaire rekenkunde) een scherpe overgang vertonen van memorisatie naar generalisatie, lang nadat de trainingsfout verzadigd is. Hoewel mechanistische studies de vorming van interne circuits (bijvoorbeeld Fourier-representaties) en induktiekoppen hebben geïdentificeerd als drijvende krachten achter deze overgang, ontbreekt het in de bestaande literatuur aan een "goedkope online diagnose" om deze dynamiek in real-time te volgen. Bovendien is de rol van gewichtsverval ($\lambda$) als controleparameter voor deze regimes niet systematisch in kaart gebracht over modelgroottes heen, en is de specifieke bijdrage van coördinatie tussen attentiekoppen niet geïsoleerd van effecten van gewichtsgrootte.

Methodologie
De auteurs trainen decoder-only transformers op taken uit de modulaire rekenkunde (mod+, mod-, mod$\times$, mod$\div$) met behulp van de AdamW-optimizer. De studie maakt gebruik van een dichte scan van gewichtsvervalwaarden en varieert modelgroottes (van 0,82 miljoen tot 85 miljoen parameters) en architecturen (Transformers, MLP's, LSTMs, Mamba).

Belangrijke methodologische componenten zijn:

1. Online Ordeparameters: Twee scalair diagnostische waarden worden bij elke trainingsstap berekend met uitsluitend attentie-activaties:
   • Gemiddelde Paarsgewijze Attentie-Kop Cosinus-Soortgelijkheid ($\bar{s}(t)$): Meet de coördinatie (synchronisatie) tussen attentiekoppen.
   • Entropie Standaardafwijking ($\sigma_H(t)$): Meet de spreiding van Shannon-entropie over de koppen, en volgt differentiatie.
2. Finite-Size Scaling: De auteurs passen een logistische functie aan op de waarschijnlijkheid van grokking ($P(\text{grok})$) versus gewichtsverval om een kritieke drempel $\lambda_c$ te schatten. Ze passen ook een machtswet toe op de "tijd tot grokking" ($t_{\text{grok}}$) om de kritieke exponent $\nu$ te bepalen.
3. Causale Interventies: Om te onderscheiden tussen koppatroonstructuur en gewichtsgrootte, voeren de auteurs gepaarde experimenten uit waarbij attentiekoppen opnieuw worden geïnitieerd versus waarbij gewichtsnormen worden afgeknipt op het hoogtepunt van de differentiatiefase.
4. Formele Verificatie: Diagnostische identiteiten (grenzen en algebraïsche eigenschappen) zijn machine-gecontroleerd in Lean 4 om wiskundige correctheid te waarborgen.

Belangrijkste Bijdragen

1. Kwantitatieve Kritieke Drempel: De studie identificeert een scherpe kritieke drempel voor gewichtsverval voor de canonieke 4-laagse, 8-koppige transformer op mod+ taken: $\lambda_c = 0,0158$ (95% CI [0,0109, 0,0200]).
2. Regime-diagram: Een tweedimensionale kaart ($\lambda$, $N$) onthult drie distincte regimes:
   • Memorisatie: $\lambda < \lambda_c$ (bijna geen grokking).
   • Ontwikkelings-Grokking: $\lambda \ge \lambda_c$ (bereikt ~100% grokkingssnelheid, waarbij de tijd tot grokking afneemt naarmate $\lambda$ toeneemt).
   • Instorting: $\lambda = 10$ (alle koppen convergeren naar identieke patronen, $\bar{s}=1,0$).
3. Twee-fase Dynamiek: De trainingsbaan wordt gekenmerkt door:
   • Fase 1 (Synchronisatie): Koppen convergeren ($\bar{s}$ stijgt van ~0,93 naar 0,995), wat samenvalt met het begin van grokking.
   • Fase 2 (Differentiatie): Koppen divergeren ($\bar{s}$ daalt naar ~0,88, $\sigma_H$ neemt toe) terwijl de testnauwkeurigheid hoog blijft.
4. Per-kop Dimensie als Amplitudemodulator: De amplitude van differentiatie ($\sigma_H$) wordt empirisch gedomineerd door de per-kop dimensie ($d/H$), en monotoon verzadigend naarmate $d/H$ toeneemt.
5. Cross-architectuur Reikwijdte: De door gewichtsverval gecontroleerde overgang wordt waargenomen in niet-attentie architecturen (MLP, LSTM, Mamba), hoewel de kritieke drempel $\lambda_c$ aanzienlijk verschuift (een orde van grootte beslaand) afhankelijk van de architectuur.

Resultaten

• Schaal-exponent: De tijd tot grokking volgt een machtswet $t_{\text{grok}} \propto (\lambda - \lambda_c)^{-\nu}$ met een empirische exponent $\nu = 0,757$ (95% CI [0,725, 0,799]). De auteurs merken expliciet op dat deze waarde niet overeenkomt met standaard referentie-exponenten (bijvoorbeeld $\nu=1/2$ of 3D Ising $\nu \approx 0,63$) onder hun huidige grid-dichtheid, en stellen de identificatie van universaliteitsklassen uit voor toekomstig werk.
• Kwetsbaarheid op Lange Horizont: Bij verlengde horizons (20.000 epoch) wordt in sommige trajecten een zaad-afhankelijke "anti-grokking" instorting waargenomen, gekenmerkt door een vijf-staps patroon dat eindigt in afname van nauwkeurigheid. Dit wordt geïdentificeerd als een kwetsbaarheid en niet als een universeel cyclus.
• Causale Isolatie: Het opnieuw initialiseren van attentiekoppen vermindert de amplitude van Fase 2 ($\sigma_H$) aanzienlijk in vergelijking met controles, terwijl het afknippen van gewichtsnormen dit niet doet. Dit isoleert het effect tot koppatroonstructuur in plaats van gewichtsgrootte.
• Multi-taken Replicatie: Het door gewichtsverval gecontroleerde patroon geldt voor vier modulaire bewerkingen en twee modelgroottes, wat de robuustheid van het regime-diagram bevestigt, verder dan alleen optelling.

Betekenis en Claims
Het artikel claimt een goedkoop, online berekenbaar diagnostisch raamwerk te bieden voor het volgen van grokking-dynamiek met uitsluitend attentie-activaties, wat meer kostbare diagnostische methoden voor het leerlandschap, zoals de verfijnde lokale leercorrectie (rLLC), aanvult. Het stelt gewichtsverval vast als een primaire empirische controleparameter die memorisatie-, ontwikkelings-grokking- en instortingsregimes scheidt.

De auteurs zijn bescheiden in hun reikwijdte:

• Ze claimen geen universaliteitsklasse voor de waargenomen exponent $\nu \approx 0,76$, en merken op dat geteste referentie-exponenten buiten hun betrouwbaarheidsintervallen vallen.
• Ze claimen niet dat gewichtsverval noodzakelijk is voor generalisatie in alle architecturen, maar alleen dat het fungeert als een controleparameter binnen standaard attentie-architecturen.
• Ze stellen expliciet claims over causale mechanismen over waarom differentiatie optreedt uit voor vervolgonderzoek, en richten zich in plaats daarvan op wanneer en hoe het zich manifesteert.
• De bevindingen zijn beperkt tot modulaire rekenkunde in kleine transformer-attentiemodellen; claims over grote taalmodellen of bredere universaliteit vallen buiten de reikwijdte.

Het werk positioneert zichzelf als een fenomenologische en diagnostische bijdrage, en biedt een "regime-kaart" en "ordeparameters" die concurrent theoretisch werk over finite-size scaling, Lyapunov-contractie en glas-relaxatie aanvullen.