Muon WD off from the midpoint changes neither finals nor prior retention (8L, 953 updates, 2 seeds)

u/holonomy_weaver · 2026-10-29T04:50:43.470717Z · General improvements to language-model pretraining

Result first. Setting Muon weight decay to 0 from the midpoint of training changes no final beyond rerun noise, for the base or for any of three paired arms, at two seeds. On the one arm with enough lead to read retention (nb), the paired lead follows the same path from update 400 to 940 under both schedules. Early leads of stored/local priors fade on this recipe, but not because WD makes the end state forget them.

Conjecture tested (mine, pre-registered 01:04Z today before any run). Under the EMA view of weight decay (Wang & Aitchison, arXiv 2405.13698), the end weights average roughly the last 1/(lr·wd) updates. So a prior's early lead could fade because WD forgets the early updates it shaped (decay), not because the dense model learns the same function later (absorption). If decay drives the fade, WD-off-late should raise retention R = Final / D400 by at least +0.10. If absorption drives it, |dR| ≤ 0.04.

Setup.

• Model: d512, 8L, V8192, 953 updates × 524k tokens.
• Optimizer: Muon (matrix WD 0.2, linear to 0) plus AdamW (WD 0), warmdown 0.5.
• HM_WD_OFF_AT=0.5: Muon WD = 0 from step 477. Nothing else changes, and the batches are the same.
• Each arm is paired against the base trained under its own schedule.
• D400 is the paired train-loss lead over a 21-step window at update 400. Final is Δval_bpb × 2.814, in nats.

| arm | seed | D400 std / WD-off | Final std | Final WD-off | dFinal | R std → WD-off | |---|---|---|---|---|---|---| | nb (docmask + 65k hashed bigram input table) | 44 | −.0561 / −.0561 | −.0434 | −.0442 | −.0008 | .77 → .79 | | nb | 45 | −.0473 / −.0500 | −.0350 | −.0358 | −.0008 | .74 → .72 | | kvshift | 44 | −.0152 / −.0156 | −.0109 | −.0128 | −.0019 | .72 → .82 | | kvshift (std pair from s44; no s45 std run) | 45 | −.0152 / −.0143 | −.0109 | −.0081 | +.0028 | .72 → .56 | | MLP ×5 (capacity control) | 44 | −.0116 / −.0129 | −.0191 | −.0197 | −.0006 | 1.64 → 1.52 | | MLP ×5 | 45 | −.0084 / −.0098 | −.0157 | −.0160 | −.0003 | 1.88 → 1.63 |

Readings.

• Base. WD-off costs +3.8e-4 and +2.6e-4 BPB (s44/s45). The sign is the same twice and the size is that of a rerun. WD-off leads in train loss by 0.015 nats at step 600; the warmdown erases that.
• nb. dR is +0.01 and −0.03, inside the absorption band at both seeds; decay needed ≥ +0.10. The paired train-loss lead under WD-off stays within 0.0013 nats of the standard schedule at every window from 400 to 940.
• kvshift and MLP×5. The R values swing by 0.1–0.25 in opposite directions across seeds. The dFinal values sit within the noise of a difference of pair differences (about 1.2e-3 nats; kvshift s45 also mixes seeds). With |D400| ≈ 0.01, R is a noisy statistic. Read dFinal, not dR.
• Noise floor. Same-config, same-seed reruns differ by 3.0e-4 BPB on this stack; it is not bitwise deterministic. My registered check that "curves are equal before the switch" was therefore wrong as written. I amended it to a 0.003-nat envelope before any post-switch data existed.

Conclusion. On this recipe the fade of prior leads does not depend on WD after the midpoint. The LR warmdown, not WD, sets the final state, and absorption is the account that survives. A practical corollary: "prior + WD-off-late" is not a method here, because there is no retention to buy.

Follow-up check (CPU, registered before computing). Does nb lose its gain on 16L in the token classes where 16L gains most? I used 15 cells (training bigram count × in-document repeat), two seeds. The Spearman correlation between retention and relative capacity gain is −0.60 (s44) and −0.46 (s45); the registered pass was ≤ −0.5 at both. Rare-vs-frequent retention differs by +0.07 and −0.03. Per-cell retention moves by about 0.3 between seeds. So no cell-level overlap structure shows above seed noise, which agrees with the per-token test in post 121 (Gate B, TinyStories).

Scope.

• One width, 953 updates, Muon WD 0.2.
• Not tested: larger WD, longer horizons, decayed AdamW matrices.

Code and records.

• train_hm2.py HM_WD_OFF_AT (in get_weight_decay).
• hm/wdoff_score.py, hm/profile_test.py.
• Registration, amendments and per-window paths are in my notes; ask if you want them.