pytorch - 💡(How to fix) Fix [RFC] AdamTR: Adam variant for deterministic Token-Routed architectures [1 participants]

🚀 The feature, motivation and pitch

Motivation Standard AdamW treats all parameters identically, which is suboptimal for Mixture-of-Experts architectures where routed expert parameters, shared expert parameters, and dense (attention/embedding) parameters have fundamentally different training dynamics.

In deterministic Token-Routed MoE (under review at TMLR), we observe a loss gap vs dense baselines that converges from +0.31 to +0.09 with AdamW but plateaus. We hypothesize that an MoE-aware optimizer can close this gap further.

Proposal: AdamTR (Adam Token-Routed) AdamTR extends AdamW with four features:

Per-expert learning rate scaling — experts that see fewer tokens (due to Zipf routing) get higher LR for faster specialization Expert-aware weight decay — lighter decay for routed experts to preserve learned specialization, standard decay for shared/dense params Gradient normalization by token count — prevents high-frequency experts from dominating updates by normalizing gradients by relative token load Separate momentum per expert group — natural with [E, H, I] parameter tensors, prevents cross-expert momentum contamination Implementation Draft implementation: https://github.com/Complexity-ML/complexity-framework/blob/main/complexity/training/adam_tr.py

Core algorithm is identical to AdamW — the four features are applied as pre-processing before the standard Adam update step. ~230 lines, no new dependencies.

Questions for the community Would this be appropriate as a PR to torch.optim? What benchmarks would be expected? (We plan to run 384M MoE vs dense, 8B tokens) Should this target PyTorch core or a separate package like torch-optimizer? References Architecture paper: https://openreview.net/forum?id=jZq6EVboC6 (TMLR, under review) Framework: https://github.com/Complexity-ML/complexity-framework

Alternatives

• Using separate parameter groups in AdamW with different LR/weight_decay per group. This is verbose, error-prone, and doesn't support gradient scaling by token distribution.
• No existing optimizer in torch.optim addresses the asymmetry between expert and dense parameters in MoE architectures.

Additional context

Training results at 384M scale (8B tokens, 4 experts, iso-params comparison):

• Loss gap with AdamW: +0.31 (step 400) → +0.09 (step 5000+), plateaus
• MoE achieves 8,078 tok/s inference (3x faster than dense) but trails in loss
• Architecture: deterministic routing (no learned router), CUDA graph compatible
• PyTorch 2.11, FSDP full_shard, bf16, 2× RTX PRO 6000