Replicated some of the results from this cool work with SGD, Muon, AdamW, SOAP
Yes! We train a transformer on cyclic sequences. SGD generalizes to an unseen skip. Muon doesn't — it memorizes each pattern separately instead of learning shared, generalizable representations.
No Digg Deeper questions have been answered for this story yet.
Some other curves + ce losses
Replicated some of the results from this cool work with SGD, Muon, AdamW, SOAP
As I understand it, the paper is saying that vanilla SGD learns modes sequentially in order of the magnitude of each mode in the data. In this case, the general skip rule has the highest mode followed by the task specific skips so SGD learns this general mode first.
The other optimizers precondition which kills this sequential-mode-learning-based-on-magnitude property. This keeps updates uniform and kills this curriculum.
In the case of muon, the update becomes isotropic so both the general high mode and the task specific smaller modes are learned all at once leading to poor generalization.
This is more evident in this sort of specifically constructed task but it probably doesnt matter for llm training since data is so diverse and vast that it regularizes against this.
(corrections welcome!)
Replicated some of the results from this cool work with SGD, Muon, AdamW, SOAP