Kernel math / optimizer lab

When we trained one line, y = wx + b, the optimizer's job looked almost trivial: get a gradient, move the weight, repeat. Muon only becomes interesting when that one weight turns into a matrix. This is the missing bridge between my smallest backpropagation loop, matrix Wx+b and GEMM, and the earlier SGD-to-AdamW guide. The first Muon research note compares optimizers broadly; here I want to put numbers on the page and train two small models from scratch.

The central question is not “Can Muon replace AdamW everywhere?” It cannot. The useful question is: when the gradient of a hidden-layer weight is a matrix, why might we transform the whole update matrix instead of scaling each entry independently?

1. One line, one gradient, one update

Take x = 2, w = 0, b = 0, and target t = 4. The prediction is zero. With half-squared error, the loss is L = 0.5(0 - 4)^2 = 8. The chain rule gives dL/dw = (wx+b-t)x = -8 and dL/db = -4. At learning rate 0.1, one SGD step makes w = 0.8 and b = 0.4. The next prediction is 0.8*2+0.4 = 2. We moved toward the answer.

Nothing about that example needs Muon. A scalar has no matrix directions to balance. Read first, if usefulThe tiny training loopThe wx+b backprop post derives the gradient rather than asking you to memorize it.

2. Turn wx+b into a matrix multiplication

Now collect examples into rows of X and let a layer have several outputs. With Y = XW + b, an arriving output gradient D = dL/dY produces dW = X^T D. That backward GEMM is the point at which a matrix-valued optimizer has something to work with. The forward GEMM computes predictions; the backward GEMM accumulates evidence about every connection in W; the optimizer changes W for the next batch.

X = [[1, 2],       D = [[-1, 2],       dW = X.T @ D
     [3, 4]]           [ 0, 1]]          = [[-1, 5],
                                            [-2, 8]]

Every entry of dW is a legitimate derivative. But entries are not the only way to describe a matrix. Its singular directions tell us which combinations of input and output features dominate the proposed change. AdamW mainly maintains statistics per entry. Muon takes the proposed whole matrix update and changes its singular-value profile. That is the “matrix-aware” part; no forward-pass architecture changes.

One gradient, three different first steps

To see what that decision means, forget the housing data for one minute. Suppose a hidden-layer weight is a 2 x 2 matrix, initially all zeros. Backprop gives the same rotated gradient G = [[2.390, -1.030], [1.498, -0.269]] to each optimizer. For this one illustration only, set the learning rate to 0.1, disable weight decay, and ignore production Muon shape scaling. These are not tuned training runs; the point is to compare directions.

SGD changes each cell in proportion to its raw gradient. AdamW's first bias-corrected step is approximately -0.1 sign(G), so the four nonzero cells move by equal magnitude. Muon starts from momentum B = G; its Nesterov-style direction is 1.95G, then matrix multiplication in Newton-Schulz changes its singular-value profile. Scaling G by 1.95 does not change the normalized input to that operation. The resulting direction is not the AdamW sign matrix. Section 9 explains the polynomial and why a diagonal example alone can hide this difference.

What does the optimizer remember?

State is simply information carried from one minibatch to the next. Here is a ledger for the same illustrative matrix, with G1 = G and a hypothetical next gradient G2 = 0.5G. Use AdamW beta1 = 0.9, beta2 = 0.99, and Muon momentum 0.95. The coefficients below are multipliers of the full matrix, not performance results; G ⊙ G means square each entry.

For example, AdamW's step-2 mean is 0.9(0.1G) + 0.1(0.5G) = 0.14G. Muon's buffer is 0.95G + 0.5G = 1.45G. Both optimizers also update the weight on each step; this table shows stored state, not a waiting period. AdamW retains two arrays per eligible matrix, while Muon retains one and performs extra matrix work when forming its update. Real consecutive gradients need not be scaled copies, so later directions will usually change too.

Keep this comparison in mind as the model grows. The optimizer changes how a gradient becomes an update; it does not change the forward equation or the loss being differentiated.

3. Eight housing measurements turn the line into GEMM

Before adding an activation or a feed-forward network, let us train something with real, multi-dimensional input. I use the California Housing dataset: 20,640 California census-block records, eight numeric features, and a target of median house value for the block in units of $100,000. This is historical 1990 census data, not a dataset of individual home sale prices and not a present-day valuation tool. The features are median income, housing age, average rooms, average bedrooms, population, average occupancy, latitude, and longitude.

We use 16,512 rows for training and 4,128 for validation, with a fixed random split. On the training rows only, each feature is clipped at its 1st and 99th percentiles and standardized; the same learned limits, means, and scales are applied to validation. The target is centered and scaled by the training-target standard deviation for stable optimization. All reported RMSE values below are converted back to $100,000 units. A random block-level split can overstate geographic generalization because nearby blocks may appear in both sets; this is a teaching dataset, not a spatially held-out housing benchmark.

In the first model, 512 rows form a minibatch X [512,8]. One vector of weights makes y_hat = X @ w + b, with w [8,1] and y_hat [512,1]. It is still a straight line in eight dimensions, computed as GEMM rather than eight handwritten scalar multiplies. With half mean-squared error, D = (y_hat - y) / 512, dw = X^T @ D, and db = sum(D). No activation has entered the picture.

4. Follow one real gradient-descent step

Initialize w = 0 and b = 0. On the first seeded minibatch, the normalized target and standardized housing features produce a weight gradient of approximately:

dw = [-0.6689, -0.0413, -0.2324, +0.0955,
      +0.0190, +0.3124, +0.0878, +0.1185]
db = -0.0226

learning_rate = 0.05
w = w - learning_rate * dw
b = b - learning_rate * db

The first coordinate belongs to median income. Its negative gradient means gradient descent initially makes that coefficient positive. Average occupancy, in contrast, gets a positive gradient and initially moves negative. These are directions in standardized feature space, not causal claims about housing. At zero weights, validation RMSE is about 1.159 in $100,000 units; after this one minibatch update it is about 1.122. That is one concrete step, not a magic straight-line fit.

After 30 epochs, direct minibatch SGD reaches validation RMSE about 0.658. A least-squares fit to the same preprocessed training features gives about 0.657 on validation. Least squares is a reference linear fit, not a bound on validation error; a different fit could score better on this particular held-out sample. The important point is that ordinary gradient descent has nearly reached the performance of this linear function class. There is no hidden “Muon trick” needed to find a line.

5. A GEMM-only matrix stack creates an internal Muon target

Why not run Muon directly on w [8,1]? It is technically a 2D array, but it is also the output head, and its rank is at most one. The Muon authors recommend a standard optimizer for output heads. For a useful matrix-level teaching comparison without sneaking in an activation, factor the same linear mapping into three matrix products:

H1    = X  @ W1            # [512,8]  @ [8,16]  -> [512,16]
H2    = H1 @ W2            # [512,16] @ [16,16] -> [512,16]
y_hat = H2 @ W3 + b        # [512,16] @ [16,1]  -> [512,1]

W_effective = W1 @ W2 @ W3 # [8,1], still one linear rule

The backward pass is just as explicit. Let D = (y_hat-y)/batch_size [512,1]. Then dW3 = H2^T D [16,1]; dH2 = D W3^T [512,16]; dW2 = H1^T dH2 [16,16]; dH1 = dH2 W2^T [512,16]; and dW1 = X^T dH1 [8,16]. The bias gradient is a row sum. These shapes are the same GEMM and reduction story as the matrix Wx+b post, now traced through several products. I checked the hand-written gradients against finite differences; the largest absolute discrepancy in a sampled check was about 5e-11.

Crucially, there is no nonlinearity between W1, W2, and W3. The entire model collapses algebraically to X @ W_effective + b. It cannot represent interactions that a single linear regression cannot. The factorization only changes the parameterization and optimization path. This is why it gives us an honest GEMM-only laboratory before we later add ReLU-squared to build an FFN.

6. AdamW versus Muon on the same linear stack

Here is the preview of the two update rules before we unpack them: AdamW remembers a moving mean and moving squared mean for every matrix entry, while Muon remembers one momentum matrix and applies a short polynomial of matrix multiplications to it. The next three sections derive that distinction. First, let us see whether it changes this GEMM-only training run at all.

For each of three paired initialization seeds, both runs start from the same Xavier-uniform W1, W2, and W3 and see the same minibatch order for 30 epochs. AdamW updates every weight at learning rate 0.01. The hybrid applies Muon with learning rate 0.04 only to hidden W2, and AdamW at 0.01 to the input matrix, output matrix, and bias. Both apply 0.01 decoupled matrix weight decay. The direct-SGD line is a separate, differently parameterized baseline, so do not treat its per-epoch speed as an optimizer-only comparison with the factorized stack.

There is no Muon victory to manufacture here. At epoch 1, the factorized optimizers move quickly from their Xavier starts. By epoch 30, all three trained models are close to the same least-squares linear reference, and AdamW and Muon on the identical stack differ very little. For this single-output linear task, balancing the internal matrix update does not create a more expressive function. The gradient-descent story is visible, but the gain we hope for in richer hidden representations has no obvious reason to appear here.

This negative result is useful. Before the FFNMatrix shape is not enoughMuon has an eligible internal matrix here, but a stack of linear GEMMs is still one linear map. The next experiment adds a nonlinear activation so hidden feature learning can matter.

7. What AdamW remembers

For a gradient g_t, AdamW carries a first moment and a second moment for each parameter. After bias correction it divides the first by the square root of the second, then applies decoupled weight decay. Here is the compact version:

m = beta1 * m + (1 - beta1) * g
v = beta2 * v + (1 - beta2) * g * g
m_hat = m / (1 - beta1**step)
v_hat = v / (1 - beta2**step)
W = (1 - lr * weight_decay) * W - lr * m_hat / (sqrt(v_hat) + eps)

AdamW is valuable because its coordinate-wise scaling handles noisy, uneven gradients remarkably well. This is not “waiting for old gradients before it can train”: both buffers update and the weights change on every step. The price is two persistent moment arrays, normally the same shape as the parameter. For a 32-by-32 fp32 hidden weight, those two arrays alone are 2 * 1024 * 4 = 8192 bytes. The parameter, gradient, master weights, and implementation-specific scratch are separate costs.

There is a subtle first-step trap here. If AdamW starts with zero moments and sees G = diag(3, 0.3), its bias-corrected first step is approximately diag(1, 1) before multiplying by the learning rate. That resembles the polar-factor story below. It does not make AdamW a matrix orthogonalizer: AdamW has normalized two coordinates in the current parameter basis. Rotate the same update and the distinction becomes visible.

8. What Muon changes, and where it is allowed

Muon means MomentUm Orthogonalized by Newton-Schulz. It first builds a momentum buffer B for an eligible 2D hidden-layer weight. A common variant takes a Nesterov-style lookahead update, then applies five Newton-Schulz polynomial steps to that matrix before changing the weight. It keeps one persistent momentum array for that weight, not AdamW's first and second moments.

B = 0.95 * B + dW
G = dW + 0.95 * B               # Nesterov-style update direction
O = newton_schulz_5(G)          # matrix operation, not elementwise division
W = (1 - lr * decay) * W - adjusted_lr * O

The original Muon guidance is deliberately selective: use it on hidden-layer matrices; keep biases, scalar and vector parameters, embeddings, and output heads on a conventional optimizer such as AdamW. A matrix shape alone does not make a parameter suitable. In particular, the first and last 2D weights in our little network stay on AdamW. Production implementations also apply shape-dependent learning-rate scaling for rectangular matrices; the square hidden matrix in our experiment does not need that distinction.

Less persistent state is not “free training.” Newton-Schulz needs matrix multiplications and scratch space on every optimizer step. Whether the whole training run gets faster depends on convergence, batch size, matrix shape, hardware, and implementation. BoundaryMuon is not the forward GEMMGEMM builds activations and gradients. Muon runs after backprop to choose a matrix weight update.

9. The matrix “aha”: change singular values, keep directions

Imagine a proposed update G = diag(3, 0.3). Its first singular direction is ten times stronger than its second. The exact polar factor would replace both nonzero singular values with 1, giving diag(1, 1). Muon avoids explicitly computing an SVD. It normalizes the matrix, then repeatedly applies a polynomial made of matrix multiplications:

X = G / (norm_frobenius(G) + 1e-7)
repeat 5 times:
    A = X @ X.T
    X = 3.4445*X + (-4.7750*A + 2.0315*(A@A)) @ X

For a tall matrix, the implementation transposes it during this operation and transposes the result back. The polynomial acts on singular values while retaining the singular-vector directions. Five steps with these tuned coefficients are an approximation, not an exact polar decomposition: in a float64 calculation for diag(3, 0.3), the resulting diagonal is about diag(0.702, 0.709). What matters is that the two directions are now of comparable size, not that each is exactly one.

In SVD notation, if G = U S V^T, the polynomial maps each singular value s through 3.4445s - 4.775s^3 + 2.0315s^5 at each step while preserving U and V. That is the conceptual difference from AdamW's coordinate-wise division by sqrt(v). The optimizer is changing the geometry of a matrix direction, not just choosing a different learning rate for each cell.

For a more revealing 2-by-2 example, rotate the diagonal diag(3, 0.3) by 30 degrees on the left and -20 degrees on the right. The resulting update is approximately [[2.390, -1.030], [1.498, -0.269]]. At AdamW's first step, coordinate-wise normalization points approximately along [[1,-1],[1,-1]]. The exact polar factor points along [[0.643,-0.766],[0.766,0.643]]. No single scalar learning rate makes those the same direction. Muon's finite-step result is an approximation to the latter, not to the former. This is why “matrix-aware” is a useful description: a change of basis affects AdamW's elementwise statistics differently from a singular-value operation.

10. Now add an activation: a deliberately small FFN

The housing experiment establishes the boundary: GEMM alone makes a linear model even when several matrices are multiplied. A diagonal matrix also proves the update mechanism, not that a nonlinear model learns. Now we add ReLU-squared after two of the GEMMs and switch to a deterministic two-moons classification dataset: 512 training points and 256 held-out points per seed, with Gaussian noise. The network has three matrix multiplications, ReLU-squared activations, and a two-class softmax loss:

X [64, 2]
  @ W1 [2, 32]  + b1  -> ReLU² -> [64, 32]
  @ W2 [32, 32] + b2  -> ReLU² -> [64, 32]
  @ W3 [32, 2]  + b3  -> logits [64, 2]
  -> softmax cross-entropy against class labels

We start both runs from the same Xavier-uniform weights for each seed. The backward pass computes the logits gradient, then dW3 = A2^T dlogits, dW2 = A1^T dz2, and dW1 = X^T dz1. For ReLU-squared, d/dz ReLU(z)^2 = 2 ReLU(z). This is the same chain rule as wx+b, repeated through three layers. Every epoch visits eight minibatches of 64 examples.

Here is the middle-layer chain rule without skipping steps. Softmax-cross-entropy sends D3 = (P - one_hot(y)) / batch_size into the output head. The head's backward GEMM sends dA2 = D3 W3^T into the second activation. Its local derivative makes D2 = dA2 * 2 ReLU(Z2); the hidden-weight gradient is then dW2 = A1^T D2. This dW2 is the exact 32-by-32 object fed to Muon's momentum buffer. Continuing backward gives dA1 = D2 W2^T, D1 = dA1 * 2 ReLU(Z1), and dW1 = X^T D1. No optimizer can repair a wrong backward pass, so I also checked these hand-written derivatives against finite differences: the largest absolute discrepancy in a small check was about 1.5e-10.

Why cross-entropy rather than the housing model's squared error? The moons task predicts a probability for each of two classes. For one correct class, assigning it probability q costs -ln(q) nats: q = 0.8 costs about 0.223, while q = 0.6 costs about 0.511. For a one-hot target distribution, its entropy is zero, so this one-example cross-entropy also equals KL(one_hot || predicted). More generally, H(p,q) = H(p) + KL(p || q): KL is the avoidable mismatch when the true distribution p is fixed. We do not know the true conditional distribution for each moon point, so the table below reports held-out empirical cross-entropy, not a measured population KL. See the cross-entropy and KL walkthrough for the full derivation.

We compare (A) AdamW on every parameter at learning rate 0.01 with (B) Muon on hidden W2 at learning rate 0.04 plus AdamW at 0.01 on W1, W3, and all biases. Both use 0.01 decoupled weight decay on matrices. Muon uses momentum 0.95, Nesterov, and five Newton-Schulz steps. These are fixed illustrative settings, not a matched hyperparameter sweep or proof that Muon is generally superior. The NumPy implementation uses float64 for clarity; the original reference uses lower-precision GPU arithmetic.

The chart stops at epoch 30. The epoch-100 row comes from a separate 100-epoch rerun with the same seeds, initialization, data order, model, and fixed learning rates. Its full epoch-by-epoch results are reproducible with --epochs 100; it is not an extrapolation of the 30-epoch plot. Each epoch visits 512 training examples per seed, so 1, 10, and 100 epochs mean 512, 5,120, and 51,200 example-visits respectively. Those are repeated visits to the same 512 points, not fresh data.

The useful observation is nuanced. Muon has lower validation loss at epochs 10, 15, and 30 in this small run, but AdamW is lower at epoch 20; both reach essentially the same classification accuracy. A lower loss after the same number of epochs is sample efficiency on this task, not evidence of faster wall-clock training or generalization to a language model.

Reading the curve as a training story: at epoch 1, both networks barely know the moons; by epoch 5, they are learning the broad boundary. Around epoch 10, the Muon hybrid's mean cross-entropy is lower, even though it touched the same 5,120 training examples per seed. By epochs 15 to 30, accuracy has saturated and small loss differences mostly reflect how confident the classifier is on a few difficult points. The held-out set is only 256 examples, so a single changed prediction moves accuracy by roughly 0.39 percentage points. This is why I show both cross-entropy and accuracy, three seeds, and the full trajectory instead of announcing a “winner” from the last row.

At epoch 100 the order reverses: AdamW's mean held-out cross-entropy is lower, while accuracy remains nearly identical. Muon's lower epoch-10 loss is not a permanent advantage. Longer training can change calibration, confidence on a few points, and the relative ranking; this tiny fixed-split experiment cannot tell us which optimizer wins on a different dataset.

11. Speed has two clocks

The optimizer itself is more expensive in this tiny NumPy CPU experiment. In one rerun, the average optimizer-only step across three seeds took roughly 0.09 ms for AdamW and 0.24 ms for the Muon hybrid. These sub-millisecond timings move with process noise and BLAS settings, so do not read them as a hardware benchmark. They do show the right accounting question: does a more expensive step reduce the number of steps enough to win in total time? Our 30-epoch toy does not establish that it does.

At larger matrix and batch sizes, the forward and backward GEMMs can dominate the step, making the extra optimizer matrix multiplies a smaller fraction. Conversely, small batches, unusual shapes, distributed synchronization, or a poor implementation can make Muon less attractive. A fair production comparison needs tuned baselines, paired seeds, loss versus tokens processed, loss versus wall time, peak memory, and optimizer-step profiles. It should report failure as readily as success.

There is also no universal “Muon learning rate.” The orthogonalized direction has a different scale from an AdamW update, and rectangular layers can need shape-aware scaling. To compare optimizers responsibly, tune each within a declared search budget on a validation task, then run the chosen settings on fresh seeds. Keep batch size, initialization, data order, and model fixed. Report the full tuning budget. Our fixed 0.01 versus 0.04 settings are a teaching control, not that study.

A scaling spectrum, not a magic crossover

Does a larger GEMM make Muon faster? Not automatically. For a square hidden matrix of width n and a batch of B rows, the Muon design analysis bounds the extra five-step Newton-Schulz matrix work relative to that layer's forward/backward GEMMs by about 5n/B. Increasing the batch amortizes the matrix-polynomial work; increasing width at a fixed batch makes it more prominent. This is a FLOP ratio, not a measured speedup, and it does not say AdamW itself gets slower relative to Muon.

To make the boundary concrete, I ran a separate synthetic one-layer sweep on this machine: fp32 NumPy, single-thread OpenBLAS, square weights, and batch sizes from 64 to 4,096. I timed a shared forward/backward GEMM and each optimizer update separately, then summed their median component times. The first plot is that CPU estimate; the second is the analytic extra-FLOP bound. Neither plot includes data loading, a full neural network, GPU/BF16 kernels, or an accuracy target.

These are medians of seven component timings in one process, not a hardware benchmark or a prediction for a transformer. The FLOP bound differs from the wall-clock ratio because memory traffic, Python allocations, BLAS efficiency, and elementwise work also cost time. The key observation is narrower: the extra optimizer work becomes a smaller fraction as B/n grows, but Muon does not become faster per step in this sweep. In the first and last rows, for example, the optimizer's mathematics is unchanged; only the matrix and batch dimensions differ.

The real crossover is time to a target loss. If a Muon step costs r times an AdamW step, Muon must reach the same prespecified validation target in less than 1/r as many steps to win on training time, before accounting for setup and evaluation. At the 64 x 4,096 point, r = 1.07, so roughly 7% fewer steps could offset the per-step cost. At 512 x 4,096, r = 1.66, so it would need roughly 40% fewer steps. This synthetic sweep does not measure whether either saving occurs. The earlier 100-epoch two-moons result also does not establish a win: it shows a different loss ranking after longer training on a tiny task.

12. Why train modern AI with Muon?

The strongest case is not that Muon makes every optimizer step cheaper. It can make the path to a chosen loss more data- or compute-efficient on some hidden-matrix-heavy models. By changing the singular-value profile of a proposed update, it can reduce the dominance of a few strong matrix directions, a pattern its designers observed for many transformer weights. That is a hypothesis about useful learning directions, backed by particular training results, not a guarantee for every layer or dataset. The one-buffer state saving is real for eligible matrices, but it is a secondary and narrower benefit than “half the model memory.”

DeepSeek's optimizer section and Algorithm 1 make the GEMM connection precise. The model's forward and backward GEMMs still compute activations and gradients. Muon runs after backprop on each eligible weight matrix: momentum, Nesterov-style direction, matrix products inside Newton-Schulz, shape-aware rescaling, then the weight update. DeepSeek's implementation is not our five-step teaching code. It uses eight iterations with the fast coefficients followed by two with stabilizing coefficients, and rescales the update so it can reuse its AdamW learning-rate schedule. Its attention design uses RMSNorm on queries and KV entries, so the report says it does not use Kimi's QK-Clip.

That is also a systems problem. The training-framework section says the whole matrix needed for Muon conflicts with ordinary elementwise ZeRO sharding. The team assigns logically independent matrices to ranks, batches same-shaped Newton-Schulz work, and uses BF16 matrix multiplication for those iterations. These are optimizer-side GEMM-like kernels, not a replacement for the model's forward GEMM kernels. DeepSeek attributes faster convergence and stability to its V4 recipe, but the paper's architecture and training system changed too; it does not make our CPU sweep a DeepSeek-scale benchmark or prove that the optimizer alone caused every gain.

These are the teams' reported results under their models, tuning, hardware, and training recipes, not a claim that our 512-point moons experiment reproduces them. They also explain why the little CPU spectrum and real large-model outcomes can disagree: a slower optimizer step can still win if it needs sufficiently fewer steps or tokens to hit the same quality target. NVIDIA separately reports near-parity Muon and AdamW training throughput on its optimized GB300 setup; that is an implementation-throughput result, not by itself proof of better final model quality.

The practical answer is selective. Use a conventional optimizer for embeddings, heads, and vectors; consider Muon for eligible hidden matrices when a tuned, paired experiment shows a better loss-versus-time or loss-versus-token curve. Expect to solve update scaling, numerical stability, and distributed matrix-work problems at scale. Do not switch optimizers solely because a matrix is large or because its state arrays are smaller.

13. The state-memory difference is real, but specific

On our W2 [32,32], AdamW keeps two 1024-element moment arrays. Muon keeps one 1024-element momentum array. In fp32, that is 8 KiB versus 4 KiB of persistent optimizer state for that weight. It does not halve total model memory: weights, gradients, activations, AdamW state for the input/output layers, and Newton-Schulz temporary matrices still exist. Our educational NumPy run uses float64, so its actual byte counts are double the fp32 illustration.

That state trade is one reason Muon attracts attention, but the more interesting idea is the update geometry. If the raw momentum is dominated by a few singular directions, the polynomial makes weaker nonzero directions matter more. That can change how quickly hidden features diversify. It can also be the wrong choice for a particular task or layer. This is an empirical optimizer, not a theorem that every matrix should be orthogonalized.

14. Where this fits beside C-Kernel-Engine

The C-Kernel-Engine optimizer kernels provide the existing SGD/AdamW side of this story, and the matrix Wx+b post explains the GEMM shapes that produce weight gradients. CKE does not currently ship the Muon kernel used in this experiment; these plots come from a separate, transparent NumPy implementation. A CKE Muon path would need explicit parameter-group selection, a momentum buffer, the five matrix polynomial steps, shape scaling, numerical parity tests, and wall-time/peak-memory measurements. Until those are implemented and tested, I will not label these results as CKE results.

The progression I want readers to keep is simple: scalar wx+b gives a scalar gradient; housing data turns it into an eight-feature GEMM; a linear matrix stack gives Muon an internal weight but no new expressivity; an activation turns the stack into an FFN. Backprop tells us what each gradient is; the optimizer decides how to move the weight. AdamW scales each cell with its own history. Muon asks whether a hidden-layer matrix update should share its strength more evenly across singular directions. That is the “aha,” and also the limit of what these small experiments can prove.

Reproduce and go deeper

Both experiments use Python 3, NumPy, and Matplotlib with OPENBLAS_NUM_THREADS=1. Run OPENBLAS_NUM_THREADS=1 python3 scripts/experiments/muon_gemm_housing.py for the checksum-verified housing archive, GEMM-only metrics, SVG, and plots; then run OPENBLAS_NUM_THREADS=1 python3 scripts/experiments/muon_walkthrough.py for the 30-epoch FFN classification metrics and plots. To reproduce the longer checkpoint separately, run OPENBLAS_NUM_THREADS=1 python3 scripts/experiments/muon_walkthrough.py --epochs 100 --output workspace/blog-writing/muon-100 and inspect its JSON history. The source archive and feature transformations match scikit-learn's California Housing loader; scikit-learn itself is not needed to run the scripts. If you vary BLAS threads, focus on curves rather than treating tiny timing changes as meaningful.

For the scaling-spectrum figure and its full 16-point JSON table, run OPENBLAS_NUM_THREADS=1 python3 scripts/experiments/muon_scaling_spectrum.py. It measures update cost, not convergence or peak memory. To answer “when is Muon better?” for a larger model, pair this cost profile with a fixed-target loss study using tuned AdamW and Muon learning rates, multiple seeds, equal data order, wall-clock time, and actual peak-memory measurements.

For the canonical algorithm and its limitations, read Keller Jordan's Muon writeup, the reference implementation, and the PyTorch Muon documentation. The earlier ShivasNotes Muon comparison places the idea next to SGD and AdamW; this post is the worked numerical companion.

Keep reading: from one derivative to an optimizer

If one step in this article felt too fast, follow the mathematical chain in order:

1. Backpropagation from wx+b: the smallest training loop derives the scalar gradients used in section 1.
2. Matrix Wx+b: from scalars to transformers turns that derivative into the forward and backward GEMMs in sections 2 and 5.
3. dL/d(LLM): the full backward pass shows where the matrix gradients come from in a transformer.
4. The SGD-to-AdamW guide builds the optimizer state that Muon is compared against here.
5. Where cross-entropy loss comes from explains the classification loss and KL connection in the FFN experiment.
6. SGD vs AdamW vs Muon is the earlier broad comparison; the present article supplies the worked examples and scaling caveat.