From 68ac1f9d1fa8944a216964479e49e813de6370ac Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Wed, 5 Aug 2026 08:18:11 -0500 Subject: =?UTF-8?q?=E5=81=8F=E7=BD=AE=E6=94=BE=E5=A4=A7=E5=85=AC=E7=90=86(?= =?UTF-8?q?=E7=94=A8=E6=88=B708-06=E6=8C=87=E5=87=BA):=20=E9=80=90?= =?UTF-8?q?=E5=9D=90=E6=A0=87=E5=BD=92=E4=B8=80=E5=8C=96=E6=8A=8A=E6=8C=81?= =?UTF-8?q?=E4=B9=85=E5=B0=8F=E5=81=8F=E7=BD=AE=E6=94=BE=E5=A4=A7=E6=88=90?= =?UTF-8?q?=E6=BB=A1=E5=B9=85=E6=BC=82=E7=A7=BB,=20sign=E5=85=A8=E5=AE=B6?= =?UTF-8?q?=E5=90=8C=E7=BD=AA(=E5=90=A6=E5=86=B3=E5=8F=AA=E6=8C=A1=20?= =?UTF-8?q?=E6=96=B9=E5=B7=AE=E4=B8=8D=E6=8C=A1=E5=81=8F=E7=BD=AE);=20?= =?UTF-8?q?=E6=9B=B2=E7=8E=87=E9=80=9A=E9=81=93=3D=E5=81=8F=E7=BD=AE?= =?UTF-8?q?=E7=A3=81=E9=93=81(E[s=C2=B7=C4=9D]=3D=CE=B2c+=CE=B4/=CE=B2,=20?= =?UTF-8?q?=E6=B1=A1=E6=9F=93=E8=A2=AB1/=CE=B2=E6=94=BE=E5=A4=A7333=C3=97)?= =?UTF-8?q?,=20=E5=AD=98=E6=B4=BB=3D=E4=B8=A4=E6=A1=A3|=CE=B2|=E9=94=81?= =?UTF-8?q?=E7=9B=B8=E5=88=86=E7=A6=BB;=20=E8=AF=84=E5=88=86=E8=A1=A8?= =?UTF-8?q?=E5=A2=9E=E7=AC=AC=E4=B8=89=E8=BD=B4=3D=E5=81=8F=E7=BD=AE?= =?UTF-8?q?=E6=94=BE=E5=A4=A7=E7=B3=BB=E6=95=B0,=20=E5=99=A8=E4=BB=B6?= =?UTF-8?q?=E7=94=B5=E6=B1=A0=E9=A2=84=E6=B3=A8=E5=86=8C=E6=8E=92=E5=BA=8F?= =?UTF-8?q?=E9=87=8D=E6=8E=92=E9=A2=84=E6=B5=8B(per-matrix=E4=BC=98?= =?UTF-8?q?=E9=9B=85=E9=80=80=E5=8C=96/sign=E6=BB=A1=E5=B9=85/Dither?= =?UTF-8?q?=E7=BF=BB=E8=BA=AB)=E3=80=82=20=E9=99=84:=20--dump=5Fgrad?= =?UTF-8?q?=E6=8E=A2=E9=92=88+cautlion=E5=B7=B2=E5=AE=9E=E8=A3=85=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/muon.py | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) (limited to 'ep_run/muon.py') 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 -- cgit v1.2.3