How Muon Lost Its Geometry

Arda Goreci · June 10, 2026

Muon is a new optimizer that has become quite popular recently. It is the first optimizer since Adam to be adopted by very large scale training runs (> 1T parameter LLMs). Its construction is theoretically motivated, and it's a breath of fresh air in the deep learning world where most choices are heuristic-based and empirical An iconic example is Noam Shazeer’s SwiGLU proposal, where he wrote that “we attribute its success, as all else, to divine benevolence.” .

In this blog, we will try to follow the popularization of Muon and give an account of how we think a scaling factor "bug" crept in to basically all implementations of Muon. This eventually made its way into very large scale training runs like Kimi-K2 and DeepSeek-V4, and we'll provide evidence for why it made ugly patches like MuonClip necessary. The theoretically correct way of doing this was already described by Jeremy Bernstein, one of the key authors of Muon and related ideas (who now works at Thinking Machines), but somehow the community seems to have overlooked this.

During the popularization of Muon, it seems to have accidentally lost is its maximal-update geometry Maximal update parameterization (µP) is a scaling rule for neural networks that keeps activations and update-induced activation changes comparable as width changes. In LLM training, it is used to tune hyperparameters at small scale and transfer them to larger models. : the scaling rule that makes an optimizer step produce the right-size perturbation in activation space. This was the reason we got interested in Muon in the first place. Neural networks used to model biomolecular space, such as AlphaFold3 architectures, have all sorts of layers: down-projections, up-projections, gates, pair updates, attention blocks, and geometry modules. It is very likely that the global learning rate in these models is an uneasy compromise between many layers that each want a different effective learning rate. If maximal-update parameterization were built into Muon correctly, we thought it might fix a lot of the quiet numerical issues in these networks.

Before we begin, we must give due credit to Jeremy Bernstein, whose blog post "Deriving Muon" is the only reason we were able to understand Muon's theoretical foundations well enough to notice something was off in the implementations. This also saved us from pretty nasty easter eggs in our codebase - silent numerical issues that would have sat in our model training stack which we probably never would have found.

A series of unfortunate events

Muon went from a training trick in CIFAR and NanoGPT speedruns to training trillion-parameter frontier models in about 18 months. That's quite impressive. This likely also means that people haven't quite dialed in how to work with it properly, iron out issues, make it efficient, etc.

We will describe the timeline that seems to have introduced the problematic scaling factor messing up the maximal update property.

One learning rate, many shapes

The problem we are going to discuss with the current implementations of Muon is about the maximal update parameterization (µP). This was an influential scaling recipe which allowed LLMs to be tuned at small scales, and then the hyperparameters such as the learning rate be transferred directly to the giant models with much larger widths. At the moment, none of the implementations of Muon offers the correct µP version, which could be problematic for researchers trying to do hyperparameter transfer with Muon or apply it to other domains such as protein co-folding. Before we talk about Muon itself, we need to discuss what kind of behavior we want from a neural network as it trains, which is related to the core idea of µP.

What we want during neural network training

At a high level, a neural network is a long sequence of vector states being multiplied by matrices, with nonlinearities sprinkled between the multiplications. The vector might represent a token, an amino-acid residue, an atom, or any other object the model is carrying through its layers. The first thing we want is simple: the individual coordinates of those vectors should stay normal-sized.

If a vector has width N, and its entries are roughly the same size, say Gaussian-ish around zero with order-one variance, then the vector magnitude is roughly √N. That is usually true at initialization, and we would like training to preserve that scale. We also want the change in activations after a training step to scale the same way: ||Δx|| should be roughly √N, so the individual coordinate updates are nonzero rather than vanishing as the model gets wider.

The point of using N is scaling behavior. We want to change the width of the model and keep the same qualitative training dynamics. But activation changes are not a knob we turn directly. They are caused by the weights. So if we want stable feature learning across widths, we need a condition on how the weight matrices are allowed to change during training. This is where Muon enters.

What Muon is trying to do

Muon is designed to limit output perturbations in a gradient update: it updates a linear layer in a loss-reducing direction, subject to a constraint on how much the layer’s outputs are allowed to change. This turns out to be quite useful.

Neural networks are compositional. They have many layers. Each layer is trying to learn a good mapping between its input distribution and its outputs. This is hard to do during training, because your input distribution is changing after every gradient update as lower layers learn.

Normalization tricks such as LayerNorm and RMSNorm try to mitigate this problem by mapping the activations back to a hypersphere of radius √n where n is the number of neurons in the layer. This makes sure layer inputs remain well-behaved numerically and they don’t become too small / big as the network grows. However, you can still be moving very fast on the surface of this hypersphere during training if your weight updates are large.

Muon has a couple tricks to bound the change in the Root Mean Square (RMS) RMS is the square root of the mean of squared entries: \(\sqrt{\frac{1}{n}\sum_i x_i^2}\). It measures the typical size of a vector entry. of a linear layer’s outputs. Combined with the right LR scaling based on fan-in and fan-out Fan-in and fan-out are the input and output dimensions of the matrix, respectively. , Muon can give tight bounds on the RMS change after an optimizer step. This fan-out / fan-in based scaling is not quite right in popular Muon implementations. You get in trouble if your matrix is rectangular. Moonshot AI got in exactly this kind of trouble when they combined multi-head latent attention (from DeepSeek V3) + wrong Muon scaling.

To see this problem, I find it the most intuitive to reason with RMS→RMS norms as done in Jeremy Bernstein’s blog post on deriving Muon.

Think of a simple linear layer that produces the output \(\mathbf{y}\), and let’s reason about what happens element-wise to the outputs of \(\mathbf{y}\) after an optimizer step:

The optimizer step changes the weight matrix with an update \(\Delta \mathbf{W}\). The resulting change in the network output is:

We do not want \(\Delta \mathbf{y}\) to be too big after a weight update. If certain layers make very large jumps after every training step, every layer after them receives noisy inputs, which makes learning harder. A nice way to reason about this change is the RMS→RMS operator norm of a matrix, which is the maximum amount by which a matrix can increase the average entry-size of its inputs after the matmul:

The spectral norm Spectral norm is a linear algebra quantity that measures the maximum factor by which a matrix can stretch its input vector. \(\|\cdot\|_{*}\) matters here because the Newton-Schulz orthogonalization step normalizes the update so that this quantity is 1.

The RMS→RMS norm of \(\Delta \mathbf{W}\) gives an upper bound on the change in output:

This means that as long as \(\|\Delta \mathbf{W}\|_{\mathrm{RMS}\to\mathrm{RMS}}\) is bounded, the change in the layer outputs remains bounded too.

Now suppose autodiff computes the raw gradient \(\nabla_{\mathbf{W}}\mathcal{L}\). Muon applies a few Newton-Schulz iterations to orthogonalize the gradient before the update step. This gives us the update direction that Muon will use. Because the routine includes a spectral-norm normalization, that update has:

If we calculate the RMS→RMS norm of this update, it comes out to:

This is the move: the geometric correction. Newton-Schulz gives us an update with spectral norm 1, but the quantity we care about is RMS→RMS. For a rectangular matrix, those are not the same. The update has RMS→RMS norm \(\sqrt{\frac{\text{fan-in}}{\text{fan-out}}}\), so we add the reciprocal factor \(\sqrt{\frac{\text{fan-out}}{\text{fan-in}}}\) before applying it:

This update rule satisfies \(\|\Delta \mathbf{W}\|_{\mathrm{RMS}\to\mathrm{RMS}} \le \eta\), and hence \(\|\Delta \mathbf{y}\|_{\mathrm{RMS}} \le \eta\). In words: the update keeps the RMS change in the outputs of any linear layer bounded by the learning rate \(\eta\) The core idea of µP with optimizers like Adam is also to make sure the singular values of weight matrices and updates scale like \(\sqrt{\frac{\text{fan-out}}{\text{fan-in}}}\). See Jeremy Bernstein’s note. .

This \(\sqrt{\frac{\text{fan-out}}{\text{fan-in}}}\) scaling is not being done right in Keller’s speedrun implementation or Moonshot AI’s training.

Where implementations diverge

When does the scaling matter?

The first large-scale stress test: Moonlight

Soon after its proposal, Muon was shown to scale to LLM training by Moonshot AI, an artificial intelligence company based in China. They wrote a paper literally called "Muon is Scalable for LLM Training", in which the Moonshot team progress the field of open-source LLM research by introducing Muon as their optimizer. However, the paper makes a couple of modifications to Muon based on a bunch of assumptions that are demonstrably wrong. This added a theoretical blunder to Muon which has now made it into frameworks like PyTorch and Optax. We dedicate the rest of the blog to explain the mistake and suggest a correction, paraphrasing Jeremy Bernstein's work.

The crucial modification is their per-matrix rescaling. The authors write, “In order to maintain consistent update RMS among matrices of different shapes, we propose to scale the Muon update for each matrix by its \(\sqrt{\max(A, B)}\).” At first glance, this sounds close to Keller’s \(\max(1, \sqrt{\frac{\text{fan-out}}{\text{fan-in}}})\), but the two are equivalent only if all matrices have the same second dimension. Moonlight uses the DeepSeek-V3 architecture, and that definitely does not have matrices with the same fan-in everywhere: Multi-head Latent Attention uses up- and down-projections to compress the KV cache during inference. So the assumption does not even make sense for the architecture they are using.

The authors then report that all is not well with Muon over Adam: attention becomes unstable. This should not be surprising. The dodgy scaling factor distorts the learning-rate geometry of MLA, and the symptom shows up as attention-logit explosions. By our calculation, the MLA down-projection is updated around 5x larger than it should be with respect to the global scale, which, when combined with the corresponding up-projection and the rest of the attention block, could easily destabilize MLA.

AdamW Muon

AdamW Muon

A CIFAR-10 learning-rate transfer test

To empirically confirm the failure of hyperparameter transfer with Moonshot AI's settings, we made a learning-rate transfer experiment using CIFAR-10. We constructed MLPs with a bottleneck in the middle of 6 layers to make sure we had rectangular matrices, and we trained models with 5 different optimizer settings: out-of-the-box Adam, µP learning-rate Adam, Muon with the correct \(\sqrt{\frac{\text{fan-out}}{\text{fan-in}}}\) scaling, Muon with \(\max(1, \sqrt{\frac{\text{fan-out}}{\text{fan-in}}})\) scaling, which gets the wrong learning rate in downprojections, and finally Muon with Moonshot AI’s proposed scaling \(\sqrt{\max(\text{fan-out}, \text{fan-in})}\). In a hyperparameter transfer experiment, the goal is come up with a recipe so that even when we scale up the size of the neural networks several fold, the optimal learning rate stays the same. If there is any drift, this would (1) require re-tuning of a larger network (2) if different layers prefer very different learning rates, the global learning rate is going to be an uneasy compromise between the layers.

The results are quite striking. Perhaps the most noticable jump is from Moonshot AI's scaling factor to the correct maximal update parameterization. The correct version produces the lowest test loss at roughly the same learning rate at every width, whereas Moonshot's version systematically shifts to lower learning rates as we increase the width. The scaling factor proposed by Keller Jordan fares much better than Moonshot's, but it still drifts a little bit compared to the µP version. The drift in production settings would depend on the number of down-projection matrices and their down-projection factor. It's also noticeable that Muon with the correct scaling produces the lowest absolute test loss of any setting.

When everything gets the wrong learning rate

The implications for not getting the maximal update parameterization exactly right matters most if your network architecture has plenty of rectangular matrices. This matters for more experimental LLM architectures like DeepSeek-V3, with its multi-head latent attention logic, and for custom neural networks for biomolecular structure prediction, like AlphaFold3. To make the comparison concrete, we took a bunch of layers from the AlphaFold3 architecture and computed what the maximal-update (µP) scaling would prefer for the learning rate to be with respect to the global learning rate, versus what it actually ends up being. Here is the percentage error of the popular scaling rules relative to the correct \(\sqrt{\frac{\text{fan-out}}{\text{fan-in}}}\) factor for some rectangular projections in an AlphaFold3-style architecture:

When the per-layer learning rates are wrong, the global optimum is going to be an uneasy compromise between all the layers of the network, and at lower precision training modes like bfloat16, could be potentially destabilizing for the network. We hypothesize that correct maximal update scaling with Muon could be a very good way of stably training very heterogeneous networks like biomolecular structure prediction models.

Conclusion

Muon is one of the most exciting algorithmic ideas proposed in deep learning. It is theoretically motivated, elegant, and beats strong baselines in speedruns. That almost never happens for optimizer ideas. It seems like as Muon moved from speedruns to large-model training runs to framework implementations, the community seems to have copied the most visible part of the method, the Newton-Schulz orthogonalization trick, while muddying the scaling rule that made the geometry meaningful. Square matrices hide this mistake, because most of the scaling rules agree there. Modern networks are rarely just square matrices; they are full of up-projections, down-projections, gates, heads, bottlenecks, and coordinate maps. Once those pieces start getting different effective learning rates, the global learning rate becomes an uneasy compromise: turn it down to save the hot layers and the cold layers stop moving; turn it up to wake up the cold layers and the hot layers destabilize. In that light, patches like MuonClip look less like new optimizer ideas and more like symptoms of a broken scaling rule. The good news is that the fix is conceptually simple: just restore the maximal-update geometry with the \(\sqrt{\frac{\text{fan-out}}{\text{fan-in}}}\) correction. This can be a one-line fix in libraries like PyTorch and Optax that implement Muon.

We are excited to see the impact of Muon and related ideas in the near future!

Arda Goreci, "How Muon Lost Its Geometry",
Ligo Biosciences Blog, June 10, 2026.

@article{goreci2026muonlostgeometry,
  author = {Arda Goreci},
  title = {How Muon Lost Its Geometry},
  journal = {Ligo Biosciences Blog},
  year = {2026},
  month = jun,
  day = {10},
  publisher = {Ligo Biosciences}
}