AF-Muon: An AdamW-Free Muon Optimizer for Tied-Embedding Models
Muon trains hidden matrices with a spectral-norm steepest-descent update but falls back to AdamW for embeddings, output heads, and auxiliary parameters. This paper focuses on the tied vocabulary table—one tensor aliased between sparse token lookup and dense softmax classification—and proposes AF-Muon, an AdamW-free Muon variant that assigns a finite-cap clipped-row linear-minimization oracle to the tied table, an RMS-normalized update to one-dimensional parameters, and the standard spectral update to hidden matrices. By eliminating all second-moment state, AF-Muon cuts optimizer-state memory by roughly 20% and improves mean validation loss across nine tied-token settings spanning text, image, and protein modeling.
The paper presents a cogent, geometry-aware solution to a real gap in the Muon recipe. Its strongest asset is a sharp theoretical formulation of the tied-table update as an aliased multi-role problem, backed by an exact LMO derivation and a broad empirical sweep. Yet its scope is deliberately narrow—tied tables only—and its practical guarantees are explicitly qualified, with the authors noting that 'Consequently, we make no claim of unconditional convergence for this stochastic setting' (Section 2.3). The results are robust for the evaluated 124M–1B scale, but do not establish behavior at frontier pretraining scale.
The finite-cap LMO for tied tables is derived cleanly: Proposition 1 gives an exact, support-aware oracle that interpolates between Sign and row-RMS, and the alias-pullback framework formally justifies intersecting row and entry constraints. Empirically, the ablation design is commendable; the 2×2 factorial test shows that replacing both the tied-table AdamW update and the 1-D AdamW update jointly improves more than either isolated change, indicating that 'the joint improvement is larger than the sum of the isolated effects' (Section 3.4). The authors also note that their main configuration '$c=3$, $s=0.5$, which for all evaluated widths ($d>9$) lies strictly between Sign and row-RMS' (Section 2.2).
Three issues limit the strength of the claims. First, the proposed geometry is a modeling choice, not a derived optimum: the authors candidly admit that 'Weight tying does not uniquely imply the intersection of max-row-$\ell_2$ and max-entry constraints' (Section 4), leaving open whether another feasible set would perform better. Second, theoretical guarantees are conditional at best; aside from idealized deterministic bounds in the appendices, the main text notes that 'the practical implementation relies on momentum and approximate oracle maps (including finite-step Newton–Schulz and numerical bisection)' (Section 2.3), so convergence is not assured.
Third, the empirical ceiling is modest—1B parameters and roughly 10B tokens—meaning the results do not speak to trillion-parameter scaling dynamics where optimizer interactions may differ.
The evidence strongly supports superiority over the two main baselines within the Muon family. Across all nine settings, AF-Muon achieves the lowest mean validation loss against Hybrid Muon and the SCION-style Sign endpoint under a carefully matched protocol; the authors stress that 'All rows compare matched tied-table optimizer arms under the same model, tokenizer or token cache, data order, token budget, batching, precision, evaluation shard, and matrix-update protocol' (Table 1 caption). A weakness is that only the protein setting includes a full AdamW baseline; the others compare only Muon-family variants, leaving open how AF-Muon would fare against heavily tuned AdamW at equal compute. Still, in that benchmark the text reports that 'At 458.8M training tokens, AF-Muon achieves the best metrics, and all Muon-family arms substantially outperform tuned AdamW' (Section 3.3).
Reproduction is substantially aided by the public codebase and explicit Algorithm 1, but several practical hurdles remain. The tied-table oracle requires per-row bisection in FP32 for intermediate caps, implemented 'in row-preserving chunks so that chunking never changes the computed update' (Section 2.3), which introduces a custom, precision-sensitive kernel not found in standard optimizers. The Newton–Schulz spectral approximation and the approximate LMO mean exact bitwise reproduction depends on implementation subtleties. Furthermore, some key sweeps—batch size and auxiliary learning-rate sensitivity—are reported from a single seed and presented as 'sensitivity evidence rather than multi-seed uncertainty estimates' (Section 3.4), which weakens confidence in the claimed robustness.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.