Estimator arms machinery: --est centered/richardson (two-pass, O(b^2) bias), --muon_mom/--adam_b1 knobs; smoke passed (resume+centered cos 0.998)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn

1 files changed, 6 insertions, 4 deletions

diff --git a/ep_run/muon.py b/ep_run/muon.py
index 7281811..d65a231 100644
--- a/ep_run/muon.py
+++ b/ep_run/muon.py
@@ -56,15 +56,17 @@ class MultiSched:
         for s in self.scheds: s.step()
 
 
-def build_hybrid(blocks, other_params, lr_adamw, lr_muon, warmup, total_steps=0, lr_min_ratio=0.1):
+def build_hybrid(blocks, other_params, lr_adamw, lr_muon, warmup, total_steps=0, lr_min_ratio=0.1,
+                 muon_mom=0.95, adam_b1=0.9):
     """Muon(2D block matrices) + AdamW(everything else). Scheds: linear warmup, then cosine decay to
-    lr_min_ratio*peak if total_steps>0 (long runs), else constant after warmup (legacy)."""
+    lr_min_ratio*peak if total_steps>0 (long runs), else constant after warmup (legacy).
+    muon_mom/adam_b1: momentum knobs (late-SNR noise-averaging arms, 2026-07-13)."""
     import math as _m
     mats = [p for p in blocks.parameters() if p.ndim == 2]
     mat_ids = {id(p) for p in mats}
     rest = [p for p in other_params if id(p) not in mat_ids]
-    om = Muon(mats, lr=lr_muon)
-    oa = torch.optim.AdamW(rest, lr=lr_adamw, weight_decay=1e-4)
+    om = Muon(mats, lr=lr_muon, momentum=muon_mom)
+    oa = torch.optim.AdamW(rest, lr=lr_adamw, weight_decay=1e-4, betas=(adam_b1, 0.999))
     if total_steps > 0:
         def fn(s):
             if s < warmup: return (s + 1) / max(warmup, 1)