Lucas Nestler released a code implementation called SFMuon.
Cool + Prof. Chulhee was my defense's committee and he is really kind
🚨New Optimizer Paper AMUSE: Anytime MUon with Stable gradient Evaluation
AMUSE combines Muon with Schedule-Free-style gradient evaluation for stable anytime training without LR decay.
• Stronger 124M / 720M / 1B pretraining • Strong ImageNet / ViT fine-tuning performance.
Positive users praise the AMUSE Optimizer paper because they view the choice of prof chulhee as solid and call the name clever.
class SFMuon(C.ScheduleFree): def __init__(self, params, **kw): d = dict(lr=.02, beta=.95, weight_decay=0, warmup_steps=0, weight_lr_power=2., r=0.) d.update(kw) super().__init__(params, d, fns=(C.nesterov_ema, C.orthogonalize_update, C.update_by_schedule_free))
🚨New Optimizer Paper AMUSE: Anytime MUon with Stable gradient Evaluation
AMUSE combines Muon with Schedule-Free-style gradient evaluation for stable anytime training without LR decay.
• Stronger 124M / 720M / 1B pretraining • Strong ImageNet / ViT fine-tuning performance.
@cloneofsimo fancy having him on your committee mustve felt good lol
paper name is clever too ngl
@cloneofsimo prof chulhee sounds like a solid pick
This implements the fixed-beta version which performs similarly to the hand-tuned-schedule variant.
The full implementation is:
_raw_sf = C.update_by_schedule_free.fn.fn
def _amuse_beta(group): step = group.get('_group_step') if step is None: return group['amuse_beta1'] step, w = max(int(step), 1), group['warmup_steps'] if step <= w or w <= 1: return group['amuse_beta1'] return 1 - ((w - 1) / (step - 1)) ** group['amuse_rho'] * (1 - group['amuse_beta1'])
@C.zero_guard('momentum') @C.no_state def muon_ema(group, update, grad, param, momentum): return utils.nesterov_ema(momentum, update, group['muon_mu'])
@C.copy_guard(2, 'z') @C.no_state def amuse_sf(group, update, grad, param, z): group['beta'] = _amuse_beta(group) return _raw_sf(group, update, grad, param, z)
class AMUSE(C.ScheduleFree): def __init__(self, params, **kw): d = dict(lr=.02, beta=.6, muon_mu=.95, weight_decay=0, warmup_steps=0, weight_lr_power=2., r=0., amuse_beta1=.6, amuse_rho=.8) d.update(kw) super().__init__(params, d, fns=(muon_ema, C.orthogonalize_update, amuse_sf))
@cloneofsimo looking at the 50k step chart, most of these converge to basically the same loss by step ~40k. the real difference is the first 10-15k steps where they separate. so optimizer choice seems more about training speed to a given loss than final loss itself
Lucas Nestler released a code implementation called SFMuon.
Cool + Prof. Chulhee was my defense's committee and he is really kind
🚨New Optimizer Paper AMUSE: Anytime MUon with Stable gradient Evaluation
AMUSE combines Muon with Schedule-Free-style gradient evaluation for stable anytime training without LR decay.
• Stronger 124M / 720M / 1B pretraining • Strong ImageNet / ViT fine-tuning performance.
