Muon Optimizer Guide: Getting Started Quickly and Key Details
Recently, I imagine many readers have come across news about the Muon optimizer. In fact, Muon was first proposed around October last year by Keller Jordan on Twitter, so it's only been a little over a year since then. Yet in that single year, Muon has already been battle-tested on models with tens of billions, hundreds of billions, and even trillions of parameters — more than enough to show that it's a genuinely competitive optimizer.
Today, Muon has been built into training frameworks like Torch and Keras, and even a heavyweight framework like Megatron has gradually started supporting it, which means it has already won widespread recognition across the industry. That said, for readers who are only familiar with Adam, figuring out how to quickly and effectively switch over to Muon can still be a confusing affair. So this post tries to offer a quick-start guide.
A Brief Introduction
The person who formally proposed Muon is Keller Jordan, who currently works at OpenAI. As mentioned at the start, Muon first appeared on Twitter, and even now the author has only written a follow-up blog post, Muon: An optimizer for hidden layers in neural networks, rather than a formal paper. The author's own view is that "whether or not it's written up as a paper has nothing to do with whether the optimizer actually works" [original text]. more
Muon is an optimizer custom-built for matrix parameters. There are some related lines of work with similar characteristics, such as Shampoo, and the even earlier Stochastic Spectral Descent, among others. Many of these works can be connected to Muon to a greater or lesser extent, but none of them fully covers what Muon does, so in my view Muon counts as a genuinely new piece of work.
In China, the earliest article introducing Muon to a broad audience should be my own blog post Appreciating the Muon Optimizer: The Essential Leap from Vectors to Matrices, and the first validation of Muon at a reasonably large scale should be our Moonlight release back in February. The Moonlight variant of Muon proposed there was later used in the trillion-parameter K2 model. After K2, GLM-4.5 likewise adopted this Muon variant.
As Jeremy Bernstein, one of the co-authors of Muon, put it in his blog post Deriving Muon — and I share this view — what makes Muon special is that it can be derived from more fundamental optimization principles, and it works well in practice. By contrast, while Adam is also effective, it feels more like a heuristic scheme.
Four Versions
This post isn't going to go into the mathematical details of Muon, nor its implementation; instead, it will mainly cover the technical details and caveats involved in switching from Adam to Muon. As just mentioned, Muon is specifically designed for optimizing matrix parameters, and its update rule is not element-wise — which can be a source of confusion for new users getting started.
Also, as far as I know, there are currently at least four slightly different versions of Muon in circulation, and this multiplicity of versions only adds to the confusion. If a user doesn't understand these details, they may end up with poor results simply from mis-tuning the hyperparameters (especially the learning rate). Let's clear this up below. First, for a matrix $\boldsymbol{W}\in\mathbb{R}^{d_{in}\times d_{out}}$, with $\boldsymbol{G}$ being its gradient, the four Muon variants are:
$$\begin{aligned}\newcommand{msign}{\mathop{\text{msign}}} &\quad\boldsymbol{M}_t \quad=\quad \beta \boldsymbol{M}_{t-1} + \boldsymbol{G}_t \\[10pt] &\quad\boldsymbol{W}_t = \left\{ \begin{aligned} &\boldsymbol{W}_{t-1} - \eta_t \left(\msign(\boldsymbol{M}_t) + \lambda \boldsymbol{W}_{t-1}\right) &\color{skyblue}{(\text{naive version})} \\[5pt] & \boldsymbol{W}_{t-1} - \eta_t \left(\sqrt{\max(1, d_{out}/d_{in})}\msign(\boldsymbol{M}_t) + \lambda \boldsymbol{W}_{t-1}\right) &\color{skyblue}{(\text{KellerJordan version})} \\[5pt] & \boldsymbol{W}_{t-1} - \eta_t \left(\sqrt{ d_{out}/d_{in}}\msign(\boldsymbol{M}_t) + \lambda \boldsymbol{W}_{t-1}\right) &\color{skyblue}{(\text{MuP version})} \\[5pt] & \boldsymbol{W}_{t-1} - \eta_t \left(0.2\times\sqrt{\max(d_{out},d_{in})}\msign(\boldsymbol{M}_t) + \lambda \boldsymbol{W}_{t-1}\right) &\color{skyblue}{(\text{Moonlight version})} \end{aligned}\right. \end{aligned}$$
If you want to enable Nesterov momentum, replace $\msign(\boldsymbol{M}_t)$ with $\msign(\beta\boldsymbol{M}_t + \boldsymbol{G}_t)$, where $\msign$ is usually named zeropower_via_newtonschulz in implementations — ordinary users don't need to worry about these implementation details.
The only difference among the four versions is the scaling factor in front of $\msign$. The "Keller Jordan version" and the "MuP version" are largely similar, while the "Moonlight version" is a bit different. Keras only implements the "Keller Jordan version," while Torch implements both the "Keller Jordan version" and the "Moonlight version." The plain/naive version seems fairly rare in practice; the one I personally use most often is the "MuP version" that I wrote myself.
Two Dimensions
There's an important detail we need to pay attention to here: both the "Keller Jordan version" and the "MuP version" are sensitive to the order of $d_{in},d_{out}$, so the first thing to get straight is what $d_{in},d_{out}$ actually means — it is not necessarily the case that the first dimension of the matrix is always $d_{in}$ and the second dimension is always $d_{out}$.
$d_{in}$ and $d_{out}$ refer respectively to the input and output dimensions of a linear layer, so which one is $d_{in}$ and which one is $d_{out}$ depends on how the linear layer is actually implemented. For instance, Keras's Dense layer implements $\boldsymbol{x}\boldsymbol{W}$, so for the matrix $\boldsymbol{W}$ the first dimension is $d_{in}$ and the second is $d_{out}$. Torch's Linear layer, however, implements $\boldsymbol{x}\boldsymbol{W}^{\top}$, so for the matrix $\boldsymbol{W}$ the second dimension is $d_{in}$, and it's the first dimension that is $d_{out}$.
So, to implement the "Keller Jordan version" of Muon, for Torch's Linear layer the scaling factor should be max(1, W.shape[0]/W.shape[1])**0.5, whereas for Keras it should be max(1, W.shape[1]/W.shape[0])**0.5. This means the current Muon implementation in Keras (version 3.12) is actually incorrect, since it directly copied Torch's scaling-factor implementation [source code].
If you're working with a model you wrote yourself, you'll need to judge carefully based on your own implementation — for example, it's entirely possible to mix Torch's built-in Linear layer with a hand-written x @ W, in which case you can't simply assume it's W.shape[0]/W.shape[1] or W.shape[1]/W.shape[0] across the board. Of course, if you'd rather not bother sorting all this out, you can just use the "Moonlight version," whose scaling factor is symmetric with respect to $d_{in},d_{out}$.
Hyperparameter Settings
Once $d_{in},d_{out}$ is sorted out, all that remains is deciding on the learning rate $\eta_t$ and the weight decay coefficient $\lambda$. Here we assume the user already has experience tuning Adam, has gotten good results with it, and now wants to quickly move over to Muon to try it out.
Let's start with the "Moonlight version." Its scaling factor is obtained by matching Adam's update RMS; if you want the details, see Muon Sequel: Why Did We Choose to Try Muon?, and for the "magic number" $0.2$, see Why Is Adam's Update RMS 0.2?. In short, the "Moonlight version" of Muon matches Adam's update magnitude, so the simplest way to migrate from Adam is: change nothing — just reuse Adam's $\eta_t$ and $\lambda$.
Now let's look at the other three versions. As we know, mainstream models typically have a hidden_size (denote it $d$), and the shapes of most matrices in the model don't deviate much from $d\times d$, so we can approximate using $d_{in}=d_{out}=d$. Under this approximation, these three versions all coincide, differing from the "Moonlight version" only by a factor of $0.2\sqrt{d}$. Since the "Moonlight version" matches Adam's update magnitude without requiring any hyperparameter changes, it follows that for these three other versions, the learning rate should be scaled up by a factor of $0.2\sqrt{d}$ in order to match Adam's update magnitude; correspondingly, $\lambda$ should be divided by $0.2\sqrt{d}$.
Substituting $d=1024,2048,4096$, the results are $6.4, 9, 12.8$ respectively. If you can't remember $0.2\sqrt{d}$, a simple rule of thumb is: when using one of the other three versions of Muon, just multiply Adam's learning rate by 10 to get Muon's learning rate. If you plug Adam's learning rate directly into Muon without adjustment, you'll end up underfitting and conclude that Muon is much worse than Adam — as far as I know, some of the negative reviews of Muon stem from exactly this mistake.
Does this mean the "Moonlight version" is simply more convenient to use? The "Moonlight version" does indeed perform well in practice, but if we say it's "more convenient" on that basis, we're really evaluating it from Adam's point of view. The advantage of the "MuP version" or "Keller Jordan version" is that the learning rate transfers across scales — that is, once you've tuned the learning rate on a small model, using the same value directly on a large model often still works well. For more on this, see Jeremy Bernstein's blog post Deriving Muon or my own post Higher-Order MuP: A Simpler Yet Smarter Spectral-Condition Scaling.
Other Parameters
If Muon only handles matrix parameters, what about everything else? For example, the bias term in a linear layer, the gamma parameter in RMSNorm — these are 1D parameters; and convolutional layers might have 3D or 4D array parameters.
Let me correct something here first: Muon doesn't only handle matrix parameters in general — it only handles the "matrix parameters of linear layers with dense inputs." If that sounds a bit abstract, just remember this: the matrix parameters of the embedding layer and the final classification layer (including the GPT LM head) should not use Muon, or performance will noticeably suffer. For these matrix parameters that shouldn't use Muon, as well as any 1D, 3D, or higher-dimensional parameters, if you don't want to spend too much effort thinking it through, just use Adam for them — most Muon implementations already mix in Adam, letting users choose which layers use Adam.
If you're willing to tinker, then 3D and 4D parameters like convolutional kernels can also be made to work with Muon. Take Conv2D as an example: the convolution kernel shape is typically $(w, h, d_{in}, d_{out})$, and its equivalent implementation actually flattens the $(w, h, d_{in})$ patch input into a $w \times h \times d_{in}$ vector, then reshapes the kernel into $(w\times h \times d_{in}, d_{out})$ as well, before performing matrix multiplication. So to use Muon on it, you'd first need to reshape the momentum into $(w\times h \times d_{in}, d_{out})$, compute $\msign$, and then reshape it back before applying the update.
Similarly, there's the gamma parameter of RMSNorm, which can be regarded as multiplication with a diagonal matrix; if we treat its momentum as a diagonal matrix, we can likewise compute $\msign$, and the result turns out to be equivalent to SignSGDM. The embedding layer can be treated as a collection of $(1,d)$ matrices, on which we compute $\msign$, giving Normalized SGDM (see Appreciating the Muon Optimizer: The Essential Leap from Vectors to Matrices). And if you want to keep tinkering — say, with multi-head attention — couldn't you consider pulling out each head's projection matrix and running $\msign$ on it separately...
Life goes on, and so does the tinkering~
What Results Should You Expect
Finally, if you've set everything up correctly according to the instructions above and gotten it running, all that's left is to pray for the favor of Lady Luck.
What kind of outcome should we expect? Barring anomalies like gradient explosions, in most cases Muon will do somewhat better than Adam — though it's not impossible for Muon to be slightly worse in some situations. Either way, the gap between them shouldn't be very large. If one turns out to be dramatically better than the other, that's probably a sign that something is misconfigured on one side or the other.
That said, none of this is absolute. Under certain extreme settings, it is indeed possible to see Muon substantially outperform Adam, with Adam failing no matter how it's tuned. In any case, good luck. If you run into anything interesting, I'd love to hear about it and discuss.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.