MTP + Muon (Exp09)¶
Paper Section 2.1 (MTP), Section 2.4 + Algorithm 1 (Muon) ·
Code: src/deepseek_v4_lab/mtp/head.py, src/deepseek_v4_lab/optim/muon.py
Evidence tags: PAPER stated by DeepSeek · DERIVED mathematically implied · IMPLEMENTED reproduced in this project · MEASURED observed in our experiments · INFERRED our interpretation · UNKNOWN not established. Full definitions: evidence ledger
Why this exists
Problem: element-wise optimizers ignore matrix geometry; unnormalized update directions let attention logits grow until training needs patch fixes.
Response: orthogonalize each update matrix via hybrid Newton-Schulz iterations, rescale update RMS for AdamW-LR reuse, keep AdamW for embeddings/heads/norms.
New cost: orthogonalization assumes dense gradients; bf16 iterations have null-space quirks; more compute per step than element-wise updates.
Our experiment: Exp09 verifies singular values land within ±0.15 of 1 and decoupled decay is exact; toy loss ordering reported honestly [MEASURED].
Exchange rate: CONSTRAINT (spectral) buys healthy update geometry; spends per-step compute and gradient-density assumptions.
MTP [IMPLEMENTED + MEASURED]¶
Depth-1 head per V3 design: predict t+2 from h_t ⊕ embed(t+1).
Markov-LM result: NTP-component loss 2.1379 (NTP-only) vs 2.1403 (NTP+MTP training), parity at toy scale, at the cost of one extra loss term. The paper keeps MTP for its speculative-decoding payoff and richer representation supervision; our toy neither confirms nor denies scale benefits [MEASURED-at-toy-scale, no generalization].
Muon [IMPLEMENTED + MEASURED]¶
- Hybrid Newton-Schulz: 8 fast iters (3.4445, −4.7750, 2.0315) then 2 fine iters (2, −1.5, 0.5) [PAPER Section 2.4]; singular values land in ±0.15 of 1.
- Update RMS rescaled to 0.18 for AdamW-LR reuse [PAPER Alg 1 line 7].
- Decoupled weight decay W(1−ηλ) − ηO verified exactly.
- AdamW split for embeddings/heads/norms [PAPER Section 2.4].
| optimizer | final loss (Markov LM) |
|---|---|
| SGD | 2.309 |
| AdamW | 2.138 |
| Muon (untuned LR) | 2.174 |
At toy scale with untuned hyperparameters the ordering AdamW < Muon < SGD is expected. The architectural claim is different: Muon's isometric updates remove the logit-growth pathologies that motivated QK-normalization in the first place; and V4's RMSNorm'd q/KV make QK-Clip unnecessary [PAPER Section 2.4].
Lab finding [MEASURED]¶
bf16 NS iterations amplify exact-zero singular directions to ~1 (rounding seeds the null space). Real gradients are dense so production never sees this, but it explains why implementations normalize by Frobenius norm and rely on gradient density.
Limitations¶
- 400 steps, one seed; no LR tuning for Muon (paper uses 0.02-class LRs at cluster scale).
- No attention-logit growth instrumentation yet (needs the integrated model).