Trasformatori Ottimizzati con Muon Mostrano Collasso Post-Grokking all'Interfaccia Rappresentazione-Readout

Un recente preprint su arXiv (2608.07436) rivela che i modelli transformer che utilizzano l'ottimizzatore Muon per compiti di addizione modulare subiscono un collasso post-grokking, con conseguente perdita della generalizzazione precedentemente raggiunta. La ricerca impiega una configurazione in cui Muon ottimizza le matrici nascoste mentre AdamW si concentra sugli embedding e sulla testa di output. Mostra che tutti e nove gli setup per (a+b) mod 113 inizialmente grokkano ma successivamente perdono la loro generalizzazione. Tra cinque semi, il riferimento AdamW scende sotto la soglia in quattro casi, raggiungendo solo il 27,59% di accuratezza. Questa instabilità è evidente attraverso vari moduli, larghezze, frazioni di addestramento, sottrazione e profondità. Il fallimento si verifica all'interfaccia rappresentazione-readout, che è identificabile solo congiuntamente attraverso una mappa invertibile non selezionata dalla perdita. Dopo aver padroneggiato il set di addestramento, il gradiente diminuisce fino all'ordine 10^-6, con risposte distinte dagli ottimizzatori: Muon mostra un'elasticità del passo di -0,03 rispetto a +1,5 per AdamW, e il gruppo Muon progredisce 8,0 volte più velocemente per parametro. Prevenire il fallimento può essere ottenuto congelando uno dei due gruppi da stati bit-identici, mentre congelare embedding/readout elimina il problema in cinque esecuzioni su 451.400 passi post-grokking e cinque p.