diff options
Diffstat (limited to 'ep_run/muon.py')
| -rw-r--r-- | ep_run/muon.py | 26 |
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 |