Spectral descent was the slowest method at 8,225 steps.
I've added two optimizers to the public benchmark:
(1) Shampoo (with its original 1/4 power). (2) Spectral descent, which is equivalent to both Muon(mu=0) and Shampoo(b1=b2=0).
Result: Shampoo falls halfway between Muon & Adam; Spectral descent is ~2x slower.
Thread below 1/6
One motivation for me to add these optimizers was seeing @_arohan_, a senior researcher whose contributions I respect, repeatedly claim that Muon is Shampoo.
IIUC, his argument is that if we disable accumulation in both Muon and Shampoo, then they become... 4/6
For this comparison, I kept Shampoo's exponent at its original value of 1/4. It is well-known that the modern "Shampoo^2" variant which uses 1/2 is more efficient, but this modification breaks the mathematical relationship to Muon and Spectral Descent, so I kept the original. 3/6
...equivalent. This is correct. The problem is that Muon without accumulation is not Muon: It is Spectral Descent, which is >2x slower.
To go fast we need accumulation, and -- as shown in the figure -- the way it's added is what makes the difference between Muon and Shampoo. 5/6
One motivation for me to add these optimizers was seeing @_arohan_, a senior researcher whose contributions I respect, repeatedly claim that Muon is Shampoo.
IIUC, his argument is that if we disable accumulation in both Muon and Shampoo, then they become... 4/6
Being a bit pedantic and finally got a few mins to run this. kthx bye.
I remember in August 2024 (before the anthology paper) when I was attending the modula-in-numpy sessions of Scale ML it was even called ShampooLinear
Here you go, sir. Muon is a good optimizer.
I think Keller's attempt at implementing it is great -- this primarily helped me look at why his effort produced much worse looking curve than what I could get.
This ended up being a nerdsnipe into hyper parameters, grafting, and eigh calls in the distributed shampoo package from Meta in my car ride home.
The main delta's are below:
I've added two optimizers to the public benchmark:
(1) Shampoo (with its original 1/4 power). (2) Spectral descent, which is equivalent to both Muon(mu=0) and Shampoo(b1=b2=0).
Result: Shampoo falls halfway between Muon & Adam; Spectral descent is ~2x slower.
Thread below 1/6
@_arohan_ My response to the claim that Muon is a renamed version of Shampoo
I've added two optimizers to the public benchmark:
(1) Shampoo (with its original 1/4 power). (2) Spectral descent, which is equivalent to both Muon(mu=0) and Shampoo(b1=b2=0).
Result: Shampoo falls halfway between Muon & Adam; Spectral descent is ~2x slower.
Thread below 1/6
Citations:
Shampoo is Gupta et al. (2018) https://arxiv.org/abs/1802.09568 and Anil et al. (2020) https://arxiv.org/abs/2002.09018
Spectral descent is Carlson et al. (2015a) https://proceedings.mlr.press/v38/carlson15.html and Carlson et al. (2015b) https://papers.nips.cc/paper_files/paper/2015/hash/f50a6c02a3fc5a3a5d4d9391f05f3efc-Abstract.html 6/6
...equivalent. This is correct. The problem is that Muon without accumulation is not Muon: It is Spectral Descent, which is >2x slower.
To go fast we need accumulation, and -- as shown in the figure -- the way it's added is what makes the difference between Muon and Shampoo. 5/6
Reproducible logs:
Shampoo: https://github.com/KellerJordan/modded-nanogpt/blob/master/records/track_3_optimization/results/20260513_shampoo_1_4_power/503575c5-6dde-425a-b461-2df4d99db974.txt Spectral descent: https://github.com/KellerJordan/modded-nanogpt/blob/master/records/track_3_optimization/results/20260517_ortho/d5098d67-7c1b-47b4-8833-80960d633d33.txt
As part of the public benchmark, further hyperparameter improvements are welcomed for any of these runs. All four use the same WSD lr schedule. 2/6
I've added two optimizers to the public benchmark:
(1) Shampoo (with its original 1/4 power). (2) Spectral descent, which is equivalent to both Muon(mu=0) and Shampoo(b1=b2=0).
Result: Shampoo falls halfway between Muon & Adam; Spectral descent is ~2x slower.
Thread below 1/6
For this comparison, I kept Shampoo's exponent at its original value of 1/4. It is well-known that the modern "Shampoo^2" variant which uses 1/2 is more efficient, but this modification breaks the mathematical relationship to Muon and Spectral Descent, so I kept the original. 3/6
Reproducible logs:
Shampoo: https://github.com/KellerJordan/modded-nanogpt/blob/master/records/track_3_optimization/results/20260513_shampoo_1_4_power/503575c5-6dde-425a-b461-2df4d99db974.txt Spectral descent: https://github.com/KellerJordan/modded-nanogpt/blob/master/records/track_3_optimization/results/20260517_ortho/d5098d67-7c1b-47b4-8833-80960d633d33.txt
As part of the public benchmark, further hyperparameter improvements are welcomed for any of these runs. All four use the same WSD lr schedule. 2/6
I went to a few optimization talks at 2018 ICML and I'm also reminded of the GGT optimizer (same as Shampoo^2?). In my mind the biggest deviation of Muon is that breaks with the "preconditioning" view. It's not preconditioning to normalize gradient by gradient in the same batch. The preconditioning direction is still open, how to best utilize correlation statistics from training history
@kellerjordan0 @_arohan_ Muon and Shampoo are indeed related but definitely not to the extent to call them the same.
@kellerjordan0 @jeffreycider @_arohan_ This paper suggests Muon is performing the input side whitening of shampoo, so like “Left-side-only-shampoo”
https://arxiv.org/abs/2604.01472
One motivation for me to add these optimizers was seeing @_arohan_, a senior researcher whose contributions I respect, repeatedly claim that Muon is Shampoo.
IIUC, his argument is that if we disable accumulation in both Muon and Shampoo, then they become... 4/6
First one is: don't take square roots of small eps. do psuedo inverse
inv_power_L = L.pow(-1.0 / root)
to:
positive_eigenvalue_mask = L > 1e-15 inv_power_L = torch.zeros_like(L) inv_power_L[positive_eigenvalue_mask] = L[positive_eigenvalue_mask].pow(-1.0 / root)
Here you go, sir. Muon is a good optimizer.
I think Keller's attempt at implementing it is great -- this primarily helped me look at why his effort produced much worse looking curve than what I could get.
This ended up being a nerdsnipe into hyper parameters, grafting, and eigh calls in the distributed shampoo package from Meta in my car ride home.
The main delta's are below:
Second one is that you do Nesterov momentum from Ilya/Marten/Hinton’s paper on importance of momentum from 2011 and pass that into shampoo.
First one is: don't take square roots of small eps. do psuedo inverse
inv_power_L = L.pow(-1.0 / root)
to:
positive_eigenvalue_mask = L > 1e-15 inv_power_L = torch.zeros_like(L) inv_power_L[positive_eigenvalue_mask] = L[positive_eigenvalue_mask].pow(-1.0 / root)
Third one is that that if [m, n] tensor take the side thats larger, especially good because other there are some numerical instability, since its low rank/ plus indefiniteness creeps in. Didnt spend too much time.
Second one is that you do Nesterov momentum from Ilya/Marten/Hinton’s paper on importance of momentum from 2011 and pass that into shampoo.
Last one is that with Adam grafting from Meta’s impl, means the size of update is O(sqrt(size)) - which you have to set different lr and weigth decay. The Muon implementation uses different lr / wd for various layers. I just used it, and rescaled it as appropriate.
Third one is that that if [m, n] tensor take the side thats larger, especially good because other there are some numerical instability, since its low rank/ plus indefiniteness creeps in. Didnt spend too much time.
@kellerjordan0 do you use a smaller training subset when iterating on optimizer architectures? I'm wondering what's the fastest training run that would give useful information for comparing optimizer architectures
@kellerjordan0 Muon: Civil War
@kellerjordan0 Why not also add Shampoo^2 curve just for completeness?
Spectral descent was the slowest method at 8,225 steps.
I've added two optimizers to the public benchmark:
(1) Shampoo (with its original 1/4 power). (2) Spectral descent, which is equivalent to both Muon(mu=0) and Shampoo(b1=b2=0).
Result: Shampoo falls halfway between Muon & Adam; Spectral descent is ~2x slower.
Thread below 1/6
