Muon：当 LLM 优化器把权重矩阵当作矩阵处理

Hello, I'm Shrijith Venkatramana, and I'm building LiveReview — a blast-radius aware AI code review built for your business-critical systems. Star us to help devs discover the project, give it a try, and share your feedback to help improve the product. Most LLM developers know the AdamW update by heart: take the gradient, keep moving averages, normalize the update, change the weights. But there is a question hiding underneath all of this: What exactly is a 4096 x 4096 weight matrix? AdamW mostly treats it as 16 million scalar coordinates. Muon treats it as a matrix. That distinction is the whole story. Muon is an optimizer for the hidden matrix parameters of neural networks. Its core operation takes the momentum update, looks at its singular directions, throws away the singular-value magnitudes, and keeps the directions. It then approximates this operation cheaply with a few Newton-Schulz iterations. The result is an optimizer that has produced faster training in small-model competitions, scaled to multi-billion-parameter language models, and is now part of the mainstream PyTorch optimization stack. The interesting part is not merely that "Muon beats AdamW." The interesting part is why someone thought to optimize a neural network matrix as a matrix in the first place. 1. The story starts with a matrix Consider a transformer linear layer: y = W x Enter fullscreen mode Exit fullscreen mode where W might be a 4096 x 4096 matrix. During backpropagation we obtain a gradient: G = dL/dW Enter fullscreen mode Exit fullscreen mode AdamW maintains statistics for each scalar element of G. Very roughly: m_t = beta1 * m_(t-1) + (1-beta1) * G_t v_t = beta2 * v_(t-1) + (1-beta2) * G_t^2 update = m_t / sqrt(v_t) Enter fullscreen mode Exit fullscreen mode The important detail is G_t^2: the second-moment estimate is elementwise. This gives Adam a coordinate-wise view of the parameter space. Muon starts from a different observation: W is not merely a bag of numbers. It represents a linear transformation. For a matrix, one of the natural ways to understand that transformation is through its singular value decomposition: G = U Sigma V^T Enter fullscreen mode Exit fullscreen mode Here: U and V = directions Sigma = magnitudes along those directions Enter fullscreen mode Exit fullscreen mode So there are two different questions: Which directions does the gradient want to move? How large should each of those directions be? Muon makes a very particular choice: preserve the directions, but approximately equalize their magnitudes. That is what orthogonalizing the update means. The idea appeared publicly in October 2024, when Keller Jordan and collaborators were working on the NanoGPT speedrunning competition. On October 15, 2024, a Muon-based run set a new training-speed record, improving the previous result by about 35%. The project then became a collaboration involving people such as Jeremy Bernstein, Laker Newhouse, Yuchen Jin, Vlado Boza, Jiacheng You, and Franz Cesista. Jordan's account is unusually concrete about the engineering: Boza found that treating Q, K, and V separately worked better; Jin pushed experiments to larger models and supplied much of the H100 compute; Bernstein, You, and Cesista reduced the cost of the matrix orthogonalization itself. That history matters because Muon did not emerge from a giant benchmark suite first. It emerged from people trying to make a tiny training program go faster. 2. The key idea: keep the directions, flatten the singular values Suppose an update matrix has the form G = U diag(20, 3, 0.2) V^T Enter fullscreen mode Exit fullscreen mode The gradient has three principal matrix directions. One direction has magnitude 20. Another has magnitude 3. Another has magnitude 0.2. Muon constructs an approximation of U diag(1, 1, 1) V^T Enter fullscreen mode Exit fullscreen mode up to the appropriate scaling for the matrix shape. The singular vectors remain. The singular values disappear. For a square matrix, this produces an ordinary orthogonal matrix. For a rectangular matrix, it produces a semi-orthogonal matrix. Another way to write the operation is: Ortho(G) ~= U V^T Enter fullscreen mode Exit fullscreen mode This has a useful geometric interpretation. Imagine that your update is saying: "Move strongly in direction A, somewhat in direction B, and barely at all in direction C." Enter fullscreen mode Exit fullscreen mode Muon says: "Those are the important directions. Let's give them roughly equal opportunity to affect the layer." Enter fullscreen mode Exit fullscreen mode The original Muon write-up notes that transformer updates often have high condition numbers: a few directions can dominate the update while other directions have much smaller singular values. The authors proposed that orthogonalization may help those lower-magnitude directions contribute more. That explanation is an empirical hypothesis rather than the complete theoretical justification for Muon. This also explains why Muon is fundamentally different from simply changing Adam's hyperparameters. Adam changes how each coordinate is scaled. Muon changes the matrix geometry of the update. There is an analogy to dimensionality reduction, but in reverse. PCA asks: Which directions contain most of the variation? Enter fullscreen mode Exit fullscreen mode Muon asks: What if the update contains a few dominant directions, but I want the matrix update to retain all of its principal directions at comparable scale? Enter fullscreen mode Exit fullscreen mode 3. But an SVD every training step would be ridiculous There is an obvious problem. If we literally want: G = U Sigma V^T G -> U V^T Enter fullscreen mode Exit fullscreen mode we could compute an SVD. For a huge transformer, doing a full SVD for every large weight matrix at every optimizer step would be an unattractive idea. Muon's practical insight is: we do not need to compute the SVD explicitly. Instead, use Newton-Schulz iteration. Start by normalizing the matrix: X = G / ||G||_F Enter fullscreen mode Exit fullscreen mode Then repeatedly apply a matrix polynomial. The production Muon implementation uses: X_next = a X + b (X X^T) X + c (X X^T)^2 X Enter fullscreen mode Exit fullscreen mode with coefficients approximately: a = 3.4445 b = -4.7750 c = 2.0315 Enter fullscreen mode Exit fullscreen mode and typically five iterations. Why does this work? Take the SVD: X = U Sigma V^T Enter fullscreen mode Exit fullscreen mode The polynomial operation preserves the singular vectors: p(X) = U p(Sigma) V^T Enter fullscreen mode Exit fullscreen mode So instead of manipulating the whole matrix conceptually, we can think about what the polynomial does to each singular value. The quintic mapping is: p(s) = a s + b s^3 + c s^5 Enter fullscreen mode Exit fullscreen mode Repeatedly apply it. The goal is for the singular values to converge toward 1. So: U Sigma V^T | v U p(Sigma) V^T | v U p(p(Sigma)) V^T | v U I V^T Enter fullscreen mode Exit fullscreen mode The fascinating implementation detail is that we never explicitly calculate U, Sigma, or V. We only perform matrix multiplications. This is where numerical linear algebra meets GPU engineering. The early Muon work considered several ways of doing the orthogonalization. SVD was too slow. Other Newton-style methods had numerical problems in lower precision. Newton-Schulz could be run efficiently in bfloat16, which made it much more suitable for modern accelerators. The coefficients themselves were tuned experimentally; Jordan describes researchers using Desmos to explore polynomial shapes during the NanoGPT speedrun. That is a useful general lesson for ML engineers: an algorithm that is mathematically expensive may become practical when you find a formulation that maps onto the hardware's favorite operations. 4. The deeper idea: Muon is really about choosing a geometry There is a deeper way to understand all of this. Suppose a linear layer is: y = W x Enter fullscreen mode Exit fullscreen mode and we change the weights by: W -> W + dW Enter fullscreen mode Exit fullscreen mode The resulting change in the output is: dy = dW x Enter fullscreen mode Exit fullscreen mode So we can ask: How large should a weight update be if I care about controlling the change it causes to the layer's output? Suppose we measure vectors using RMS: ||x||_RMS = sqrt((1/d) sum_i x_i^2) Enter fullscreen mode Exit fullscreen mode Then a matrix has an operator norm describing its maximum RMS-to-RMS amplification: ||W||_(RMS->RMS) Enter fullscreen mode Exit fullscreen mode Now imagine the optimization problem: minimize subject to ||dW||_(RMS->RMS) Muon embeddings -> AdamW normalization -> AdamW biases -> AdamW LM head -> AdamW Enter fullscreen mode Exit fullscreen mode With current PyTorch, a schematic setup looks like this: muon_opt = torch.optim.Muon( muon_params, lr=3e-4, weight_decay=0.01, momentum=0.95, nesterov=True, adjust_lr_fn="match_rms_adamw", ) adamw_opt = torch.optim.AdamW( adamw_params, lr=3e-4, weight_decay=0.01, ) Enter fullscreen mode Exit fullscreen mode The important part is not the exact numbers above. The important part is constructing muon_params deliberately rather than doing: [p for p in model.parameters() if p.ndim == 2] Enter fullscreen mode Exit fullscreen mode because a tensor being 2D does not automatically mean that its optimization geometry should be Muon's. Also, do not blindly copy early Muon examples that use a fixed lr=0.02. There have been several generations of scaling conventions. Keller Jordan's original implementation, Moonshot's RMS-matching variant, and Bernstein's theoretical scaling rule use different parameter-shape-dependent factors. Current PyTorch exposes all three approaches. For a developer evaluating Muon, I would therefore treat the optimizer as part of the training configuration, not as a one-line AdamW replacement. A reasonable experiment is: same architecture same data same tokens same batch size same hardware same evaluation checkpoints compare: AdamW Muon + AdamW hybrid Enter fullscreen mode Exit fullscreen mode Then measure: validation loss vs tokens validation loss vs FLOPs validation loss vs GPU-hours peak memory step time Enter fullscreen mode Exit fullscreen mode The last three matter because an optimizer can improve sample efficiency while making each step slower, or reduce FLOPs without reducing wall-clock time on a communication-bound cluster. 7. The larger lesson Muon is interesting because it exposes a broader idea that is easy to miss when working with large neural networks: the optimizer contains assumptions about what a parameter means. Adam implicitly says: parameters are coordinates gradients are coordinates normalize coordinates independently Enter fullscreen mode Exit fullscreen mode Muon says: these parameters form matrices matrices represent linear operators linear operators have meaningful singular directions optimize those operators using a matrix-aware geometry Enter fullscreen mode Exit fullscreen mode That shift is bigger than the particular Newton-Schulz polynomial. The Newton-Schulz iteration is the implementation technique. The deeper idea is choosing a geometry that matches the structure of the object being optimized. That may also explain why optimizer research sometimes looks disconnected from ordinary software engineering. An optimizer sounds like a small implementation detail: optimizer.step() Enter fullscreen mode Exit fullscreen mode Yet changing that one line changes the effective geometry of a trillion-parameter computation. Muon began as an October 2024 experiment in a community speedrun, accumulated contributions from researchers and engineers working on the math and GPU implementation, scaled into the Moonlight training run, and subsequently became available in major training infrastructure. The interesting question now is less "Is Muon the replacement for AdamW?" and more: How many other parts of deep-learning systems are still being optimized using mathematical abstractions chosen for convenience rather than for the structure of the object itself? What other neural-network components do you think deserve their own optimization geometry? Your team's attention is limited, and the deluge of AI-generated code is making it harder to keep production reliable and secure without slowing you down. I'm building LiveReview, a blast-radius aware AI code review built for your business-critical systems. Instead of presenting every diff with equal emphasis, LiveReview scores each change by blast radius — how far its impact reaches through your call graph — so you can focus attention where it actually matters. Spend code review effort where business risk is highest — not spread evenly across every diff. ⭐ Star it on GitHub: HexmosTech / LiveReview Blast-Radius Aware AI Code Review for Business-Critical Systems LiveReview: Blast-Radius Aware AI Code Review for Business-Critical Systems LiveReview is an AI code reviewer that scores every hunk of a diff by blast radius: how far a change reaches through your call graph, how much persistent state it touches, and how well-tested it is. A 3-line change to a shared auth check can outrank a 300-line UI tweak. Your team's attention goes to the highest-risk code first, not spread evenly across every diff. blast-radius-demo.mp4 LiveReview's Blast Radius & Review Priority scoring, live in the diff viewer. The exact math, not a black box Visualize blast radius at a glance Every factor that feeds the score How does Blast Radius scoring work? (a more technical explanation) Here's the goal: A 3-line fix in a function used by 40 other files, that also writes to a database, should score high. A 300-line UI change in one file, fully covered by… View on GitHub Click below to try LiveReview with your codebase: