AF-Muon: An AdamW-Free Muon Optimizer for Tied-Embedding Models
AuthorsArash Lagzian, Paniz Halvachi, Junming Zhang, Zhouhan Lin, Dianbo Liu
AffiliationsIndependent Researcher · LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University, Shanghai, China*Corresponding authors: arashlagzian94@nus.edu.sg, dianbo@nus.edu.sg
AF-Muon adapts Muon to tied embeddings, reducing optimizer memory while improving training across language, image, and protein models.
Key results
Experiments span language, image-token, protein, encoder–decoder, and sparse-MoE settings.
AF-Muon’s main configuration uses c=3 with tied-update scale s=0.5.
Approximate reduction relative to Hybrid Muon; measured savings range from 17.6% to 22.2%.
Approximate AF-Muon overhead relative to Hybrid Muon.
AF-Muon score; SCION-style Sign scores 33.87 in the matched comparison.
What the paper found
AF-Muon extends Muon without AdamW by assigning updates according to parameter role: hidden matrices retain Muon’s spectral update, one-dimensional parameters use RMS normalization, and a tied vocabulary table gets a support-aware clipped-row linear minimization oracle. That table receives both sparse token-lookup gradients and dense output-classifier gradients, so the method constrains updates with the intersection of row-wise ℓ2 and coordinate-wise max-entry bounds; its cap interpolates between Sign and row-RMS. Across 9 tied-token settings on text, ImageNet-32, and UniRef50, including Qwen2.5-0.5B, Meta’s Llama-3.2-1B, T5Gemma2-style models, and sparse mixture-of-experts, AF-Muon improves mean validation results over Hybrid Muon and a SCION-style Sign baseline. On Llama-3.2-1B, validation perplexity is 33.30 for AF-Muon versus 33.87 for Sign. Ablations support the finite cap—not simply removing AdamW—as a source of gains, while the optimizer uses about 20% less optimizer-state memory and adds about 1% step-time overhead relative to Hybrid Muon. The approach is relevant to the broader MoE direction represented by DeepSeek-V3, though that model is not evaluated here.
Original abstract
Muon improves large-scale training by applying a spectral-norm steepest-descent update to matrix parameters, but practical models also contain parameter blocks that do not fit dense-matrix geometry. One important case is the tied vocabulary table, which appears in language models and other token generators and can receive multiple structurally different gradient sources, from sparse input lookups to dense output-classifier updates. In the reference recipe these blocks are handed to an auxiliary AdamW optimizer, which restores second-moment state and updates the aliased table as a generic tensor. We propose AF-Muon, an AdamW-free extension of Muon that keeps the Muon matrix update for hidden weight matrices while using a support-aware finite-cap linear minimization oracle for tied vocabulary tables and an RMS-normalized update for one-dimensional auxiliary parameters. AF-Muon therefore trains every parameter class with a single first-moment buffer and no second-moment state, saving around 20% optimizer-state memory relative to Hybrid Muon in our benchmark. Across nine tied-token settings - decoder-only language models from 124M to 1B parameters, a fully shared T5-style encoder-decoder, and ImageGPT-style image-token, protein, and sparse-MoE variants, spanning text, image, and protein-sequence data - AF-Muon improves mean validation loss and perplexity over both Hybrid Muon and a SCION-style Sign endpoint. Long-horizon runs and hyperparameter sensitivity studies confirm the gain is robust, and identical-momentum diagnostics attribute it to the finite cap, which preserves more within-row magnitude than Sign while bounding the coordinate concentration of row-RMS. These results identify tied vocabulary tables as a distinct optimizer geometry and yield a robust AdamW-free Muon variant across models, modalities, and architectures, with about 1% step-time overhead in matched training.
Read the original paperMore in Optimization
Browse all 39 papers →A provable quantum advantage for approximate optimization via decoded quantum interferometry
Maximilian J. Kramer, Elies Gil-Fuster, Benjamin D. M. Jones, Jens Eisert, Franz J. Schreiber
This paper proves that, for a specially constructed optimization problem, quantum interferometry can achieve approximation scores that classical algorithms cannot match efficiently.
Planning to Learn
Ian Osband
This paper makes classification more forward-looking by designing a loss that values not just immediate accuracy gains, but how much future learning an update enables.
An $Ω(κ_y^8ε^{-6})$ Lower Bound for Stochastic NC-SC Bilevel Optimization with First-order Oracles
Zhihao Gu, Qilong Wu, Junchi Yang
This work proves that stochastic bilevel optimization fundamentally requires up to epsilon^{-6} oracle queries, showing existing methods are asymptotically optimal.