🤖 machine learning

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

Questo articolo dimostra che i transformer addestrati con Muon esibiscono un "collasso post-grokking" in cui la generalizzazione fallisce dopo l'apprendimento dell'aritmetica modulare a causa dell'instabilità all'interfaccia tra rappresentazione e readout, un fenomeno caratterizzato da dinamiche dell'ottimizzatore divergenti ed effetti di mascheramento che possono essere risolti congelando componenti specifiche del modello o riscalando il circuito allineato al compito.

Autori originali: Ali Janati, Kaoutar El Maghraoui, Andrei Kanavalau, Anass Belfatmi

Pubblicato 2026-08-10

📖 7 min di lettura🧠 Approfondimento

CC BY 4.0

Autori originali: Ali Janati, Kaoutar El Maghraoui, Andrei Kanavalau, Anass Belfatmi

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Immagina di insegnare a un robot a risolvere un enigma matematico molto specifico: sommare due numeri e trovare il resto della divisione per un numero primo molto grande. All'inizio, il robot sembra solo memorizzare le risposte, come uno studente che studia a memoria con le flashcard. Ma poi, accade qualcosa di magico chiamato "grokking". Improvvisamente, dopo un lungo periodo di apparente mancanza di progressi, il robot smette di memorizzare e capisce davvero il modello. Riesce a risolvere l'enigma perfettamente, anche con numeri che non ha mai visto prima. Questo è un grande evento per gli scienziati perché aiuta a capire come l'intelligenza artificiale impari a pensare invece di limitarsi a memorizzare.

Tuttavia, c'è un problema. A volte, dopo che il robot ha capito la soluzione, inizia a dimenticarla, anche se viene ancora addestrato sullo stesso enigma. È come uno studente che finalmente capisce l'algebra, ma poi inizia a fare errori banali sugli stessi problemi che aveva già risolto. Gli scienziati usano strumenti speciali per guardare dentro il "cervello" del robot per vedere cosa sta succedendo. Osservano i modelli specifici che il robot usa per risolvere la matematica (come le note musicali in una canzone) per vedere se il robot sta ancora usando le note giuste. La grande domanda è: perché il robot dimentica? È perché ha smesso di imparare le note, o sta succedendo qualcos'altro?

Il Mistero del Genio Svanito

In questo studio, i ricercatori hanno esaminato un tipo speciale di robot IA chiamato Transformer, addestrato a risolvere enigmi di addizione modulare. Hanno usato due diversi "insegnanti" (ottimizzatori) per insegnare al robot. Un insegnante, chiamato AdamW, è l'istruttore standard e affidabile. L'altro, chiamato Muon, è un insegnante più nuovo e veloce che sembra aiutare il robot a capire l'enigma molto più rapidamente.

I ricercatori hanno scoperto che Muon è davvero un fulmine. Aiuta il robot a raggiungere quel momento di "rivelazione" (grokking) in circa la metà del tempo necessario ad AdamW. Ma ecco il colpo di scena: le soluzioni che Muon trova sono incredibilmente fragili. Una volta che il robot risolve l'enigma, spesso inizia a fallire di nuovo, scendendo dal 100% di precisione quasi allo zero, per poi recuperare in seguito. Questo accade continuamente, anche se il robot viene ancora addestrato sullo stesso problema. AdamW non è del tutto sicuro, ma fallisce molto meno spesso e in modo molto meno severo.

La Stretta di Mano Interrotta

Quindi, cosa sta andando storto? I ricercatori hanno scoperto che il problema non è che il robot dimentica la matematica. Il robot conosce ancora perfettamente la soluzione all'interno del suo "cervello". Il problema è una stretta di mano interrotta tra due parti del robot: la parte che pensa (la rappresentazione) e la parte che risponde (il readout).

Immaginate che il cervello del robot sia una biblioteca piena di libri (i pensieri), e la bocca del robot sia il bibliotecario che legge i libri ad alta voce (la risposta).

La Biblioteca (gli strati nascosti) continua a scrivere nuove note sui libri ogni secondo.
Il Bibliotecario (la testa di output) sta cercando di leggere quei libri per dare la risposta.

Con l'insegnante Muon, la Biblioteca continua a riscrivere i libri in un linguaggio leggermente diverso ogni volta che si aggiorna, ma il Bibliotecario continua a cercare di leggerli nel vecchio linguaggio. All'inizio, riescono a capirsi. Ma poiché la Biblioteca cambia il suo linguaggio più velocemente di quanto il Bibliotecario riesca ad adattarsi, alla fine smettono di capirsi. La Biblioteca sta ancora scrivendo la matematica corretta, ma il Bibliotecario sta leggendo le parole sbagliate.

I ricercatori hanno dimostrato questo bloccando un lato. Quando hanno impedito alla Biblioteca di cambiare il proprio linguaggio, il robot è rimasto perfetto. Quando hanno impedito al Bibliotecario di cambiare il proprio stile di lettura, il robot è rimasto perfetto anche in quel caso. Ma quando entrambi erano liberi di muoversi, si sono allontanati l'uno dall'altro e il robot ha iniziato a fallire.

Il Fantasma nella Macchina

Ecco la parte più confusa: anche quando il robot sta fallendo, la "Biblioteca" contiene ancora la soluzione perfetta. I ricercatori hanno usato un filtro speciale per guardare solo la parte del cervello che conosce la matematica. Quando lo hanno fatto, la soluzione era ancora corretta al 100%!

Si scopre che il robot soffre di mascheramento del circuito (circuit masking). Immaginate che la Biblioteca stia gridando la risposta corretta, ma il resto del robot stia gridando una risposta sbagliata esattamente allo stesso volume. Il Bibliotecario sente un caos confuso e dà la risposta errata. La risposta corretta è ancora lì, ma viene sommersa dal rumore.

I ricercatori hanno dimostrato che se si aumenta semplicemente il volume della parte corretta (la famiglia della matematica) senza cambiare nient'altro, il robot torna istantaneamente al 100% di precisione. La matematica non è scomparsa; è solo andata persa nel rumore.

Perché gli Strumenti Standard hanno Fallito

Di solito, gli scienziati controllano se un robot funziona guardando le "note" che usa per risolvere il problema. Controllano se le note giuste sono presenti e se sono abbastanza forti. In questo studio, quegli strumenti dicevano che il robot stava bene. Le note erano lì, e suonavano quasi esattamente come prima del fallimento.

Ma il robot stava fallendo comunque. Questo perché gli strumenti stavano guardando le note ma non il volume o il linguaggio che il Bibliotecario stava usando. Le note erano perfette, ma il Bibliotecario aveva smesso di ascoltarle correttamente. È come controllare una radio e vedere che la stazione trasmette chiaramente, ma non rendersi conto che la radio è sintonizzata sulla frequenza sbagliata.

L'Approfondimento

I ricercatori hanno anche testato cosa succede quando rendono il robot più "profondo" (aggiungendo più strati di pensiero). Hanno scoperto che Muon è ancora più veloce nell'apprendimento, ma il problema della stretta di mano interrotta peggiora. Nei modelli più profondi, il cervello del robot scrive la matematica in una sezione, ma la parte che legge la risposta si trova in una sezione diversa, e si allontanano ancora più velocemente.

Hanno anche provato a correggere l'insegnante Muon rimuovendo il suo trucco speciale di "normalizzazione" (un modo per mantenere costanti gli aggiornamenti). Senza questo trucco, il robot ha imparato la matematica in modo molto stretto e concentrato, ma alla fine è crollato completamente e non è riuscito a recuperare. Ciò suggerisce che, sebbene il trucco speciale aiuti il robot a imparare velocemente e a diffondere la sua conoscenza, lo rende anche più propenso a questo specifico tipo di fallimento dovuto all'allontanamento.

La Conclusione

La lezione principale di questo articolo è che imparare velocemente può essere pericoloso se le parti del robot non sono sincronizzate. Il robot non ha dimenticato la matematica; ha solo perso la connessione tra la parte che conosce la matematica e la parte che la esprime.

I ricercatori hanno scoperto che se si blocca la parte che "parla" (l'output e gli embedding) dopo che il robot ha imparato la soluzione, il robot rimane stabile e non dimentica mai. Questo suggerisce che l'instabilità non deriva dal fatto che il robot sia cattivo in matematica, ma perché le due metà si muovono a velocità diverse e perdono il loro linguaggio comune.

In breve: il robot è un genio che improvvisamente inizia a balbettare perché il suo cervello e la sua bocca parlano dialetti diversi. La matematica è ancora perfetta all'interno, ma senza un modo per allineare i dialetti, il genio non può mostrare le sue abilità.

Sintesi Tecnica: Collasso Post-Grokking all'Interfaccia Rappresentazione–Readout in Transformer Addestrati con Muon

Enunciato del Problema
Il documento investiga una specifica instabilità che si verifica nei modelli Transformer addestrati sull'aritmetica modulare (specificamente addizione e sottrazione modulare) dopo aver raggiunto il "grokking"—il fenomeno per cui la generalizzazione emerge molto tempo dopo che l'accuratezza di addestramento ha saturato. Mentre lavori precedenti hanno stabilito che Muon (un ottimizzatore ortogonalizzante) accelera il grokking rispetto ad AdamW, questo studio riporta che i modelli addestrati con Muon soffrono frequentemente di un collasso ricorrente e brusco delle prestazioni di generalizzazione post-grokking. Fondamentalmente, questo fallimento non è limitato a Muon; anche i modelli AdamW mostrano instabilità. Il documento dimostra che né l'ottimizzatore né l'altro sono strettamente stabili su cinque seed: l'AdamW di riferimento selezionato registra valutazioni sotto la soglia in quattro seed su cinque (raggiungendo un minimo del 27,59%), mentre Muon le registra in tutti e cinque. Il problema centrale è che questi modelli perdono la capacità di generalizzare nonostante il set di addestramento rimanga risolto, e le misure standard di progresso meccanicistico (supporto spettrale e distribuzione della potenza) non riescono a rilevare il fallimento, riportando un circuito intatto anche quando l'accuratezza scende vicino al caso (chance levels).

Metodologia
Gli autori impiegano un regime di addestramento con ottimizzatore diviso (split-optimizer) in cui le matrici dei pesi nascosti sono aggiornate da Muon, mentre gli embedding di token/posizione e l'unembedding di output (readout) sono aggiornati da AdamW. Lo studio utilizza compiti di aritmetica modulare ( $a+b \pmod p$ e $a-b \pmod p$ ) con $p=113$ (e varianti con $p=97$ ) su Transformer decoder-only di profondità variabile (1, 2 e 4).

Componenti metodologiche chiave includono:

Branching Corrispondente (Matched Branching): Le traiettorie sono ramificate da stati dei parametri identici al bit per isolare i fattori causali. Un ramo continua con l'ottimizzazione standard, mentre altri congelano gruppi specifici di parametri (matrici nascoste, embedding o readout) per testarne la necessità.
Decomposizione di Fourier: Gli autori decompongono gli stati residui interni del modello in modi Fourier 2D sul reticolo degli operandi. Definiscono specifiche "famiglie" di modi (ad esempio, la famiglia dell'addizione $(k, k)$ , la famiglia della sottrazione $(k, -k)$ ) e utilizzano interventi di filtraggio per misurare la sufficienza e la necessità di queste famiglie per le prestazioni del compito.
Suite di Intervento: Lo studio applica interventi come lo scrambling della fase, la ricollocazione della frequenza e la sostituzione incrociata del readout per determinare la natura esatta dell'algoritmo appreso e la fonte del disallineamento.
Ablazione: Il percorso di ortogonalizzazione normalizzata di Muon viene [] ablatato per testare il suo ruolo nella dispersione spettrale e nell'instabilità osservate.

Risultati Chiave e Conclusioni

La Natura del Collasso: Il fallimento è un "disallineamento tra nascosto e readout" (hidden–readout misalignment). Mentre l'errore di addestramento si avvicina a zero, il gradiente svanisce, ma i due ottimizzatori rispondono diversamente. Il passo di aggiornamento di Muon diventa ampiamente indipendente dall'entità del gradiente (elasticità $\approx -0,03$ ), mentre quello di AdamW rimane proporzionale (elasticità $\approx +1,5$ ). Ciò causa la deriva del vettore di rappresentazione nascosta e del readout l'uno dall'altro a un tasso di 8,0 volte per parametro. Il set di addestramento rimane risolto, ma il readout non è più in grado di decodificare la rappresentazione con cui è stato addestrato.
Mascheramento vs Fallimento del Circuito: Il documento distingue due modalità di fallimento. In un fallimento del circuito, la famiglia di Fourier stessa smette di calcolare il compito. In un mascheramento del circuito, la famiglia allineata al compito (ad esempio, la famiglia dell'addizione) continua a risolvere il compito perfettamente in isolamento (accuratezza del 100%), ma il suo contributo viene sopraffatto da un resto avversario nel modello completo. Gli autori dimostrano che riscalare solo la componente allineata al compito ripristina l'accuratezza completa del modello, provando che il circuito era intatto ma è stato "voto fuori" (outvoted).
Differenze tra Ottimizzatori: Muon raggiunge il grokking significativamente più velocemente (mediana 13.011 step contro 30.486 per AdamW), ma è universalmente instabile in tutte le configurazioni testate. Anche AdamW è instabile attraverso cinque seed, sebbene la severità del collasso (misurata dal conteggio delle valutazioni sotto la soglia) sia di due ordini di grandezza inferiore rispetto a Muon. In uno sweep su singolo seed, le configurazioni AdamW hanno mostrato un'instabilità condizionale dove la severità scalava con il tasso di apprendimento, ma la replica su cinque seed ha confermato che il baseline "stabile" selezionato non era strettamente stabile tra i seed. L'instabilità in Muon è attribuita alla sua regola di aggiornamento, che diffonde la rappresentazione su uno spettro ampio (326 coppie coniugate effettive contro circa 5 per AdamW), lasciando alla famiglia allineata al compito una quota minore della potenza totale.
Selezione della Famiglia tramite il Compito: Il compito, non l'ottimizzatore, seleziona la famiglia di Fourier. I modelli addestrati all'addizione usano i modi $(k, k)$ ; quelli addestrati alla sottrazione usano $(k, -k)$ . La base all'interno della famiglia è esatta e non condivisa; due modelli addestrati dalla stessa seed divergono nell'utilizzare frequenze e fasi diverse, e i loro readout diventano quasi ortogonali, fallendo nel decodificare le rappresentazioni l'uno dell'altro.
Profondità e Allineamento: Nei modelli più profondi, il codice Fourier è scritto da un MLP nel penultimo blocco, ma diventa "leggibile" all'unembedding solo nell'ultimo blocco. Il disallineamento avviene specificamente all'interfaccia tra l'ultimo blocco del residuo e l'unembedding.
Prevenzione: Congelare o le matrici nascoste o i parametri non nascosti (embedding e readout) dopo che il circuito si è formato previene completamente il collasso. "Stable Muon", dove gli embedding e il readout vengono congelati post-grokking, raggiunge una stabilità rigorosa attraverso 451.400 step post-grokking e cinque seed, mentre il Muon standard fallisce ripetutamente.

Significato e Rivendicazioni
Il documento rivendica l'identificazione di un limite fondamentale nell'addestramento con ottimizzatore diviso e nelle dinamiche post-grokking: lo stream residuo non ha una base privilegiata e, quando la superficie di perdita si appiattisce, la perdita non fornisce alcuna pressione per mantenere l'allineamento specifico tra la rappresentazione e il readout. Gli autori sostengono che:

Le misure standard sono insufficienti: Le misure spettrali (supporto e distribuzione della potenza) possono rimanere invariate durante un collasso perché il fallimento è un cambiamento di base (disallineamento) piuttosto che una perdita dell'algoritmo stesso.
Il grokking è un mascheramento reversibile: La transizione dalla generalizzazione al collasso è l'inverso della transizione dalla memorizzazione alla generalizzazione; entrambi comportano che la quota di potenza rappresentativa della famiglia allineata al compito aumenti o diminuisca rispetto al resto avversario.
L'interfaccia è il luogo del fallimento: L'instabilità non è una proprietà della rappresentazione nascosta o del readout da soli, ma della loro identificazione congiunta. Fissare un lato dell'interfaccia stabilizza il sistema.
La velocità di Muon ha un costo: Sebbene Muon acceleri la scoperta del circuito, la sua regola di aggiornamento crea una rappresentazione che è più suscettibile al disallineamento con il readout in condizioni di addestramento standard, rendendo necessario l'intervento (congelamento) per la stabilità a lungo termine.

Lo studio non propone nuove applicazioni ma offre un resoconto meccanicistico di una specifica modalità di fallimento nella interpretabilità meccanicistica, evidenziando come un modello possa possedere un circuito corretto pur fallendo nell'utilizzarlo a causa di un drift rappresentazionale.