Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers

Technische Samenvatting: Post-Grokking Collaps aan de Representatie–Readout Interface in Muon-getrainde Transformers

Probleemstelling
Het artikel onderzoekt een specifieke instabiliteit die optreedt in Transformer-modellen getraind op modulaire rekenkunde (specifiek modulaire optelling en aftrekking) nadat zij "grokking" hebben bereikt—het fenomeen waarbij generalisatie verschijnt lang nadat de trainingsnauwkeurigheid is verzadigd. Hoewel eerder werk heeft vastgesteld dat Muon (een orthogonaliserende optimizer) grokking versnelt vergeleken met AdamW, rapporteert deze studie dat Muon-getrainde modellen frequent lijden onder een recurrente, abrupte collaps in generalisatieprestaties post-grokking. Cruciaal is dat dit falen niet beperkt is tot Muon; AdamW-modellen vertonen ook instabiliteit. Het artikel demonstreert dat noch de optimizer strikt stabiel is over vijf seeds: de geselecteerde AdamW-baseline registreert sub-threshold evaluaties op vier van de vijf seeds (met een minimum van 27,59%), terwijl Muon dit op alle vijf de seeds registreert. Het kernprobleem is dat deze modellen het vermogen verliezen om te generaliseren ondanks dat de trainingsset nog steeds opgelost is, en standaard mechanistische progressiematen (spectrale ondersteuning en vermogensverdeling) het falen niet detecteren, aangezien zij een intact circuit rapporteren terwijl de nauwkeurigheid daalt naar bijna kansniveau.

Methodologie
De auteurs hanteren een split-optimizer trainingsregime waarbij de verborgen gewichtsmatrices worden bijgewerkt door Muon, terwijl de token/positie embeddings en de output unembedding (readout) worden bijgewerkt door AdamW. De studie maakt gebruik van modulaire rekenkunde taken ($a+b \pmod p$ en $a-b \pmod p$) met $p=113$ (en varianten met $p=97$) op decoder-only Transformers van variërende dieptes (1, 2 en 4).

Belangrijke methodologische componenten zijn:

1. Matched Branching: Trajecten worden vertakt vanuit bit-identieke parametertoestanden om causale factoren te isoleren. Eén tak vervolgt de standaard optimalisatie, terwijl andere takken specifieke parametergroepen bevriezen (verborgen matrices, embeddings of readout) om noodzakelijkheid te testen.
2. Fourier Decompositie: De auteurs deconstrueren de interne residu-toestanden van het model in 2D Fourier-modi over het operand-grid. Zij definiëren specifieke "families" van modi (bijv. de optellingsfamilie $(k, k)$, aftrekkingsfamilie $(k, -k)$) en gebruiken filtering-interventies om de voldoendeheid en noodzakelijkheid van deze families voor taakprestaties te meten.
3. Interventie Suite: De studie past interventies toe zoals fase-scrambling, frequentie-verplaatsing en cross-readout substitutie om de exacte aard van het geleerde algoritme en de bron van misalignement te bepalen.
4. Ablatie: Het genormaliseerde-orthogonalisatiepad van Muon wordt geableerd om de rol ervan in de geobserveerde spectrale dispersie en instabiliteit te testen.

Belangrijkste Bevindingen en Resultaten

• De Aard van de Collaps: Het falen is een "hidden–readout misalignement." Terwijl de trainingsverlies de nul nadert, verdwijnt de gradiënt, maar de twee optimizers reageren verschillend. Muon's update stapgrootte wordt grotendeels onafhankelijk van de gradiëntmagnitude (elasticiteit $\approx -0,03$), terwijl die van AdamW proportioneel blijft (elasticiteit $\approx +1,5$). Dit veroorzaakt dat de verborgen representatie en de readout uit elkaar drijven met een snelheid van 8,0 keer per parameter. De trainingsset blijft opgelost, maar de readout kan de representatie die zij samen met de training heeft ontwikkeld niet langer decoderen.
• Masking vs. Circuit Failure: Het artikel maakt onderscheid tussen twee foutmodi. In circuit failure stopt de Fourier-familie zelf met het berekenen van de taak. In circuit masking blijft de taak-gealigneerde familie (bijv. de optellingsfamilie) de taak perfect oplossen in isolatie (100% nauwkeurigheid), maar wordt de bijdrage ervan overstemd door een adversariële remainder in het volledige model. De auteurs tonen aan dat het herschalen van de taak-gealigneerde component alleen de volledige modelnauwkeurigheid herstelt, wat bewijst dat het circuit intact was maar werd "overstemd".
• Verschillen tussen Optimizers: Muon bereikt grokking aanzienlijk sneller (mediaan 13.011 stappen vs. 30.486 voor AdamW), maar is universeel instabiel over alle geteste configuraties. AdamW is ook instabiel over vijf seeds, hoewel de ernst van de collaps (gemeten als het aantal sub-threshold evaluaties) twee orde grootheden lager is dan bij Muon. In een single-seed sweep vertoonden AdamW-configuraties voorwaardelijke instabiliteit waarbij de ernst schaalde met de learning rate, maar replicatie over vijf seeds bevestigde dat de geselecteerde "stabiele" baseline niet strikt stabiel was over de seeds. De instabiliteit in Muon wordt toegeschreven aan de update-regel, die de representatie verspreidt over een breed spectrum (326 effectieve geconjugeerde paren vs. ~5 voor AdamW), waardoor de taak-gealigneerde familie een kleiner deel van het totale vermogen krijgt.
• Taakselectie van Familie: De taak, niet de optimizer, selecteert de Fourier-familie. Modellen getraind op optelling gebruiken de $(k, k)$ modi; modellen getraind op aftrekking gebruiken $(k, -k)$. De basis binnen de familie is exact en niet gedeeld; twee modellen getraind vanaf dezelfde seed divergeren naar het gebruik van verschillende frequenties en fasen, en hun readouts worden bijna orthogonaal, waardoor ze elkaars representaties niet kunnen decoderen.
• Diepte en Alignment: In diepere modellen wordt de Fourier-code geschreven door een MLP in de voorlaatste block, maar wordt deze pas "leesbaar" voor de unembedding in het laatste block. De misalignement vindt specifiek plaats op de interface tussen het laatste block's residu en de unembedding.
• Preventie: Het bevriezen van ofwel de verborgen matrices of de niet-verborgen parameters (embeddings en readout) na de vorming van het circuit voorkomt de collaps volledig. "Stable Muon", waarbij de embeddings en de readout na grokking worden bevroren, bereikt strikte stabiliteit over 451.400 post-grokking stappen en vijf seeds, terwijl standaard Muon herhaaldelijk faalt.

Betekenis en Claims
Het artikel claimt een fundamentele beperking in split-optimizer training en post-grokking dynamiek te hebben geïdentificeerd: de residu-stream heeft geen geprivilegieerde basis, en wanneer het verliesoppervlak afvlakt, biedt het verlies geen druk om de specifieke alignment tussen de representatie en de readout te behouden. De auteurs stellen dat:

1. Standaardmaten zijn onvoldoende: Spectrale maten (support en vermogensverdeling) kunnen invariant blijven tijdens een collaps omdat het falen een verandering van basis is (misalignment) in plaats van een verlies van het algoritme zelf.
2. Grokking is reversibele masking: De transitie van generalisatie naar collaps is de inverse van de transitie van memorisatie naar generalisatie; beide houden in dat het aandeel van de taak-gealigneerde familie in de representatieve kracht stijgt of daalt ten opzichte van de adversariële remainder.
3. De interface is de locus van falen: De instabiliteit is niet een eigenschap van de verborgen representatie of de readout alleen, maar van hun gezamenlijke identificatie. Het fixeren van één zijde van de interface stabiliseert het systeem.
4. Muon's snelheid komt met een prijs: Hoewel Muon de ontdekking van het circuit versnelt, creëert de update-regel een representatie die gevoeliger is voor misalignement met de readout onder standaard trainingscondities, wat interventie (bevriezen) noodzakelijk maakt voor langetermijnstabiliteit.

De studie stelt geen nieuwe toepassingen voor, maar biedt een mechanisch verslag van een specifieke foutmodus in mechanistische interpreteerbaarheid, waarbij benadrukt wordt dat een model een correct circuit kan bezitten terwijl het er niet in slaagt dit te gebruiken vanwege representatieve drift.