{"slug": "muon-what-happens-when-an-llm-optimizer-treats-a-weight-matrix-like-a-matrix", "title": "Muon: What Happens When an LLM Optimizer Treats a Weight Matrix Like a Matrix", "summary": "Developer Shrijith Venkatramana explains Muon, an optimizer that treats neural network weight matrices as matrices rather than collections of scalar coordinates, orthogonalizing momentum updates by preserving singular directions while discarding singular-value magnitudes. The technique, introduced publicly in October 2024 by Keller Jordan and collaborators during the NanoGPT speedrunning competition, set a training-speed record about 35% faster than the previous result and has since scaled to multi-billion-parameter language models and entered the mainstream PyTorch optimization stack.", "body_md": "*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](https://github.com/HexmosTech/LiveReview/) to help devs discover the project, give it a try, and share your feedback to help improve the product.*\n\nMost LLM developers know the AdamW update by heart:\n\ntake the gradient, keep moving averages, normalize the update, change the weights.\n\nBut there is a question hiding underneath all of this:\n\n**What exactly is a 4096 x 4096 weight matrix?**\n\nAdamW mostly treats it as 16 million scalar coordinates.\n\nMuon treats it as a matrix.\n\nThat distinction is the whole story.\n\nMuon 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.\n\nThe 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.\n\nThe interesting part is not merely that \"Muon beats AdamW.\"\n\nThe interesting part is **why someone thought to optimize a neural network matrix as a matrix in the first place.**\n\nConsider a transformer linear layer:\n\n```\ny = W x\n```\n\nwhere `W` might be a 4096 x 4096 matrix.\n\nDuring backpropagation we obtain a gradient:\n\n```\nG = dL/dW\n```\n\nAdamW maintains statistics for each scalar element of `G`.\n\nVery roughly:\n\n```\nm_t = beta1 * m_(t-1) + (1-beta1) * G_t\n\nv_t = beta2 * v_(t-1) + (1-beta2) * G_t^2\n\nupdate = m_t / sqrt(v_t)\n```\n\nThe important detail is `G_t^2`: the second-moment estimate is elementwise.\n\nThis gives Adam a coordinate-wise view of the parameter space.\n\nMuon starts from a different observation:\n\n`W` is not merely a bag of numbers. It represents a linear transformation.\n\nFor a matrix, one of the natural ways to understand that transformation is through its singular value decomposition:\n\n```\nG = U Sigma V^T\n```\n\nHere:\n\n```\nU and V = directions\nSigma  = magnitudes along those directions\n```\n\nSo there are two different questions:\n\nMuon makes a very particular choice:\n\n**preserve the directions, but approximately equalize their magnitudes.**\n\nThat is what orthogonalizing the update means.\n\nThe 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.\n\nThat history matters because Muon did not emerge from a giant benchmark suite first.\n\nIt emerged from people trying to make a tiny training program go faster.\n\nSuppose an update matrix has the form\n\n```\nG = U diag(20, 3, 0.2) V^T\n```\n\nThe gradient has three principal matrix directions.\n\nOne direction has magnitude 20.\n\nAnother has magnitude 3.\n\nAnother has magnitude 0.2.\n\nMuon constructs an approximation of\n\n```\nU diag(1, 1, 1) V^T\n```\n\nup to the appropriate scaling for the matrix shape.\n\nThe singular vectors remain.\n\nThe singular values disappear.\n\nFor a square matrix, this produces an ordinary orthogonal matrix. For a rectangular matrix, it produces a semi-orthogonal matrix.\n\nAnother way to write the operation is:\n\n```\nOrtho(G) ~= U V^T\n```\n\nThis has a useful geometric interpretation.\n\nImagine that your update is saying:\n\n```\n\"Move strongly in direction A,\n somewhat in direction B,\n and barely at all in direction C.\"\n```\n\nMuon says:\n\n```\n\"Those are the important directions.\nLet's give them roughly equal opportunity to affect the layer.\"\n```\n\nThe 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.\n\nThis also explains why Muon is fundamentally different from simply changing Adam's hyperparameters.\n\nAdam changes how each coordinate is scaled.\n\nMuon changes the **matrix geometry of the update**.\n\nThere is an analogy to dimensionality reduction, but in reverse.\n\nPCA asks:\n\n```\nWhich directions contain most of the variation?\n```\n\nMuon asks:\n\n```\nWhat if the update contains a few dominant directions,\nbut I want the matrix update to retain all of its principal directions\nat comparable scale?\n```\n\nThere is an obvious problem.\n\nIf we literally want:\n\n``` php\nG = U Sigma V^T\n\nG -> U V^T\n```\n\nwe could compute an SVD.\n\nFor a huge transformer, doing a full SVD for every large weight matrix at every optimizer step would be an unattractive idea.\n\nMuon's practical insight is:\n\n**we do not need to compute the SVD explicitly.**\n\nInstead, use Newton-Schulz iteration.\n\nStart by normalizing the matrix:\n\n```\nX = G / ||G||_F\n```\n\nThen repeatedly apply a matrix polynomial.\n\nThe production Muon implementation uses:\n\n```\nX_next = a X\n          + b (X X^T) X\n          + c (X X^T)^2 X\n```\n\nwith coefficients approximately:\n\n```\na =  3.4445\nb = -4.7750\nc =  2.0315\n```\n\nand typically five iterations.\n\nWhy does this work?\n\nTake the SVD:\n\n```\nX = U Sigma V^T\n```\n\nThe polynomial operation preserves the singular vectors:\n\n```\np(X) = U p(Sigma) V^T\n```\n\nSo instead of manipulating the whole matrix conceptually, we can think about what the polynomial does to each singular value.\n\nThe quintic mapping is:\n\n```\np(s) = a s + b s^3 + c s^5\n```\n\nRepeatedly apply it.\n\nThe goal is for the singular values to converge toward 1.\n\nSo:\n\n```\nU Sigma V^T\n        |\n        v\nU p(Sigma) V^T\n        |\n        v\nU p(p(Sigma)) V^T\n        |\n        v\nU I V^T\n```\n\nThe fascinating implementation detail is that we never explicitly calculate `U`, `Sigma`, or `V`.\n\nWe only perform matrix multiplications.\n\nThis is where numerical linear algebra meets GPU engineering.\n\nThe 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.\n\nThat is a useful general lesson for ML engineers:\n\n**an algorithm that is mathematically expensive may become practical when you find a formulation that maps onto the hardware's favorite operations.**\n\nThere is a deeper way to understand all of this.\n\nSuppose a linear layer is:\n\n```\ny = W x\n```\n\nand we change the weights by:\n\n``` php\nW -> W + dW\n```\n\nThe resulting change in the output is:\n\n```\ndy = dW x\n```\n\nSo we can ask:\n\nHow large should a weight update be if I care about controlling the change it causes to the layer's output?\n\nSuppose we measure vectors using RMS:\n\n```\n||x||_RMS = sqrt((1/d) sum_i x_i^2)\n```\n\nThen a matrix has an operator norm describing its maximum RMS-to-RMS amplification:\n\n``` php\n||W||_(RMS->RMS)\n```\n\nNow imagine the optimization problem:\n\n``` php\nminimize      <G, dW>\n\nsubject to    ||dW||_(RMS->RMS) <= eta\n```\n\nIn plain English:\n\n```\nChoose the update that gives the largest first-order\ndecrease in loss, while limiting how much the layer\ncan change its outputs.\n```\n\nThe solution involves the orthogonalized gradient:\n\n```\ndW ~= -eta * scale * U V^T\n```\n\nSo the `U V^T` operation is not merely an arbitrary trick.\n\nIt arises from asking what \"the biggest useful update\" means under a matrix norm that is tied to the behavior of a linear layer.\n\nThis is part of Jeremy Bernstein and Laker Newhouse's broader work on **modular duality**. Their 2025 ICML paper develops a framework in which different neural-network modules can be assigned different geometries, with GPU-friendly dualization procedures for layers such as Linear and Conv2D. Newton-Schulz appears naturally in that construction.\n\nThis perspective also connects Muon to Shampoo.\n\nWithout its accumulation mechanism, the Shampoo update can be algebraically reduced to an orthogonalized gradient:\n\n```\nG\n |\n v\nU Sigma V^T\n |\n v\nU V^T\n```\n\nSo Muon can be viewed as a particularly cheap, momentum-based way of getting this matrix-aware behavior.\n\nThat is one reason the optimizer is intellectually interesting.\n\nIt is less about inventing another collection of moving averages and more about asking:\n\n**What metric should a neural-network layer use for optimization?**\n\nThe early results were promising, but there was a serious problem:\n\nWould this thing actually scale?\n\nMoonshot AI addressed that question in the 2025 paper *Muon is Scalable for LLM Training*.\n\nThey identified two practical issues that mattered at larger scale:\n\n```\n1. Weight decay\n2. Correct scaling of the Muon update\n```\n\nThe second point is particularly important.\n\nMuon produces an orthogonalized matrix whose RMS behavior depends on the dimensions of the matrix.\n\nA 1024 x 1024 matrix and a 8192 x 8192 matrix cannot simply receive the identical raw update scale and be expected to behave identically.\n\nMoonshot introduced an update scaling rule designed to make Muon's update RMS comparable to AdamW's. Their experiments reported roughly 2x computational efficiency at compute-optimal training, with comparable performance reached using roughly 52% of the training FLOPs of the AdamW counterparts in their scaling experiments.\n\nThey also trained Moonlight, a 3B/16B mixture-of-experts model, on 5.7 trillion tokens using Muon.\n\nThis is where the distinction between \"interesting optimizer paper\" and \"useful engineering technique\" becomes important.\n\nA 2x efficiency result is economically meaningful only when the comparison is properly controlled.\n\nFor example, if a training run costs:\n\n```\n$1,000,000\n```\n\nand the compute requirement genuinely falls by 48%, the idealized savings are:\n\n```\n$1,000,000 * 0.48 = $480,000\n```\n\nBut actual GPU spend is not a pure FLOP meter.\n\nYou also have:\n\n```\nGPU utilization\ncommunication\ncheckpointing\ndata loading\noptimizer implementation\nnetwork topology\nengineering time\nfailed runs\n```\n\nSo \"48% fewer FLOPs\" should be read as an opportunity for lower cost, rather than as a promise of a 48% lower cloud bill.\n\nThere is also a useful memory difference.\n\nAdam-like optimizers commonly maintain two moment tensors:\n\n```\nm\nv\n```\n\nMuon's core optimizer state contains one momentum buffer.\n\nIgnoring parameter replicas, master weights, sharding, and datatype choices:\n\n```\nAdamW optimizer state: 2 x parameter bytes\nMuon optimizer state:  1 x parameter bytes\n```\n\nAt 100B parameters, if those state tensors were stored in fp32:\n\n```\n100B * 4 bytes = 400 GB\n\nAdamW moments:\n2 * 400 GB = 800 GB\n\nMuon momentum:\n1 * 400 GB = 400 GB\n```\n\nThat difference becomes relevant when optimizer state is one of the constraints determining how many GPUs a training job needs.\n\nThe computational cost of Newton-Schulz is also less frightening than the name suggests.\n\nFor an n x m matrix, with `m <= n`, the Muon write-up derives an extra cost of roughly:\n\n```\n6 T n m^2\n```\n\nFLOPs for `T` Newton-Schulz steps.\n\nThe corresponding forward-plus-backward cost for the linear layer scales roughly like:\n\n```\n6 n m B\n```\n\nwhere `B` is the number of tokens processed by the layer in the batch.\n\nThe ratio is therefore approximately:\n\n```\noverhead ~= T m / B\n```\n\nTake a hypothetical training setup:\n\n```\nmodel width m = 4096\ntokens per batch B = 4,000,000\nNewton-Schulz steps T = 5\n```\n\nThen:\n\n```\noverhead ~= 5 * 4096 / 4,000,000\n         ~= 0.00512\n         ~= 0.51%\n```\n\nThat is the operational trick.\n\nThe expensive-looking matrix computation is amortized across millions of tokens.\n\nFor the actual NanoGPT speedrun configuration discussed by Jordan, the corresponding estimate was about 0.7%.\n\nAs of 2026, the idea has also moved into mainstream infrastructure. Current PyTorch documentation exposes `torch.optim.Muon`, including different update-scaling modes, and the DeepSpeed team added Muon support in June 2026.\n\nThe first mistake would be:\n\n```\n\"Replace AdamW everywhere with Muon.\"\n```\n\nThat is not how the original method is intended to be used.\n\nMuon is primarily for 2D hidden weight matrices.\n\nParameters such as:\n\n```\nembeddings\nbiases\nLayerNorm parameters\nother 1D parameters\ninput layers\noutput heads\n```\n\nremain on a conventional optimizer such as AdamW.\n\nThe original experiments also found that Q, K, and V were better treated as separate matrices rather than as one fused QKV matrix.\n\nConceptually, your optimizer setup looks like:\n\n``` php\nhidden Linear weights  -> Muon\nembeddings              -> AdamW\nnormalization           -> AdamW\nbiases                  -> AdamW\nLM head                 -> AdamW\n```\n\nWith current PyTorch, a schematic setup looks like this:\n\n```\nmuon_opt = torch.optim.Muon(\n    muon_params,\n    lr=3e-4,\n    weight_decay=0.01,\n    momentum=0.95,\n    nesterov=True,\n    adjust_lr_fn=\"match_rms_adamw\",\n)\n\nadamw_opt = torch.optim.AdamW(\n    adamw_params,\n    lr=3e-4,\n    weight_decay=0.01,\n)\n```\n\nThe important part is not the exact numbers above.\n\nThe important part is constructing `muon_params` deliberately rather than doing:\n\n```\n[p for p in model.parameters() if p.ndim == 2]\n```\n\nbecause a tensor being 2D does not automatically mean that its optimization geometry should be Muon's.\n\nAlso, do not blindly copy early Muon examples that use a fixed `lr=0.02`.\n\nThere 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.\n\nFor a developer evaluating Muon, I would therefore treat the optimizer as part of the **training configuration**, not as a one-line AdamW replacement.\n\nA reasonable experiment is:\n\n```\nsame architecture\nsame data\nsame tokens\nsame batch size\nsame hardware\nsame evaluation checkpoints\n\ncompare:\n    AdamW\n    Muon + AdamW hybrid\n```\n\nThen measure:\n\n```\nvalidation loss vs tokens\nvalidation loss vs FLOPs\nvalidation loss vs GPU-hours\npeak memory\nstep time\n```\n\nThe 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.\n\nMuon is interesting because it exposes a broader idea that is easy to miss when working with large neural networks:\n\n**the optimizer contains assumptions about what a parameter means.**\n\nAdam implicitly says:\n\n```\nparameters are coordinates\ngradients are coordinates\nnormalize coordinates independently\nthese parameters form matrices\nmatrices represent linear operators\nlinear operators have meaningful singular directions\noptimize those operators using a matrix-aware geometry\n```\n\nThat shift is bigger than the particular Newton-Schulz polynomial.\n\nThe Newton-Schulz iteration is the implementation technique.\n\nThe deeper idea is choosing a geometry that matches the structure of the object being optimized.\n\nThat may also explain why optimizer research sometimes looks disconnected from ordinary software engineering. An optimizer sounds like a small implementation detail:\n\n```\noptimizer.step()\n```\n\nYet changing that one line changes the effective geometry of a trillion-parameter computation.\n\nMuon 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:\n\n**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?**\n\nWhat other neural-network components do you think deserve their own optimization geometry?\n\nYour 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.\n\nI'm building **LiveReview**, a blast-radius aware AI code review built for your business-critical systems.\n\nInstead 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.**\n\nSpend code review effort where business risk is highest — not spread evenly across every diff.\n\n⭐ Star it on GitHub: \n\nLiveReview 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.\n\n*LiveReview's Blast Radius & Review Priority scoring, live in the diff viewer.*\n\n| The exact math, not a black box | Visualize blast radius at a glance | Every factor that feeds the score | \n|---|---|---|\n\n**Here's the goal:**\n\n**Click below to try LiveReview with your codebase:**", "url": "https://wpnews.pro/news/muon-what-happens-when-an-llm-optimizer-treats-a-weight-matrix-like-a-matrix", "canonical_source": "https://dev.to/shrsv/muon-what-happens-when-an-llm-optimizer-treats-a-weight-matrix-like-a-matrix-35n6", "published_at": "2026-09-23 19:50:25+00:00", "updated_at": "2026-09-23 20:29:04.854895+00:00", "lang": "en", "topics": ["machine-learning", "large-language-models", "ai-research", "ai-infrastructure", "developer-tools"], "entities": ["Muon", "AdamW", "Shrijith Venkatramana", "LiveReview", "Keller Jordan", "Jeremy Bernstein", "PyTorch", "NanoGPT"], "alternates": {"html": "https://wpnews.pro/news/muon-what-happens-when-an-llm-optimizer-treats-a-weight-matrix-like-a-matrix", "markdown": "https://wpnews.pro/news/muon-what-happens-when-an-llm-optimizer-treats-a-weight-matrix-like-a-matrix.md", "text": "https://wpnews.pro/news/muon-what-happens-when-an-llm-optimizer-treats-a-weight-matrix-like-a-matrix.txt", "jsonld": "https://wpnews.pro/news/muon-what-happens-when-an-llm-optimizer-treats-a-weight-matrix-like-a-matrix.jsonld"}}