模型訓練與最佳化
Muon Still Fails After Accelerating Grokking: Generalization Collapses Across All Nine Configurations
New experiments find that when Muon trains a Transformer’s hidden matrices while AdamW updates its embeddings and output head, the model learns modular arithmetic faster but cannot reliably retain the solution. Freezing either parameter group prevents collapse, suggesting that the problem stems from joint drift at the representation–readout interface.
Muon orthogonalizes momentum updates before applying them and has recently been used for hidden weight matrices in Transformers, while embeddings and the output head are handled by AdamW. A new study tested this division of labor on `(a+b) mod 113` using a decoder-only Transformer without normalization layers. It confirmed that the setup crosses the grokking threshold of 95% test accuracy more quickly. However, all nine Muon hyperparameter configurations subsequently fell back below the threshold, with some runs repeatedly dropping below 1% before recovering. AdamW was not entirely stable either, but its collapses depended more heavily on larger learning rates.
The authors traced the cause to the representation–readout interface. The final residual representation `h` and output matrix `W_U` determine the logits only through the product `W_U h`, so both can undergo a coordinated change of basis without altering the function. Applying two different update dynamics to them may gradually misalign bases that were previously coordinated. After the model had solved the training set, gradients had fallen to approximately `10^-6`, yet the per-parameter movement speed of the Muon parameter group remained eight times that of the AdamW group. Starting from exactly the same state and freezing either the hidden matrices or the embeddings and output head prevented failure. Across five runs with the latter group frozen, the researchers observed a cumulative 451,400 post-grokking steps without another drop below the threshold.
Fourier interventions also distinguished genuine damage to the computational circuit from “masking.” At some checkpoints, the task-aligned frequency component alone still achieved 100% accuracy, but the rest of the representation overwhelmed it, leaving the full model at only 45.85%. Simply amplifying that component restored accuracy to 99.9%. This indicates that tracking only frequency support or representational power may miss competition at the readout.
From an engineering perspective, the study covers only small modular-arithmetic Transformers, so its findings cannot be directly extrapolated to claim that large language models will collapse in the same way. The public repository includes training and intervention-analysis code, 191 CSV files, and 20 HTML summaries, but no checkpoints, meaning full reproduction still requires retraining. The next step is to examine relative update magnitudes across parameter groups at more realistic pretraining and fine-tuning scales, and to evaluate whether freezing or shared-basis constraints are more robust than fixed optimizer routing.