summaryrefslogtreecommitdiff
path: root/ep_run/muon.py
diff options
context:
space:
mode:
authorYuren Hao <yurenh2@illinois.edu>2026-08-05 08:18:11 -0500
committerYuren Hao <yurenh2@illinois.edu>2026-08-05 08:18:11 -0500
commit68ac1f9d1fa8944a216964479e49e813de6370ac (patch)
tree3a5aedebe30df058079ddc123fe04b63c7f5b7e2 /ep_run/muon.py
parentd0d1849e16127d4248bf5ef6c0bb1d6686aa9924 (diff)
偏置放大公理(用户08-06指出): 逐坐标归一化把持久小偏置放大成满幅漂移, sign全家同罪(否决只挡
方差不挡偏置); 曲率通道=偏置磁铁(E[s·ĝ]=βc+δ/β, 污染被1/β放大333×), 存活=两档|β|锁相分离; 评分表增第三轴=偏置放大系数, 器件电池预注册排序重排预测(per-matrix优雅退化/sign满幅/Dither翻身)。 附: --dump_grad探针+cautlion已实装。 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn
Diffstat (limited to 'ep_run/muon.py')
-rw-r--r--ep_run/muon.py26
1 files changed, 25 insertions, 1 deletions
diff --git a/ep_run/muon.py b/ep_run/muon.py
index e6c2692..1289620 100644
--- a/ep_run/muon.py
+++ b/ep_run/muon.py
@@ -176,7 +176,8 @@ def build_alt(opt_name, blocks, other_params, lr, warmup, total_steps=0, lr_min_
'adafactor': lambda: Adafactor2D(mats, lr=lm, wd=wd),
'signline': lambda: SignLine(mats, lr=lm, wd=wd),
'conslion': lambda: ConsensusLion(mats, lr=lm, wd=wd),
- 'ditherlion': lambda: DitherLion(mats, lr=lm, wd=wd)}[opt_name]()
+ 'ditherlion': lambda: DitherLion(mats, lr=lm, wd=wd),
+ 'cautlion': lambda: CautiousLion(mats, lr=lm, wd=wd)}[opt_name]()
oa = torch.optim.AdamW(rest, lr=lr, weight_decay=1e-4)
opts = [om, oa]
if total_steps > 0:
@@ -247,6 +248,29 @@ class ConsensusLion(torch.optim.Optimizer):
p.add_(u, alpha=-g_['lr'])
+class CautiousLion(torch.optim.Optimizer):
+ """C-Lion (Liang et al., arXiv:2411.16085, ICLR'26): Lion masked where the update sign
+ disagrees with the CURRENT gradient sign. The mandatory prior-art baseline for ConsensusLion;
+ the difference under test is instantaneous-gradient gating (this) vs filtered two-EMA gating."""
+ def __init__(self, params, lr=3e-4, betas=(0.9, 0.99), wd=0.0):
+ super().__init__(params, dict(lr=lr, betas=betas, wd=wd))
+
+ @torch.no_grad()
+ def step(self, closure=None):
+ for g_ in self.param_groups:
+ b1, b2 = g_['betas']
+ for p in g_['params']:
+ if p.grad is None: continue
+ st = self.state[p]
+ if 'm' not in st: st['m'] = torch.zeros_like(p)
+ m = st['m']
+ u = (b1 * m + (1 - b1) * p.grad).sign_()
+ u = u * ((u * p.grad) > 0)
+ if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd'])
+ p.add_(u, alpha=-g_['lr'])
+ m.mul_(b2).add_(p.grad, alpha=1 - b2)
+
+
class DitherLion(torch.optim.Optimizer):
"""Lion with dithered sign: sign(m + tau*noise*rms(m)). Free substrate noise turns the hard
sign into an unbiased soft-sign in expectation, letting small entries carry proportional