PaperlayerView exact revision on arXiv

Source-linked AI summary

Gluon: Making Muon & Scion Great Again! (Bridging Theory and Practice of LMO-based Optimizers for LLMs)

Artem Riabinin, Egor Shulgin, Kaja Gruntkowska, Peter Richtárik

arXiv:2505.13416v1cs.LGmath.OCstat.ML

TL;DR

Existing analyses of LMO-based optimizers miss their practical layer-wise updates and rely on classical smoothness assumptions that yield impractically small stepsizes. The paper introduces Gluon and layer-wise (L0, L1)-smoothness to analyze these methods faithfully. Its adaptive stepsizes closely match tuned values, while experiments support the assumption along training trajectories.

  • Problem

    Existing analyses fail to capture practical layer-wise updates and use classical smoothness assumptions that prescribe impractically small stepsizes.

  • Method

    The paper introduces Gluon, a layer-wise LMO-based optimizer, and a layer-wise (L0, L1)-smoothness framework that recovers Muon and Scion as special cases.

  • Results

    Gluon’s adaptive stepsizes closely match the tuned learning rates associated with state-of-the-art empirical performance.

  • Takeaways & Limitations

    The framework provides a theoretically predictive analysis of practical layer-wise optimizers and supports theory-driven stepsize selection.

  • Takeaways & Limitations

    The stochastic guarantees rely on bounded variance, and theoretical support for adaptive stepsizes is currently restricted to the deterministic setting.

Abstract

from arXiv · show

Recent developments in deep learning optimization have brought about radically new algorithms based on the Linear Minimization Oracle (LMO) framework, such as $\sf Muon$ and $\sf Scion$. After over a decade of $\sf Adam$'s dominance, these LMO-based methods are emerging as viable replacements, offering several practical advantages such as improved memory efficiency, better hyperparameter transferability, and most importantly, superior empirical performance on large-scale tasks, including LLM training. However, a significant gap remains between their practical use and our current theoretical understanding: prior analyses (1) overlook the layer-wise LMO application of these optimizers in practice, and (2) rely on an unrealistic smoothness assumption, leading to impractically small stepsizes. To address both, we propose a new LMO-based method called $\sf Gluon$, capturing prior theoretically analyzed methods as special cases, and introduce a new refined generalized smoothness model that captures the layer-wise geometry of neural networks, matches the layer-wise practical implementation of $\sf Muon$ and $\sf Scion$, and leads to convergence guarantees with strong practical predictive power. Unlike prior results, our theoretical stepsizes closely match the fine-tuned values reported by Pethick et al. (2025). Our experiments with NanoGPT and CNN confirm that our assumption holds along the optimization trajectory, ultimately closing the gap between theory and practice.

1 Introduction

LMO-based optimizers are emerging as practical alternatives to AdamW, but existing theory misses their layer-wise implementation and relies on an inaccurate smoothness model. Gluon addresses both mismatches through a layer-wise framework for deep-network optimization.

  • Gluon is formalized as an iterative layer-wise LMO-based optimizer with momentum and adaptively chosen per-layer stepsizes or radii.
  • LMO-based optimizers appear empirically stronger than AdamW in LLM training while offering simpler implementation and tuning.
  • Existing convergence analyses overlook layer-wise structure and use flawed stepsize choices stemming from an inaccurate smoothness model.
  • Gluon studies optimization over a product space of matrix-valued parameters representing the trainable layers of a deep model.
  • The framework targets continuously differentiable objectives that may be non-convex and non-smooth, while assuming the objective has a finite lower bound.

2 Theory vs. practice of Muon and Scion

Prior analyses do not faithfully represent practical Muon and Scion updates because they ignore per-layer computation and impose uniform classical smoothness. Gluon instead models layer-specific geometry, yielding assumptions and stepsizes that better align with training behavior.

  • 2.1 Layer-wise structure: Muon originally targeted hidden layers, whereas Scion extends the LMO-based rule to all layers and improves empirical performance.
  • 2.1 Layer-wise structure: Practical Gluon applies LMOs independently to each layer because a single global LMO is prohibitively expensive for extremely high-dimensional models.
  • 2.1 Layer-wise structure: The framework represents parameters in a matrix product space, enabling assumptions and hyperparameters adapted to each layer.
  • 2.2 A theory with predictive power: Classical L-smoothness imposes one constant across layers despite their differing geometries, producing tiny uniform stepsizes unlike tuned practical learning rates.
  • 2.2 A theory with predictive power: Layer-wise (L0, L1)-smoothness extends generalized smoothness to layer-specific, arbitrary-norm geometry and supports more expressive analysis.
  • 2.2 A theory with predictive power: Trajectory measurements on NanoGPT and CNN show the assumption approximately holds, with L0_i ≈ 0 and smoothness varying substantially across layers.
  • 2.2 A theory with predictive power: The theoretically derived stepsizes are almost the same as Pethick et al.’s tuned learning rates, indicating predictive value for practical selection.

3 Contributions

The paper contributes a generalized layer-wise framework and convergence theory for LMO-based optimizers, together with experiments across language modeling and image classification. These results connect formal guarantees to practical optimizer behavior.

  • Layer-wise (L0, L1)-smoothness models anisotropic deep-network structure using arbitrary norms and per-layer variation.
  • Gluon provides a faithful convergence analysis for layer-wise LMO algorithms and recovers Muon-type optimizers as special cases.
  • The convergence theory covers deterministic non-convex objectives, block-wise PŁ conditions, and stochastic optimization with momentum under bounded variance.
  • The theory establishes provable advantages over SGD through tighter convergence rates under more general assumptions and explains benefits of layer-wise learning rates.
  • Experiments on NanoGPT with FineWeb and CNNs with CIFAR-10 find that the smoothness assumption approximately holds throughout training and that prescribed stepsizes are practically useful.

4 Main theory and results

Gluon analyzes layer-wise LMO updates under general norms, recovering practical optimizers such as Muon as special cases. Its deterministic and stochastic convergence guarantees support adaptive, larger stepsizes that better reflect neural-network geometry and empirical tuning.

  • Gluon framework: Gluon independently minimizes a linearized loss for each parameter group within a layer-specific radius, with norms determining the resulting update.The framework permits layer-specific algorithmic choices through independently applied LMOs.
  • Special cases: Muon is recovered when hidden layers use the matrix spectral norm, while Euclidean and max-norm choices recover normalized GD and signGD variants.The framework also accommodates specialized norms for embedding and output layers under weight sharing.
  • Deterministic guarantees: Under layer-wise (L0, L1)-smoothness, deterministic Gluon has convergence guarantees for general non-convex objectives and block-wise PŁ settings.The theory extends classical smoothness analysis to per-layer variation and arbitrary norms.
  • Deterministic guarantees: When L0_i ≈ 0 across layers, the first convergence-rate term vanishes and the worst-case L1_max is replaced by a more favorable harmonic mean.For p = 1, the resulting rate matches the best-known stationary-point complexity O(L0Δ0/ε^2 + L1Δ0/ε).
  • Practical implications: The adaptive stepsizes can start larger when layer gradient norms are large, shrink as gradients decrease, and closely match empirically tuned values.This alignment suggests theory-driven stepsize selection can reduce reliance on costly manual tuning.
  • Stochastic guarantees: Stochastic Gluon uses unbiased bounded-variance gradient estimates, momentum, and time-varying stepsizes while retaining convergence under the generalized smoothness model.The stated schedule includes βk = 1−(k+1)^−1/2, with stepsizes proportional to k^−3/4 rather than constants tied to the total iteration count K.

5 Experiments

Experiments on NanoGPT and CNN test whether the layer-wise smoothness assumption holds along training trajectories and whether its theoretical stepsizes predict tuned values. Results support the assumption across architectures and show that non-uniform layer scaling is useful.

  • Training NanoGPT on FineWeb: The estimated trajectory-smoothness curves closely align with their approximations, indicating that Assumption 1 is approximately satisfied during NanoGPT training.The analysis examines embedding parameters and the 4th and 8th transformer blocks, with similar trends across blocks.
  • Training NanoGPT on FineWeb: Theoretical stepsizes closely match manually tuned values, supporting the predictive power of the prescriptions derived from Theorem 1.The NanoGPT configuration maps γ = 0.00036, ρ2 = 50, and ρ3 = 3000 into the paper’s notation.
  • Scaling factors: The best unScion performance uses ρ2 = 50 and ρ3 = 3000, with t_i = 0.018 for transformer layers and t_p = 1.08 for embedding and output layers.The result is consistent with the theoretical prediction and supports larger stepsizes for the embedding layer.
  • Additional ablations: Specialized norms better approximate trajectory smoothness than the standard Euclidean norm, and the layer-wise model also closely approximates smoothness during AdamW training.The transformer-to-embedding smoothness gap observed with unScion also appears during AdamW training.
  • Training CNN on CIFAR-10: CNN experiments on CIFAR-10 with full-batch unScion likewise show that Assumption 1 is approximately satisfied along the training trajectory.The experiment evaluates different parameter groups using a simplified setting with L0_i = 0.
  • Training CNN on CIFAR-10: For the CNN, estimated L1_i ≈ 3 for most parameter groups, while the classification-head value is approximately 0.03.This roughly two-orders-of-magnitude difference supports the much larger tuned radius for the head weights.

6 Conclusion and future work

The paper concludes that Gluon provides a layer-wise theoretical framework for LMO-based optimization that is empirically relevant and practically predictive. It also identifies exact-LMO computation, bounded variance, and stochastic adaptive stepsizes as remaining boundaries.

  • Conclusion: Gluon recovers Muon and Scion as special cases and introduces layer-wise (L0, L1)-smoothness for sharper convergence guarantees.The framework is designed to capture the anisotropic structure of modern deep networks.
  • Conclusion: Theoretical stepsizes closely match fine-tuned values, while experiments support the smoothness assumption throughout training.The reported evidence covers both NanoGPT language modeling and CNN image classification.
  • Future work: The analysis assumes exact LMO computations, whereas practical implementations use approximations.This discrepancy is identified as a direction for future research.
  • Future work: The stochastic convergence guarantees rely on bounded variance, which may fail in settings such as subsampling.The paper also notes that adaptive stepsizes are theoretically supported only in the deterministic setting.
  • Future work: A complete theoretical justification for adaptive stepsizes in the stochastic regime remains an open challenge.The paper reports that such methods nevertheless perform well empirically in that regime.

B Remarks on the theoretical results

The theoretical remarks explain how Gluon’s layer-wise LMO steps operate and how norm choices recover several optimizers. They also clarify why global guarantees based on max-norms do not fully reflect practical layer heterogeneity.

  • LMO interpretation: Each LMO step moves a layer’s iterate to the boundary of a norm ball, so radius, stepsize, and learning rate denote the same per-layer distance.The LMO objective is linear in the layer parameters, making the boundary property applicable at every step.
  • Layer-wise structure: Global max-norm guarantees require additional analysis to transfer to structured parameters and overlook the per-layer variability used in practical training.Per-layer tuning reflects the heterogeneous geometry of deep networks.
  • Special cases: Deterministic Gluon encompasses layer-wise normalized GD and layer-wise signGD through Euclidean- and infinity-norm choices, respectively.With suitable stepsizes, it also recovers the Gradient Method for (L0, L1)-smooth functions.
  • Special cases: Using the spectral norm for every parameter group recovers the per-layer deterministic Muon optimizer.Practical LLM variants add stochasticity and momentum to intermediate layers while using other methods for input and output layers.
  • Special cases: For LLMs, scaled spectral norms apply to transformer-block matrices, while a scaled 1→∞ norm applies to shared embedding and output parameters, yielding deterministic layer-wise unScion.The same optimizer framework also defines specialized norms for CNN biases, convolutional kernels, and classification heads.

C.2 Proof of Theorem 1

The section establishes convergence guarantees for Gluon under layer-wise smoothness and compares bounds that aggregate layer-dependent constants differently.

  • Deterministic convergence: Theorem 3 analyzes deterministic Gluon iterates under Assumption 1 and derives a stationarity guarantee after a sufficient number of iterations.The result is stated for iterates of deterministic Gluon with adaptive stepsizes.
  • Bound comparison: When L0_i ≈ 0 across layers, the first terms in the competing bounds effectively vanish.This is the empirical regime emphasized in the comparison of bounds.
  • Bound comparison: In that regime, the refined bound replaces the worst-case L1 maximum with a harmonic mean and is clearly superior according to the discussion.The comparison still depends on how the layer-wise constants are ordered when L0_i values are not negligible.
  • Bound construction: The convergence analysis offers two routes: upper-bounding each L1_i by a global maximum or reweighting layers using inverse L1_i values.The proof introduces increasing functions and their inverses to express sufficient iteration counts.
  • PŁ convergence: The section also states layer-wise PŁ convergence results, including separate cases for arbitrary nonnegative L1_i and L1_i = 0 across all layers.Theorem 5 gives corresponding sufficient iteration guarantees under Assumptions 1 and 3.
  • Stochastic convergence: For stochastic Gluon, the analysis assumes unbiased stochastic gradients with bounded relative variance and derives an adaptive-stepsize convergence theorem.The assumption is described as strong and unconventional, with a contraction-like interpretation relative to stochastic gradients.

D.2 Proof of Theorem 2

The section proves convergence for momentum-based Gluon under layer-wise smoothness, treating cases with vanishing and nonvanishing layer-wise L1 constants.

  • Momentum convergence: Theorem 8 analyzes Gluon with momentum schedule βk = 1−(k+1)^−1/2 under Assumptions 1 and 2.The proof decomposes the momentum analysis into four terms and bounds them separately.
  • Case analysis: The resulting argument considers separately the case L1_i = 0 for every layer and the case where L1_i ≠ 0.For nonzero L1_i, the proof chooses stepsizes proportional to 1/L1_i.
  • Case analysis: For nonzero L1_i, the stated bound contains layer-dependent factors 144(L1_i)^2 and a logarithmic dependence through 87 + 28 log(K).This expression appears in the theorem proof after selecting t_i = 1/(12L1_i).
  • Experimental context: The experimental setup includes NanoGPT training with distributed data parallelism on four NVIDIA A100 GPUs and CNN training on one A100 GPU.The pipelines use PyTorch and documented open-source codebases.
  • Experimental context: When an analytical LMO is unavailable, the experiments approximate it with Newton–Schulz iterations for computationally efficient orthogonalization.The passage states that this preserves the convergence behavior of the overall algorithm.

E.3.1 Empirical validation of Assumption 1

The experiments test whether layer-wise (L0, L1)-smoothness matches optimization trajectories and whether specialized norms fit those trajectories better than Euclidean norms.

  • Trajectory validation: Figures 5–7 evaluate estimated and approximated trajectory smoothness for embedding, 4th-block, and 8th-block parameter groups during unScion training.The reported fitted values include L0_p ≈ 0, L1_p ≈ 1.3 for embeddings and L0_i ≈ 0, L1_i ≈ 70 for the 4th block.
  • Trajectory validation: Across the examined parameter groups, estimated and approximated smoothness agree strongly, suggesting that Assumption 1 holds approximately along the optimization trajectory.The text reports similar patterns across all layers.
  • Norm comparison: The comparison evaluates specialized group norms against the standard Euclidean norm using relative mean squared error, where lower values indicate better fits.The Euclidean-norm experiment omits the first point and sets the underestimation penalty to λ = 0.
  • Norm comparison: Specialized norms provide better approximations than the standard Euclidean norm across several parameter groups.The relative mean squared error is consistently an order of magnitude lower with specialized norms.

E.3.3 Learning rate transfer from AdamW

The AdamW experiment examines whether trajectory smoothness can inform layer-specific learning-rate choices for Scion.

  • AdamW validation: The AdamW experiment evaluates estimated and approximated trajectory smoothness across several parameter groups using hyperparameters from Pethick et al.The results are presented in Figure 10.
  • Layer differences: The embedding layer has an L1_p value approximately 20–30 times smaller than the values for other parameter groups.This difference is observed in the fitted smoothness results.
  • Layer differences: The observed layer-wise relationship implies that the embedding-layer step size should be 20–30 times larger than the step sizes of the remaining layers.The text reports consistency with tuned parameters from Pethick et al.
  • Learning-rate transfer: Smoothness statistics collected during AdamW training can serve as a prior for initializing Scion learning rates, including larger embedding-layer stepsizes.The passage states that collecting these statistics is computationally inexpensive.

E.4 Training CNN on CIFAR-10

Additional CIFAR-10 CNN experiments examine layer-wise smoothness for unScion under full-batch and stochastic gradients. Across broader parameter groups, the trajectories support Assumption 1, including despite the added variance of stochastic gradients.

  • The experiments train a CNN on CIFAR-10 using the unScion optimizer and extend the analysis described in Section 5.2.
  • Full-batch (deterministic) gradients: Full-batch experiments evaluate estimated trajectory smoothness for a broader selection of parameter groups than in the main text.
  • Full-batch (deterministic) gradients: The deterministic results further support the validity of Assumption 1 with L0 ...
  • Stochastic gradients: Stochastic experiments use noisy gradients and momentum without a linear decay schedule, while applying the same layer-wise smoothness analysis.
  • Stochastic gradients: Despite the added variance, stochastic trajectories roughly adhere to Assumption 1 across the analyzed parameter groups.
Loading 2505.13416v1…