Guidelines for whether to use muon for parameters

[insomniaaac]

As you mentioned in README, we should divide the model parameters into three parts:

hidden_weights = [p for p in model.body.parameters() if p.ndim >= 2]
hidden_gains_biases = [p for p in model.body.parameters() if p.ndim < 2]
nonhidden_params = [*model.head.parameters(), *model.embed.parameters()])

However, I still confused about whether to use muon for parameters of different model architecture.

1. if my model is vision transformer with a encoder, a decoder and a dpt head for depth prediction task, should i set the patchify layer as nonhidden_params? should i set the entire dpt head as nonhidden_params?
2. if my model is multi-modality llm, during alignment stage, should i optimize the projector (which usually a mlp projector) with muon? or it should be treated as a head?

Further more, would you like to offer a guideline for whether to use muon for different parameters?

[Arongil]

This blog post (linked from the README) has a helpful explanation of what's going on:

https://jeremybernste.in/writing/deriving-muon

TL;DR: Muon is designed for matrix multiplications where the input and output norm are usefully measured by the RMS norm, aka L2 / sqrt(dim). Hidden linear layers fit this criterion. However, the text embedding layer reasonably has an input norm of L1, so Muon isn't designed to work well for it. A text unembedding layer has a similar discrepancy. If you're deciding whether to use Muon for a particular layer in your network, ask yourself whether the input and output make sense to be measured in the RMS norm. If so, go for it. One heuristic is if the inputs and outputs are dense, which loosely means ||x||_RMS ≈ 1. For example the vector (1, 1, 1, ..., 1) is dense but (1, 0, 0, ..., 0) is not. For vision models with dense data (image frames), dense representations (hidden layers), and dense outputs (generated image frames or depth maps), Muon likely makes sense for all layers. Similarly, projecting from dense image space to dense representation space during an alignment phase also fits the criterion for Muon.