PaperlayerView exact revision on arXiv

Source-linked AI summary

AF-Muon: An AdamW-Free Muon Optimizer for Tied-Embedding Models

Arash Lagzian, Paniz Halvachi, Junming Zhang, Zhouhan Lin, Dianbo Liu

arXiv:2610.01395v1cs.LG

TL;DR

Tied vocabulary tables merge structurally different lookup and classifier gradients, yet hybrid Muon leaves embeddings and output blocks on an auxiliary AdamW branch. AF-Muon assigns first-moment updates by parameter role, including a finite-cap LMO for tied tables. Across nine settings it improves validation results over Hybrid Muon and a SCION-style Sign endpoint.

  • Problem

    Tied vocabulary tables merge sparse lookup and dense classifier gradients, while the reference Muon recipe assigns embeddings and output blocks to an auxiliary AdamW branch.

    Source: §1 · pp. 1–2

  • Method

    AF-Muon keeps Muon for hidden matrices, uses RMS-normalized updates for one-dimensional parameters, and applies a support-aware finite-cap LMO to the tied table.

    Source: §1 · pp. 2–3

  • Results

    Across nine tied-token settings, AF-Muon achieves the lowest mean validation loss in every setting, outperforming Hybrid Muon and the SCION-style Sign endpoint.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    The results support treating functional aliasing, rather than stored tensor shape alone, as a first-class consideration in optimizer design.

    Source: §5 · p. 10

    The results establish robustness under matched fixed-token protocols, but do not tune every optimizer family or test frontier-scale pretraining.

    Source: §4 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

Muon’s hybrid recipe leaves embeddings and output blocks on AdamW, while tied vocabulary tables merge sparse lookup and dense classifier gradients on one tensor. AF-Muon assigns role-specific first-moment updates to these parameter classes and is evaluated across nine tied-token settings.

1 INTRODUCTION

  • The tied vocabulary table is an aliased multi-role parameter: sparse lookup and dense classifier gradients merge on the same stored tensor.Hybrid Muon instead assigns the table to AdamW, retaining coordinate-wise second-moment state.Source: §1 · pp. 1–3
  • AF-Muon keeps Muon for hidden matrices, uses a finite-cap LMO for tied tables, and applies RMS-normalized updates to one-dimensional parameters.The cap interpolates between Sign and row-RMS; its geometry controls direction while a separate scale controls update magnitude.Source: §1 · pp. 2–3
  • Across nine tied-token settings, AF-Muon achieves the lowest mean validation loss in every setting, outperforming Hybrid Muon and the SCION-style Sign endpoint.The evaluation includes decoder-only, fully shared encoder–decoder, image-token, protein-language, and sparse-MoE models.Source: §1 · pp. 2–3
  • The evaluation spans 124M–1B decoder-only models plus encoder–decoder, image-token, protein-sequence, and sparse-MoE variants, with three paired seeds and runs up to ∼10B tokens.The studies include tied-geometry sensitivity tests.Source: §1 · p. 3

2 METHOD

AF-Muon assigns first-moment updates by parameter role, using Muon for hidden matrices, RMS normalization for vectors, and a finite-cap LMO for tied vocabulary tables. Its table compares optimizer outcomes across tied-token models, while the method also specifies the oracle geometry and implementation scope.

2.2 PARAMETER CLASSES

  • AF-Muon retains Muon’s spectral LMO for hidden matrices and uses RMS-LMO updates for one-dimensional parameters, replacing the auxiliary AdamW branch with role-matched first-moment updates.The hidden-matrix spectral LMO is approximated by a K-step Newton–Schulz iteration.Source: §2 · p. 3, §2.2 · pp. 3–4

2.1 STEEPEST DESCENT UNDER A NORM

  • For norm-based steepest descent, the LMO maps its oracle input to a direction that maximizes linearized decrease under the selected norm.The tied table and auxiliary tensors use the unnormalized momentum buffer as input; hidden matrices use a Nesterov-style construction before the Newton–Schulz map.Source: §2.1 · p. 3

2.2 PARAMETER CLASSES

  • The tied table’s proposed feasible set intersects a max-row-ℓ2 constraint with a max-entry cap, reflecting its aliased lookup and output-classifier roles.The paper presents this intersection as a modeling choice, not a geometry implied by weight tying alone.Source: §2.2 · p. 4
  • The clipped-row LMO interpolates from Sign at c = 1 to row-RMS when the cap is removed, with c controlling direction shape and s controlling row magnitude.For dense rows, the finite-cap update bounds peak-to-average ratio by c^2, while its optimal linear objective is nondecreasing in c.Source: §2.2 · p. 4
  • Sharing identical spectral constraints across depths collapses to one spectral LMO on the accumulated gradient, unlike the tied table’s heterogeneous role constraints.In a four-depth shared-MLP diagnostic, combining per-depth Muon directions retained 74.1% of the linear objective for up-projection and 83.9% for down-projection.Source: §2.2 · p. 5

2.3 FULL OPTIMIZER

  • The optimizer is stated for tied-head RoPE decoders; untied heads and learned positional tables are outside the tied-table claim.The practical implementation uses momentum and approximate oracle maps, and the authors make no claim of unconditional convergence for this stochastic setting.Source: §2.3 · p. 5
  • Table 1 reports that AF-Muon achieves the best mean validation loss and PPL/BPD across all settings.Rows compare matched tied-table optimizer arms; ImageGPT reports validation bits/dim.Source: Table 1 · p. 6

3 EXPERIMENTS

Across matched tied-token experiments, AF-Muon achieves the best mean validation loss and perplexity across the evaluated settings. It transfers beyond decoder-only language modeling, while ablations and sensitivity studies examine its role-specific updates and robustness.

3.1 EXPERIMENTAL SETUP

  • Experiments compare nine tied-token settings using three paired seeds under matched model, data, token-budget, and training protocols.The settings span decoder-only language models, shared-vocabulary encoder–decoder models, sparse MoE, image-token modeling, and protein language modeling.Source: §3.1 · p. 6

3.2 MAIN DECODER-ONLY LANGUAGE MODELING RESULTS

  • Across all evaluated settings, AF-Muon achieves the best mean validation loss and PPL/BPD compared with Hybrid Muon.Comparisons use matched optimizer protocols; paired three-seed tests against the strongest non-AF baseline are reported separately.Source: §3.2 · p. 7
  • AF-Muon has the best validation loss and perplexity across paired seeds in NanoGPT, SmolLM2-135M, Qwen2.5-0.5B, and Llama-3.2-1B.The advantage persists in a Llama-3.2-1B continuation to 9.98B tokens.Source: §3.2 · p. 7

3.3 GENERALIZATION BEYOND DENSE DECODER-ONLY LANGUAGE MODELING

  • AF-Muon achieves the best validation loss in the encoder–decoder, sparse-MoE, image-token, and protein settings.The four non-decoder-only settings include a fully shared T5-style vocabulary, ImageGPT-style ImageNet-32, and protein modeling; Figure 3 reports the lowest final validation loss in all four panels.Source: Figure 3 · p. 8, §3.3 · p. 8

3.4 ABLATIONS AND ROBUSTNESS

  • The factorial ablation finds that both the tied-table finite-cap LMO and one-dimensional RMS-LMO contribute, with a joint improvement larger than their isolated effects combined.The dense output gradient dominates the tied-table gradient, but adding the sparse input component still improves validation; the two components are nearly orthogonal (cosine 0.003).Source: §3.4 · p. 8, Figure 4 · p. 9
  • At fixed c = 3, every tested scale s ∈ {0.125, 0.5} beats the Sign arm, while s = 0.5 losses span only 0.0069 across the tested cap range.The c = 3 configuration caps 0.29% of coordinates, affects 90.3% of rows, and attains 1.253× the same-state Sign linear objective; it is not presented as uniquely optimal.Source: §3.4 · pp. 8–9
  • AF-Muon retains its advantage across tested vocabularies, hidden widths, batch sizes, and auxiliary learning rates.Paired improvements over Sign are 0.036, 0.045, and 0.052 for 50k, 100k, and 150k vocabularies, and 0.036, 0.040, and 0.032 for widths 768, 1024, and 1536; batch-size and learning-rate sweeps are single-seed sensitivity evidence.Source: §3.4 · p. 9, Figure 5 · p. 9

4 LIMITATIONS

The proposed tied-table feasible set is one modeling choice, not a uniquely implied consequence of weight tying.

4 LIMITATIONS

  • The results establish robustness under matched fixed-token protocols, but do not test frontier-scale pretraining or show that the proposed feasible set is optimal among role-aware geometries.The default c = 3, s = 0.5 is described as robust rather than universally optimal.Source: §4 · p. 9

5 CONCLUSION

The conclusion summarizes AF-Muon’s role-aware, support-aware updates for tied vocabulary tables alongside Muon updates for hidden matrices and RMS updates for vectors, reporting improved fixed-budget validation across nine tied-token settings without second-moment state. It also characterizes the alias-geometry dual through an aggregate constraint and quotient dual, whose minimizing components are geometric decompositions rather than necessarily role-specific gradients.

5 CONCLUSION

  • AF-Muon keeps Muon’s spectral update for hidden matrices, uses RMS updates for one-dimensional parameters, and applies an exact support-aware LMO to the tied table, eliminating second-moment state.The tied-table oracle uses an intersection of max-row-ℓ2 and max-entry constraints.Source: §5 · p. 9, A RELATED WORK · p. 17
  • Across nine tied-token settings and matched three-seed protocols, AF-Muon improves fixed-budget validation over Muon and SCION-style Sign, cutting optimizer memory about 20% at roughly 1% training-time overhead.Settings span 124M–1B decoder-only models, a fully shared encoder–decoder, image-token, protein, and sparse-MoE variants across text, image, and protein data.Source: §5 · pp. 9–10

A RELATED WORK

  • AF-Muon’s tied-table contribution is an exact LMO for intersected max-row-ℓ2 and max-entry balls, containing SCION Sign and row-RMS as endpoints and exposing finite-cap directions between them.The paper distinguishes this joint oracle from prior role-specific vocabulary updates and does not claim novelty for Sign or row-RMS.Source: A RELATED WORK · p. 16
  • Tied tables combine sparse lookup gradients with generally dense output-softmax gradients, motivating an update geometry for the shared tensor rather than treating it as only an input table or output classifier (Press & Wolf (2017); Inan et al. (2016)).AF-Muon changes the LMO geometry of the merged gradient rather than reweighting its lookup and softmax components.Source: A RELATED WORK · p. 16

B SETTING, NOTATION, AND STANDARD LMO GEOMETRY

  • Hard tying requires one stored-table update, and imposing both lookup-side row and output-side coordinate constraints is a modeling choice, not a consequence uniquely determined by weight tying.The construction applies both role constraints to the same parameter object.Source: B SETTING, NOTATION, AND STANDARD LMO… · p. 17

B.2 DUAL NORMS AND LINEAR MINIMIZATION ORACLES

  • An LMO’s alignment identity applies to its input; when that input is momentum rather than the current gradient, it does not itself establish descent of the current objective.The actual minimization step is the negative learning-rate-scaled aligned LMO direction.Source: §B.2 · p. 18

C.1 WHERE DOES THE TERMINOLOGY COME FROM?

  • Alias-pullback geometry transfers each functional role’s norm constraint to the shared parameter, making its feasible set the intersection of role-wise sets and its dual an optimal decomposition into latent role covectors.Redundancy of the shared parameter must also be compatible with every role.Source: §C.1 · p. 19

C.2 PULLBACK NORM AND INTERSECTION BALL

  • Under a common-kernel condition, the alias seminorm is a norm whose unit ball is the intersection of the role-wise feasible sets; its dual is an attained minimum-cost decomposition through role covectors.The dual formula is exact under the stated common-interior qualification.Source: §C.2 · p. 20, §C.3 · p. 20

C.3 DUAL OF THE ALIAS GEOMETRY

  • The aggregate constraint follows by substitution, and the feasible infimum is attained because the norm is coercive on the closed affine feasible set.If the constraint is infeasible, both sides equal +∞; coercivity on the product space also ensures attainment.Source: §C.3 · p. 20
  • The dual support function is finite only for B in K⊥ = range(A*), where it defines a quotient dual rather than a finite norm on all of W.Source: §C.3 · p. 20
  • The minimizing Zr form an optimal geometric decomposition of a merged gradient or momentum, need not equal role-specific gradients, and need not be formed or stored by the oracle.Source: §C.3 · p. 20

C.4 REDUNDANCY UNDER ALIASING

This section characterizes redundancy when roles share one stored parameter, then applies that framework to tied-table gradients and constraints. It distinguishes role-wise redundancy from full-network invariance and motivates the tied intersection norm.

C.4 REDUNDANCY UNDER ALIASING

  • Simultaneous role-wise redundancies are exactly the intersection of the role-wise redundant perturbation sets, but this need not characterize every invariance of the composed network.Downstream maps may discard changes or effects across paths may cancel.Source: §C.4 · p. 21
  • An oracle on the projected quotient ball can differ from one restricted to the section, which may be more restrictive unless projection is nonexpansive in the alias norm.The section and projection are not generally the same set.Source: §C.4 · p. 21

D.1 GRADIENT STRUCTURE OF THE TIED TABLE

  • The lookup contribution is row-sparse over observed token indices, while the output-softmax contribution is generally dense; their gradients add on the shared table.Full-softmax support is not guaranteed for every merged-gradient coordinate because cancellation, masking, alternative losses, or numerical underflow can alter support.Source: §D.1 · pp. 21–22

D.2 WHY THE UNTIED-HEAD QUOTIENT DISAPPEARS

  • A head-shift direction can be redundant for the output role alone but not for both tied roles, so projecting the merged gradient may discard a legitimate lookup contribution.The tied update therefore uses the merged, unprojected buffer.Source: §D.2 · p. 22

D.3 THE TIED INTERSECTION NORM

  • The tied intersection norm bounds row energy and coordinate magnitude; the cap is tighter for concentrated hidden states, and neither constraint implies the other for 1 < c < d.The cap bounds a coordinate that could otherwise dominate a logit direction.Source: §D.3 · p. 23

E EXACT SUPPORT-AWARE CLIPPED-ROW ORACLE

The section constructs an exact support-aware clipped-row oracle for tied-table updates from one merged buffer. It relates this finite-cap update to Sign and row-RMS endpoints and characterizes its dual geometry and fixed-input objective retention.

E EXACT SUPPORT-AWARE CLIPPED-ROW ORACLE

  • The exact clipped-row LMO maximizes each row objective with a unique support-preserving solution, while zero-input coordinates receive no update.AF-Muon applies it to the merged buffer without splitting lookup and head contributions.Source: E EXACT SUPPORT-AWARE CLIPPED-ROW ORACLE · pp. 24–25
  • At c = 1 the oracle gives Sign, while for c ≥ √d it gives row-RMS, with finite-cap directions between these endpoints.The cap and row-radius constraints define the intersection geometry.Source: E EXACT SUPPORT-AWARE CLIPPED-ROW ORACLE · p. 25

F.1 THE TIED DUAL NORM AS AN INFIMAL CONVOLUTION √

  • The tied dual norm prices the merged momentum by its cheapest split into a row-structured component and an entrywise-spiky component, though AF-Muon does not compute that split.The primal clipped-row oracle reaches the same optimal value using one merged buffer.Source: §F.1 · p. 25

F.2 EXPLICIT ROW SUPPORT FUNCTION AND OBJECTIVE RETENTION

  • For a fixed input, same-input objective retention relative to row-RMS is nondecreasing with cap size; this geometric comparison does not predict future training loss.The support function evaluates the oracle’s optimized objective.Source: §F.2 · p. 26

F.3 SUPPORT-AWARE CONCENTRATION (PAPR) BOUNDS

The finite cap controls update-row coordinate concentration, but its PAPR guarantee depends on the row’s support. PAPR is scale-invariant, and uniform scaling does not equalize energy across rows with different supports.

F.3 SUPPORT-AWARE CONCENTRATION (PAPR) BOUNDS

  • PAPR is scale-invariant, lies in [1, d], and is at least d/k for a row with k nonzero entries.Row-RMS can reach PAPR d for a one-coordinate input and approach d for dense inputs dominated by one coordinate.Source: §F.3 · p. 27
  • When c^2k(b) ≥ d, the clipped-row update has PAPR at most c^2; the bound does not hold generally for support-limited rows.For d = 16, c = 2, and b = e1, the update is 2e1 with PAPR 16, not at most 4.Source: §F.3 · p. 27
  • Uniformly scaling the completed direction preserves its radius-to-cap ratio and PAPR, but does not force equal energy across rows with different supports.For support-sufficient rows the scaled energy is s^2d; for support-limited rows it is s^2c^2k(b).Source: §F.3 · p. 27
  • The width factor and radius ratio are design choices motivated by norm-based scaling and evaluated empirically, not consequences of the alias-pullback theorem.The theorem supplies no width-independent convergence constants.Source: §F.3 · p. 27

G SYMMETRY AND ITS LIMITATION

The finite-cap oracle preserves signed-permutation equivariance but, in its active-cap range, not general feature-rotation equivariance; full orthogonal equivariance instead uniquely selects row-RMS under the stated row-separable conditions. The appendix also gives conditional geometry and descent results for idealized variants, not a convergence guarantee for practical stochastic training.

G SYMMETRY AND ITS LIMITATION

  • In the active-cap range, the oracle is equivariant to vocabulary permutations and signed feature permutations, but not to general right rotations.When the cap is redundant, it equals row-RMS and full right-O(d) equivariance holds.Source: G SYMMETRY AND ITS LIMITATION · p. 28

Then necessarily

  • For row-separable rules, full orthogonal equivariance, fixed energy, and positive alignment uniquely select row-RMS.This result does not characterize coupled geometries or imply that signed permutations are maximal for every concentration-control method.Source: Then necessarily · pp. 28–29

H EXACT ORACLE PROPERTIES AND CONDITIONAL CONVERGENCE BOUNDS

  • The convergence analysis covers idealized exact and inexact oracle variants, not practical stochastic training without additional error-control assumptions.Unbiased stochastic gradients alone are insufficient, and the bounds do not claim dimension-independent constants.Source: H EXACT ORACLE PROPERTIES AND… · p. 29

H.1 PRODUCT GEOMETRY

  • With a shared scalar schedule, idealized role-specific AF-Muon is one linear minimization oracle over a weighted product norm; varying relative branch rates requires time-varying geometry analysis.The block radii absorb branch learning-rate ratios and shape factors.Source: §H.1 · p. 29

H.2 EXACT DETERMINISTIC DESCENT

  • The exact deterministic descent bounds rely on stated smoothness conditions and do not by themselves establish hyperparameter transfer across model widths or vocabulary sizes.The analyzed schedule differs from practical AF-Muon.Source: §H.2 · p. 30

H.3 MASTER INEXACT-ORACLE THEOREM

The appendix bounds progress for inexact oracle updates by accounting for estimator and oracle errors, and relates decoupled decay to conditional-gradient steps. A separate joint result gives conditional stationarity for exact deterministic hybrid updates, without guaranteeing convergence of practical stochastic pretraining.

H.3 MASTER INEXACT-ORACLE THEOREM

  • The master inexact-oracle bound accounts for gradient-estimation error, oracle objective error, and curvature under smoothness and feasibility assumptions.With feasibility inflation κt, the error coefficient and curvature term increase; practical Newton–Schulz errors require implementation-specific bounds.Source: §H.3 · pp. 30–31

H.4 DECOUPLED DECAY AS CONDITIONAL GRADIENT

  • For feasible iterates, decoupled decay is a Frank–Wolfe step; convexity preserves feasibility, while attraction to the radius bound requires nonsummable step sizes.Outside the constraint set, the same algebraic update is not a feasible Frank–Wolfe iterate.Source: §H.4 · p. 31

H.5 HYBRID STATIONARITY

  • The hybrid stationarity result applies to exact deterministic updates under joint smoothness assumptions, not practical momentum, numerical approximations, or stochastic pretraining with decay.It does not interpret the update as gradient descent on a regularized objective.Source: §H.5 · p. 32

I EXTENSION TO ENCODER-DECODER ARCHITECTURES (T5)

In fully shared T5-style models, the encoder and decoder lookup roles are geometrically redundant, while their gradients combine into a sparse lookup contribution alongside a dense output contribution. Across the included tied-vocabulary decoder-only experiments, AF-Muon improves validation performance, with Qwen2.5-0.5B retaining its ordering in a long-horizon check.

I EXTENSION TO ENCODER-DECODER ARCHITECTURES (T5)

  • In fully shared T5-style models, encoder and decoder lookup roles are geometrically redundant, so the three-way alias norm collapses exactly to the two-way tied intersection norm.The added lookup role does not change the oracle, dual, or symmetry results from the two-way setting.Source: I EXTENSION TO ENCODER-DECODER… · p. 33

I.1 GRADIENT STRUCTURE AND CONVERGENCE

  • Combining encoder and decoder lookup gradients preserves row sparsity on their occupied-token union, while the output-classification gradient remains generically dense across vocabulary rows.The convergence guarantees transfer under standard block reindexing.Source: §I.1 · pp. 33–34

J.1 NANOGPT

  • On NanoGPT, AF-Muon improves validation loss while keeping the shared table smaller in RMS norm than the SCION-style Sign endpoint.Hybrid Muon’s AdamW weight decay makes its table smaller still, but that additional shrinkage does not close the performance gap.Source: §J.1 · pp. 34–35

J.2 SMOLLM2-135M

  • On SmolLM2-135M, AF-Muon improves validation loss while keeping the shared table smaller in RMS norm than the SCION-style Sign endpoint.Hybrid Muon’s AdamW weight decay makes its table smaller still, but that additional shrinkage does not close the performance gap.Source: §J.2 · pp. 35–36

J.3 QWEN2.5-0.5B

  • At the near-1B evaluation point, AF-Muon improves validation loss by 0.0562 ± 0.0016 points over Hybrid Muon and 0.0278 ± 0.0022 over SCION-style Sign on Qwen2.5-0.5B.AF-Muon also has the lowest validation loss on all three paired seeds.Source: §J.3 · p. 36

J.3.1 QWEN2.5-0.5B LONG-HORIZON CONTINUATION

  • At 6.29B tokens in the seed-43 Qwen2.5-0.5B continuation, AF-Muon improves validation loss by 0.0316 points over Hybrid Muon and 0.0138 over SCION-style Sign.This single-seed continuation checks stability of the paired short-horizon result, not an independent multi-seed estimate.Source: §J.3.1 · p. 37

J.4 LLAMA-3.2-1B DECODER-ONLY RESULTS

In the Llama-3.2-1B tied-embedding experiments, AF-Muon has the best validation loss in short-horizon diagnostics and remains ahead of Hybrid Muon in the long-horizon continuation. The continuation is a single-seed stability check, not a separate multi-seed estimate.

J.4 LLAMA-3.2-1B DECODER-ONLY RESULTS

  • AF-Muon has the best validation loss in the short-horizon Llama-3.2-1B checkpoint diagnostics, with a tied-table RMS norm between Hybrid Muon and SCION-style Sign.Hybrid Muon has the smallest table because its auxiliary AdamW branch applies weight decay; Sign has the largest RMS norm.Source: §J.4 · pp. 39–40

J.4.1 LLAMA-3.2-1B LONG-HORIZON CONTINUATION

  • The long-horizon seed-43 run extends training to 9.98B tokens to check whether the short-horizon optimizer ordering persists beyond the three-seed 2B-token budget.It is not treated as an independent three-seed estimate.Source: §J.4.1 · p. 40
  • At 9.98B tokens, AF-Muon has the lowest final validation loss, with SCION-style Sign close behind and both AdamW-free tied-table arms ahead of Hybrid Muon.The AF-Muon advantage over Sign is small, and this seed-43 continuation checks stability rather than providing a separate multi-seed estimate.Source: §J.4.1 · p. 40

J.5 SHARED-VOCABULARY ENCODER–DECODER EXPERIMENTS

Across shared-vocabulary encoder–decoder, sparse-MoE, image-token, and protein experiments, AF-Muon generally achieves the strongest reported validation outcomes, though the small T5 comparison against Sign is statistically insignificant. Timings and ablations show modest overhead and contributions from both tied-table and one-dimensional updates.

J.5 SHARED-VOCABULARY ENCODER–DECODER EXPERIMENTS

  • In the fully shared small T5 setting, AF-Muon has the best mean validation loss among tied methods, but its margin over SCION-style Sign is small and statistically insignificant.The untied Hybrid Muon control increases model size from 54.6M to 104.9M parameters and is not parameter matched.Source: §J.5 · p. 42
  • In the T5Gemma2-style experiment, AF-Muon obtains the lowest validation loss and perplexity in the completed evaluated seeds, supporting transfer to a larger encoder–decoder configuration.The experiments differ in architecture, vocabulary configuration, and training budget, so they do not isolate model or vocabulary size as the cause of the larger gap.Source: §J.5 · p. 44

J.6 NANOGPT-MOE

  • In tied-embedding NanoGPT-MoE, AF-Muon improves validation loss by 0.0224 ± 0.0023 paired points over Hybrid Muon and 0.0452 ± 0.0019 over SCION-style Sign.All methods keep all four experts active, with router entropy near one and low expert-load coefficient of variation.Source: §J.6 · p. 44

J.7 TIED IMAGEGPT-STYLE IMAGENET-32

  • On tied ImageGPT-style ImageNet-32, AF-Muon improves validation loss by 0.0103 ± 0.0014 and bits/dim by 0.0049 ± 0.0007 over Hybrid Muon.Against SCION-style Sign, the respective improvements are 0.0469 ± 0.0064 validation-loss points and 0.0226 ± 0.0031 bits/dim.Source: §J.7 · p. 45

J.8 PROTEIN LANGUAGE MODELING

  • On UniRef50 protein modeling, AF-Muon improves average validation loss by 0.1147 over Hybrid Muon and 0.0569 over SCION-style Sign across three seeds.AF-Muon has the lowest validation loss and perplexity for each seed.Source: §J.8 · p. 51

J.9 TRAINING OVERHEAD

  • AF-Muon reduces optimizer-state memory by roughly 18–22% relative to Hybrid Muon, with at most 1.3% median step-time overhead across models.Peak allocated-memory reduction is smaller because these large-batch runs are activation- and temporary-buffer dominated.Source: §J.9 · p. 51

K.1 NANOGPT FACTORIAL ABLATION: TIED TABLE VERSUS 1D AUXILIARY PARAMETERS

  • In the NanoGPT factorial ablation, replacing the tied-table update alone improves validation loss by 0.0133, replacing the one-dimensional update alone by 0.0075, and replacing both by 0.0413.The four arms use zero auxiliary weight decay, isolating update geometry from an extra shrinkage term.Source: §K.1 · p. 52

K.2 REAL VOCABULARY-SIZE RE-TOKENIZATION ABLATION

A real re-tokenization test on NanoGPT finds that AF-Muon’s validation-loss advantage over the SCION-style Sign endpoint does not shrink as vocabulary size grows from 50k to 150k.

K.2 REAL VOCABULARY-SIZE RE-TOKENIZATION ABLATION

  • The ablation uses separately trained ByteLevel BPE tokenizers on the same FineWeb shard, making this a real re-tokenization test rather than a padded-vocabulary test.Token IDs reach the top of each vocabulary in every setting.Source: §K.2 · p. 52
  • AF-Muon improves validation loss over the SCION-style Sign endpoint by 0.036 points at 50k, 0.045 at 100k, and 0.052 at 150k vocabulary size.Because tokenizers induce different per-token entropy scales, absolute losses should be compared only within each vocabulary.Source: §K.2 · p. 52

K.3 TIED-STEP MAGNITUDE ABLATION

The tied-scale sweep varies AF-Muon’s tied-table scale while holding the model and training protocol fixed. AF-Muon beats the SCION-style Sign endpoint across tested scales, with the best validation loss at s = 0.25.

K.3 TIED-STEP MAGNITUDE ABLATION

  • Only the AF-Muon tied-table scale s varies in the sweep; c stays fixed at 3.The model, vocabulary, data, precision recipe, matrix weight decay, and token budget are held fixed.Source: §K.3 · p. 54, Table 30 · p. 54
  • AF-Muon obtains lower validation loss and perplexity than the SCION-style Sign endpoint across tested tied scales.Table 30 fixes c = 3 for AF-Muon and uses c = 1, s = 1 for Sign under the same protocol.Source: §K.3 · p. 54, Table 30 · p. 54
  • The best validation loss occurs at s = 0.25, while s = 0.375 and s = 0.1875 give similar results.The tied-table RMS decreases as s is reduced.Source: §K.3 · p. 54

K.4 NANOGPT REAL-WIDTH ABLATION

These ablations examine AF-Muon across NanoGPT widths and training and optimizer settings, including tied-table gradient sources and parameter sharing. A dense-sharing diagnostic probes the boundary of its tied-table geometry in a homogeneous sharing setting.

K.4 NANOGPT REAL-WIDTH ABLATION

  • At d = 768, 1024, and 1536, AF-Muon minus Sign validation-loss gaps were −0.0363 ± 0.0063, −0.0397 ± 0.0023, and −0.0321 ± 0.0085, respectively.The ablation changes the actual model width and tied-row dimension; AF-Muon remains better at all three widths.Source: §K.4 · p. 55, Table 33 · p. 56

K.5 BATCH-SIZE SENSITIVITY

  • AF-Muon gives the lowest validation loss at every tested batch size, from 32k to 524k tokens per optimizer update.Its advantage over Sign is smallest at the smallest batch and clearer at intermediate and large batches, at a matched 458.8M-token comparison point.Source: Table 34 · p. 57, §K.5 · p. 57

K.6 AUXILIARY LEARNING-RATE SENSITIVITY

  • Hybrid Muon is sensitive to the auxiliary AdamW learning rate, while AF-Muon remains competitive across the extended sweep and has the best loss at most tested settings.The largest rates degrade all methods; this is a sensitivity diagnostic under literature-motivated defaults, not compute-optimal tuning.Source: §K.6 · pp. 57–58, Table 35 · p. 58

K.7 AF-MUON CAP AND SCALE SENSITIVITY

  • At 393.2M tokens, (c, s) = (3, 0.5) achieved 3.7815 ± 0.0026 mean validation loss, versus 3.7831 ± 0.0033 for scale-matched Sign and 3.7844 ± 0.0010 for scale-matched row-RMS.Reducing s from 1 to 0.5 lowered mean loss for every tested cap, but the margins were small and the sweep does not establish a universal finite-cap advantage. At the default, the cap affected 0.0028 ± 0.0001 of coordinates and 0.8949 ± 0.0042 of rows.Source: Table 36 · p. 58, §K.7 · pp. 58–59

K.8 SPARSE-INPUT AND DENSE-OUTPUT GRADIENT SOURCES

  • At 458.8M tokens, dense output-only reached 3.7634 validation loss, while the full sparse-plus-dense tied gradient reached 3.7395; sparse-only was substantially worse.The input/output gradient cosine was near zero, and the full gradient consistently improved over dense-only during training.Source: Table 38 · p. 60, Figure 19 · p. 60, §K.8 · p. 60

K.9 BOUNDARY CASE: HOMOGENEOUS DENSE-PARAMETER REUSE

  • Repeated use of a shared dense matrix does not by itself require a new LMO: because each use has the same spectral geometry, Muon on the accumulated gradient is the appropriate shared-parameter update.Use-specific gradients differed but all pairwise cosines were positive; after matching spectral radius, summing per-depth Muon directions retained only 74.08% and 83.85% of the aggregate Muon linear objective for up and down projections, respectively. This was a one-seed boundary diagnostic, not a complete-optimizer performance comparison.Source: Table 39 · p. 62, §K.9 · pp. 62–63
  • The dense-sharing negative control shows AF-Muon’s tied-table geometry is specialized to heterogeneous lookup/classifier aliasing, not generic parameter reuse.The relevant distinction is whether aliases impose heterogeneous update geometries.Source: §K.9 · p. 63

Related papers

Loading 2610.01395v1…