Muon from First Principles.

Imagine that two people are trying to descend the same mountain. One asks, “In which direction is the slope steepest?” The other asks a more subtle question: “What counts as a step of the same size?”

The answers need not be the same.

That distinction is the route to Muon. It is often introduced as “SGD with momentum, followed by orthogonalization.” That is a useful description of its implementation, but it leaves the most interesting question unanswered: why would we want to orthogonalize an update in the first place?

We will build the answer rather than start with the recipe. We will begin with AdamW, ask what a preconditioner changes, choose a geometry for a matrix update, and derive the appearance of the SVD’s polar factor. Only then will Newton–Schulz and the actual Muon update feel inevitable.

You do not need to know Muon before reading. Some familiarity with gradients, momentum, and matrix multiplication will help; every additional mathematical object is introduced when we need it.

The question behind an update

Consider a weight matrix W and the gradient of the loss with respect to its entries, G. Close to the current weights, the loss behaves approximately like a plane. If we propose a small change ΔW, its first-order effect is:

A negative inner product means our proposed step reduces the linear approximation to the loss. But we cannot make the step arbitrarily large: the approximation eventually becomes unreliable. We need a rule for how large a step is allowed to be.

SGD uses the negative gradient. Momentum first averages information across steps, making the direction less dependent on any one noisy mini-batch. AdamW goes further. It maintains moving averages of gradients and squared gradients, and scales each coordinate using its own history (while applying weight decay separately).

That works remarkably well. Yet an important structure has been left mostly implicit. A linear layer’s weights are not merely a bag of unrelated numbers; together, they define a transformation from one vector space to another.

If the parameter is a matrix, should “a large update” be measured entry by entry—or by what the update does to vectors?

That is the question we will follow.

What preconditioning changes

A preconditioner changes how a gradient becomes a step. In its familiar linear form, if we flatten the parameters into a vector g, we might write:

The matrix P can amplify some directions and suppress others. Different methods construct it for different reasons. Newton-type methods involve curvature information; adaptive optimizers such as Adam use gradient statistics to rescale coordinates. These are not the same thing: Adam’s diagonal adaptation should not be mistaken for an exact inverse Hessian.

The immediate intuition is simple. Suppose a gradient points mostly along one coordinate. Scaling coordinates independently can change its direction and the relative sizes of its components. But doing so does not explicitly treat a weight matrix as a linear operator. It does not ask about its singular directions as a group.

Muon does not escape the general idea of changing gradient geometry. In fact, exact orthogonalization can itself be expressed as a nonlinear, matrix-dependent transformation of the update. What it does not require is estimating and inverting a Hessian. We can motivate its direction with a much simpler choice: a norm.

Choose a geometry

Let’s return to the local approximation. We want a change with the smallest possible inner product with G, subject to a fixed step budget:

The answer depends on the norm ‖·‖. There is no unique “steepest direction” until we specify what constitutes a unit-sized step.

For ordinary Euclidean geometry, use the Frobenius norm—the square root of the sum of squared matrix entries. The solution is the familiar normalized negative gradient:

Now try a different definition of size. A matrix update also acts on an input vector x, changing the linear layer’s output by ΔWx. The spectral norm measures the largest amplification over every possible unit-length input:

If ‖ΔW‖₂ ≤ ε, then ‖ΔW x‖₂ ≤ ε‖x‖₂ for every input x. This is a natural operator-level notion of update size: it controls the worst-case change in the layer’s linear output. It is not, by itself, a guarantee about the change in the entire network’s loss.

Notice how the problem has shifted. We still want to descend the local loss plane. But we no longer measure the step using the length of a flattened vector. We measure what the matrix can do.

So which update is steepest under this spectral budget?

Why UVᵀ appears

To answer that question, we need the singular value decomposition (SVD). For a real matrix G of rank r, write its compact SVD as:

Think of this as three operations. Vᵀ identifies the input directions, Σ tells us how strongly each direction is stretched, and U identifies the corresponding output directions. The columns of U and V are orthonormal. The singular values are nonnegative.

Here is the crucial observation: the singular values describe the magnitude along these directions, while the singular vectors describe the directions themselves.

For any update D with spectral norm at most one, the matrix trace inequality gives:

The quantity on the right, the sum of singular values, is the nuclear norm. It is dual to the spectral norm: this inequality tells us the largest inner product available inside the spectral unit ball.

Can we achieve the bound? Yes. Choose D = U_rV_rᵀ. Its nonzero singular values are all one, so its spectral norm is one. Its alignment with G gives exactly the sum of G’s singular values.

Taking the negative direction solves our original descent problem:

We did not assume that orthogonalization was desirable and then invent a justification. We derived UVᵀ by asking for the steepest local decrease under a spectral-norm step budget.

A question is worth answering here. Does this mean every singular direction deserves equal emphasis? No. It means the spectral constraint allows us to use the full permitted scale in several orthogonal directions at once, and the resulting update optimizes the linearized objective under that particular constraint. It does not establish that all directions are equally useful in a nonlinear network, or that this choice always outperforms AdamW.

The polar factor

We now know why UVᵀ appeared. But there is a second way of understanding the same matrix, and it will lead directly to an efficient algorithm.

The polar decomposition separates a matrix into a direction-carrying factor and a positive-semidefinite stretching factor. If G = UΣVᵀ, define:

In plain language, Q preserves the singular directions but removes the unequal nonzero singular-value magnitudes. This operation is often called orthogonalization or the matrix’s zeroth power. For a full-rank rectangular matrix it produces a semi-orthogonal factor: its columns or rows are orthonormal, depending on the shape.

That expression is a useful bridge back to preconditioning. We are effectively transforming G using a matrix built from GᵀG. But this is not an inverse-Hessian method: GᵀG here is formed from the current update matrix, not from the loss’s Hessian.

It also exposes a computational problem. Calculating an SVD or an inverse matrix square root at every optimizer step for every eligible weight matrix would be expensive. Can we move the singular values toward one without explicitly finding the singular vectors?

Yes. This is the job of Newton–Schulz.

Newton–Schulz

We want a sequence of matrix multiplications that leaves singular vectors in place and changes only singular values.

Start with a normalized matrix X₀, obtained by dividing our update matrix by its Frobenius norm. This guarantees that all its singular values are at most one, up to numerical effects. For intuition, consider the classical cubic iteration:

Why should this work? Suppose X_k = UΣ_kVᵀ. Because U and V have orthonormal columns, multiplication by X_kX_kᵀ simply inserts another factor of Σ²:

For 0 < σ ≤ 1, that scalar mapping moves σ toward one. Singular vectors stay the same. Zero singular values remain zero in exact arithmetic. In other words, the matrix iteration can approach the polar factor without computing an SVD.

That is the central trick. Matrix products, unlike repeated explicit factorizations, are particularly well suited to modern accelerators.

But practical Muon takes a deliberate shortcut. It generally uses a more aggressive quintic polynomial for a small, fixed number of iterations:

The corresponding scalar map is f(σ) = aσ + bσ³ + cσ⁵. Near zero, it has a slope of 3.4445, much steeper than the classical cubic’s slope of 1.5. Small singular values can rise rapidly in just a few iterations.

Here is the trade-off: these practical coefficients are not designed to converge exactly to one if iterated indefinitely. In fact, f(1) = 0.701, so one is not even a fixed point of this polynomial. Five steps aim to give a useful, roughly balanced spectrum, not a mathematically exact polar factor.

There is one more subtlety. A zero singular value remains zero under any of these polynomials in exact arithmetic. Orthogonalization cannot manufacture missing rank from nothing. And for rectangular matrices, the factor is generally semi-orthogonal, not a square rotation matrix.

Assemble the optimizer

We finally have the mathematics needed to read the original Muon idea without taking any component on faith.

Let G_t be the gradient of an eligible weight matrix, and let B_t be its momentum buffer:

Call the matrix we want to transform M_t: either B_t or the Nesterov direction N_t. Normalize it and apply a few quintic Newton–Schulz steps. That yields NS(M_t), an approximately orthogonalized matrix direction.

Finally, scale the update and apply the learning rate. The scaled Muon recipe used in my Kimi K3 implementation follows the Moonlight shape convention:

Here λ is the decoupled weight-decay coefficient. The shape factor is a practical scaling convention, not a consequence of the spectral steepest-descent proof. Different Muon implementations and recipes need not use exactly the same scale.

Here is the numerical core, simplified from the working PyTorch implementation. The code assumes the input contains one matrix and omits device, precision, compilation, batching, and optimizer-state details:

def approximate_polar(matrix, steps=5, eps=1e-7):
		x = matrix / (matrix.norm() + eps)
		a, b, c = 3.4445, -4.7750, 2.0315

		for _ in range(steps):
				gram = x @ x.T
				x = a * x + (b * gram + c * gram @ gram) @ x

		return x

The code should now look familiar rather than mysterious. gram implements the polynomial without finding singular vectors. The coefficients control how singular values move. The number of iterations trades approximation quality against GPU cost.

In a real model, we also need to decide what counts as one matrix. This was surprisingly important in my Kimi K3 from Scratch training setup. The fused gate/up FFN weight, for example, is stored as one tensor but represents two independent projections. Similarly, packed attention projections can contain separate heads or key/value pieces. Orthogonalizing the storage tensor as one large matrix would mix those independent update geometries.

I therefore use logical parameter views so that each intended matrix receives its own Newton–Schulz transformation, while its underlying parameter and optimizer state remain shared with their original storage. I also implement compiled and batched matrix iterations to reduce Python and kernel overhead; those are implementation optimizations, not new mathematical claims about Muon.

And why keep AdamW? Because vectors, scalars, embeddings, output classifiers, and depthwise filters do not necessarily fit the same hidden-matrix update geometry. Muon is not intended to replace every optimizer group indiscriminately.

Where the idea stops

There are three conclusions we cannot draw from this derivation.

First, spectral-norm steepest descent is optimal only for the first-order local objective under its specified norm constraint. It does not prove that this norm is always the best geometry for a particular neural network.

Second, practical Muon does not apply the exact spectral steepest direction at every step. Momentum changes the matrix before orthogonalization, a finite Newton–Schulz polynomial only approximates its polar factor, and learning-rate scaling, decay, and parameter grouping affect the applied update.

Third, even a mathematically appealing direction does not establish better convergence or compute efficiency against a well-tuned AdamW baseline. Those are experimental questions, affected by architecture, batch size, training schedule, kernel cost, and hardware.

This distinction matters to me because the Kimi K3 project demonstrates an implemented training recipe and verified numerical behavior; without a controlled optimizer ablation, it would not be honest to claim that Muon was solely responsible for the final training results.

The whole journey

We began with an innocent question: what is a unit-sized step? From there, the entire optimizer fell into place.

A coordinate-wise adaptive method such as AdamW uses gradient history to change the scaling of individual parameters. A spectral view instead treats an eligible weight update as a linear operator. Under a spectral-norm constraint, the steepest direction of the linearized loss is the negative polar factor of the gradient: −UVᵀ.

An SVD tells us what that factor is. A Newton–Schulz polynomial lets us approximate it using matrix multiplication, without computing the SVD. Momentum decides which matrix to transform. Scaling, parameter groups, and numerical kernels turn the idea into a workable optimizer.

That is Muon from first principles: not magic, not simply “make matrices orthogonal,” and not an unconditional replacement for AdamW. It is a precise choice about which geometry should shape a matrix update, implemented approximately enough to be practical on GPUs.