The performernce of Muon optimizer worse than AdamW ?

[chenzixiong808-max]

Why does the Muon optimizer seem to perform worse than AdamW in my project? The error decreases very slowly, and the final error is no better than AdamW. My project is just a very small network with six layers: P is a convolutional layer, and the middle four layers are FNO networks—which can actually be thought of as convolution in the complex domain, but with 3D weight parameters. I don't think it will matter, as 4D models are trainable. Q is an MLP. This is the network, but it doesn't work well in practice. Has anyone else had the same question?

model.py
train.py

From this training screenshot, we can see that the error in the 50th round is still a little bit higher, but for Adam, the same time may drop to about 10e-3. I have uploaded the model and train files. Is there anyone who can help me?

[axion66]

Without knowing your setup_muon_optimizer code, it's hard to tell.

Are you only using Muon Optimizer? If that's the case, Muon expects input / output / embedding layers to not be included in its parameters. That is, you need a mix of AdamW & Muon, where you use MuonWithAuxAdam.

[parlance-zz]

Additionally, beyond what axion66 has said, you should be using a completely different learn rate schedule for Muon. The base learn rate should be significantly higher than what you would use with AdamW (think 10x to 50x), and the learn rate decay should be more aggressive than what you would use with AdamW (closer to 1/step decay).

Non-matrix parameters like biases and matrix parameters with a single input or output should be optimized with AdamW.

If you're using grouped convolutions you need to expose that information for Muon to be able to optimize each individual submatrix correctly. This isn't included in the vanilla implementation, but I have an implementation here if you're interested.

Lastly, you should be aware that vanilla Muon doesn't include any kind of momentum rescaling / normalization like AdamW, which means the absolute gradient norm affects the overall learn rate. Your loss should be an NLL or something measured in information units (nats / bits) that can provide that rescaling intrinsically. A simple example for rescaling a positive definite MSE loss would be adding a learnable logvar parameter, and then transforming the loss like: nll_loss = mse_loss / logvar.exp() + logvar. This logvar would be another example of a scalar parameter that should be optimized with AdamW.