Kimi K3 · Lesson 9 of 14 · Optimisation
The last piece of the model sits outside the network entirely: how its weights get updated, and the empirical study that produced the claimed 2.5×.
Architecture is only half of what determines whether a 2.8T model trains well. This lesson covers the optimiser and the empirical study that set K3's hyperparameters. Both are short sections in the report, and both contain a lesson that transfers well beyond K3.
Start with the optimiser everyone knows. Adam-family optimisers treat every weight as an independent scalar: track a running average of each weight's gradient, scale by a running average of its magnitude, step. Nothing in that procedure knows that the weights are arranged in a matrix.
Muon takes the arrangement seriously. Before applying an update to a matrix parameter, it orthogonalises the momentum matrix.
WHY ORTHOGONALISE AN UPDATE
raw momentum matrix after orthogonalisation
─────────────────── ───────────────────────
one direction dominates all directions comparable
████████████████████ ◀── the ██████
██ update ██████
█ is mostly ██████
█ this row ██████
the step moves the weights the step uses the full
almost entirely one way capacity of the matrix
Now K3's refinement, and it is worth reasoning through because the flaw it fixes is a general one.
Attention does not use one big matrix. It uses many heads, each reading and writing its own slice of the representation. K3 has 96 of them. But the query, key, and value projections are stored as single large matrices, with the heads sitting side by side inside them.
So when Muon orthogonalises the full Q, K, or V projection, it treats all 96 heads as one coupled block. The report spells out the consequence: heads with larger gradient or momentum scales dominate the shared update direction, while smaller-scale heads receive insufficiently normalized updates.
1
FULL-MATRIX MUON PER-HEAD MUON
┌───────────────────────────┐ ┌─────┬─────┬─────┬─────┐
│ h1 │ h2 │ h3 │ ... │ h96 │ │ h1 │ h2 │ h3 │ h96 │
└───────────────────────────┘ └─────┴─────┴─────┴─────┘
│ │ │ │ │
▼ ▼ ▼ ▼ ▼
one orthogonalisation orthogonalise each
over everything head separately
a loud head sets the scale every head gets its own
for all 96; quiet heads get properly normalised update
updates that are barely
normalised at all
── and it is cheaper: Newton-Schulz on tall per-head blocks
costs less than on the full projection matrix
more balanced learning dynamics across heads and improves training stability at larger scales.
This is the same error shape you have already seen twice in this course. AttnRes exists because summing all earlier layers into one state loses the distinction between them. Channel-wise decay in KDA exists because one scalar decay rate for a whole memory cannot express different retention needs. Per-Head Muon exists because one orthogonalisation across 96 heads cannot express different update scales. When you see a single statistic standing in for many things that differ, that is a place to look for an improvement.
The report also notes it slightly reduces optimizer overhead
, since Newton–Schulz on tall thin per-head blocks is cheaper than on the full matrix. A quality improvement that also runs faster is rare enough to be worth noticing.
Muon has an awkward requirement under distributed training: orthogonalisation needs the full parameter matrix, but the optimiser shards parameters across data-parallel ranks. The obvious fix is an all-gather over the whole parameter buffer on every rank, which the report notes makes communication the primary bottleneck at scale.
K3 instead has each rank fetch only the shards of parameters it actually owns, peer to peer, then pipelines that communication against computation. The full-parameter buffer disappears entirely.1
Now the number quoted in every summary of K3, and what it actually means.
A scaling law is an empirical curve: train a family of small models at varying sizes and compute budgets, measure loss, and fit a relationship you can extrapolate. Its practical purpose is to choose hyperparameters for a run you can only afford to do once.
What the study retuned, because the architecture changes moved the optimum: batch size, learning rate, tokens-per-parameter ratio, and model shape.1 That last one is why K3 has 93 layers and a 7,168 hidden dimension rather than some other combination. The shape was fitted, not guessed.
The 2.5× is the authors' own fit, on their own held-out validation set, comparing against their own previous model. Nobody has replicated it. It is also not decomposed: you cannot tell from the report how much came from KDA versus AttnRes versus better data, and "refined data recipes" is doing unknown work inside that figure. When explaining K3 to someone, say "the report claims a 2.5× scaling-efficiency gain, attributed to the architecture and data changes together." That is both accurate and appropriately hedged.
The most interesting paragraph in §3.2 has nothing to do with 2.5×, and it is a lesson in experimental method rather than architecture.
Two learning-rate schedules compete in current practice. Cosine decay smoothly lowers the rate along a cosine curve. Warmup-Stable-Decay holds the rate constant for most of training, then drops it at the end, which is convenient because you can stop at any point. Recent work has reported WSD matching or beating cosine.
K3's study finds the opposite, and the explanation is the valuable part.
THE METHODOLOGICAL TRAP
the usual comparison what K3 did
──────────────────── ───────────
pick one set of hyperparameters search for the optimal
│ hyperparameters for
▼ EACH schedule separately
run cosine ──▶ loss A │
run WSD ──▶ loss B ▼
│ cosine at its own optimum
▼ ──▶ loss A'
declare a winner WSD at its own optimum
──▶ loss B'
but the two schedules have │
"markedly different optimal ▼
hyperparameters" — so this a fair comparison:
just measures which schedule cosine consistently wins
the shared settings suited
substantiallydifferent peak learning rates and batch sizes. Comparing them under shared settings
may unfairly favor one simply because those hyperparameters are better aligned with it.1
This is a trap worth recognising in any paper you read. Whenever two methods are compared under one shared configuration, ask whether that configuration was tuned for one of them. It very often was.
K3's final recipe, for reference: cosine decay, 1% linear warmup, weight decay 0.1 throughout, Per-Head Muon plus K2's weight-clipping mechanism, Quantile Balancing for the experts. Context starts at 8K and reaches 64K in pre-training, then 256K to 1M during cooldown.1
What problem does orthogonalising per head solve that full-matrix Muon has?
One orthogonalisation across 96 heads produces one shared scale, so loud heads set it and quiet heads get updates that are barely normalised. Partitioning along the head dimension gives each head a properly scaled update, and costs slightly less to compute.
The 2.5× scaling-efficiency claim should be read as which of these?
Read Figure 7 horizontally: for a target validation loss, K3's curve reaches it at a lower FLOP count. The gain is attributed to architecture and data changes together, and the report does not decompose it, so no single component can be credited.
Why did K3's study find cosine decay beating WSD when prior work found otherwise?
The two schedules want markedly different peak learning rates and batch sizes even at identical scale and token budget, so a shared configuration silently favours whichever it happens to suit. K3 ran an independent hyperparameter search for each, and cosine won consistently. Worth watching for in any two-method comparison.
✅ SEQUENCE KDA · chunkwise · g_min · Gated MLA · NoPE L2–L3
✅ DEPTH Block AttnRes · N = 8 · pseudo-queries L4
✅ WIDTH LatentMoE · SiTU-GLU · Quantile Balancing L6–L7
✅ INPUT MoonViT-V2 · from scratch · pixel-shuffle L8
✅ OPTIMISE Per-Head Muon · scaling law · cosine > WSD L9
────────── every architectural claim in the report ──────────
▢ REVIEW B interleaves lessons 6–9 L10
▢ SYSTEM serving · cost · cache · quantization L11
▢ BEHAVIOUR post-training · effort levels · limitations L12
▢ REVIEW C capstone L13
Lesson 10 is Review B, interleaving lessons 6 through 9 in the same free-recall format as Review A, with the discrimination questions leaning on distinctions that are easy to blur: SiTU-GLU against Quantile Balancing, latent compression in MoE against latent compression in attention, Per-Head Muon against the architectural changes.