PaperlayerView exact revision on arXiv

Source-linked AI summary

Training Deep Learning Models with Norm-Constrained LMOs

Thomas Pethick, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falls, Volkan Cevher

arXiv:2502.07529v2cs.LGmath.OC

TL;DR

Neural-network optimizers typically adapt geometry from gradients, but this paper asks whether architecture-aware norms can be chosen in advance. It develops stochastic LMO methods for both constrained and unconstrained optimization, including SCION, and reports transferable hyperparameters, strong training performance, and memory efficiency. The paper also identifies practical and analytical boundaries for spectral-norm-based methods and SCION’s Lion-K interpretation.

  • Problem

    The paper asks whether neural-network optimization should adapt geometry a priori from architecture rather than explore it on-the-fly during training.

  • Method

    The paper develops stochastic LMO-based algorithms for norm-ball geometry, extends them to unconstrained problems, and instantiates operator-norm choices in SCION.

  • Results

    SCION consistently exhibits hyperparameter transferability across model settings, with numerical evaluation spanning MLPs, CNNs, ImageNet ViTs, and NanoGPT models up to 3B parameters.

  • Takeaways & Limitations

    The framework unifies popular optimization methods through norm choice while enabling SCION to use one parameter set and one half-precision gradient set.

  • Takeaways & Limitations

    Spectral-norm projections can require expensive SVDs for dense model weights, and SCION’s Lion-K interpretation lacks the cited smooth-K continuous-time analysis.

Abstract

from arXiv · show

In this work, we study optimization methods that leverage the linear minimization oracle (LMO) over a norm-ball. We propose a new stochastic family of algorithms that uses the LMO to adapt to the geometry of the problem and, perhaps surprisingly, show that they can be applied to unconstrained problems. The resulting update rule unifies several existing optimization methods under a single framework. Furthermore, we propose an explicit choice of norm for deep architectures, which, as a side benefit, leads to the transferability of hyperparameters across model sizes. Experimentally, we demonstrate significant speedups on nanoGPT training using our algorithm, Scion, without any reliance on Adam. The proposed method is memory-efficient, requiring only one set of model weights and one set of gradients, which can be stored in half-precision. The code is available at https://github.com/LIONS-EPFL/scion .

1. Introduction

The paper asks whether neural-network optimization should adapt to architecture geometry a priori rather than infer it online, and develops LMO-based methods to do so. It introduces uSCG and SCION, unifies optimizer updates through norm choices, and reports transferable hyperparameters and strong numerical performance.

  • Motivation: The paper asks whether adapting the optimizer a priori to known neural-network structure is preferable to adapting geometry on-the-fly.It focuses on choosing the underlying norm used to measure parameter-space distances.
  • Method: uSCG extends stochastic conditional-gradient methods to unconstrained problems, improving theoretically and practically when the norm-ball matches the problem’s natural geometry.The approach uses linear minimization oracles over norm balls.
  • Theoretical rates: O(n^-1/4) is the order-optimal convergence rate established for uSCG under general nonconvexity and stochasticity.The stepsize is agnostic to the Lipschitz constant, unlike steepest descent.
  • Unification: The LMO framework unifies popular algorithms through norm choice and provides the first provable rate for Muon with and without weight decay.It also motivates SCION, an operator-norm method for deep learning.
  • Numerical validation: SCION stores only one parameter set and one gradient, which can be kept in half-precision to reduce memory use in large-scale training.The paper presents this as an implementation benefit of the framework.
  • Numerical validation: SCION consistently exhibits hyperparameter transferability across MLPs, CNNs, ViTs on ImageNet, and NanoGPT models up to 3B parameters.The experiments also report tolerance to large batch sizes and superior performance attributed to a priori adaptation.

2. Preliminaries

The preliminaries formulate smooth possibly nonconvex optimization in general norms, define the LMO over norm balls, and contrast constrained conditional-gradient updates with the paper’s unconstrained stochastic setting. Momentum is introduced because stochastic LMOs are biased even when stochastic gradients are unbiased.

  • Problem setting: The paper studies smooth, possibly nonconvex objectives in a non-Euclidean norm, either unconstrained or restricted to a norm-ball set.The constrained set is denoted D, while the unconstrained case can use X = R^d.
  • Linear minimization oracle: The linear minimization oracle minimizes a linear objective over a norm-constrained set, and its output has norm at most ρ.Under norm constraints, the oracle is scale invariant, so only the input direction matters.
  • Conditional Gradient: Conditional Gradient solves the constrained problem using an LMO and preserves feasibility through simplicial combinations.The method proceeds with stepsizes γ_k in (0, 1).
  • Unconstrained setting: uSCG applies LMO-based updates to unconstrained problems, extending the usual constrained use of Conditional Gradient.The preliminaries distinguish this setting from the constrained problem that classical CG solves.
  • Stochastic setting: Even unbiased stochastic gradients yield biased LMO directions, so the algorithms average past gradients using momentum coefficients α_k in (0, 1].The averaging is used to help establish algorithmic convergence.

3. Our Methods

The paper introduces unconstrained and constrained stochastic conditional-gradient methods using norm-constrained LMOs, then designs operator-norm choices for deep networks. These choices provide norm control, layerwise computable updates, and width-independent maximal-update learning rates.

  • uSCG and SCG: uSCG extends stochastic conditional-gradient optimization to unconstrained problems by summing LMO updates instead of using the constrained convex combination.Its update uses momentum-based stochastic gradients and returns a uniformly sampled iterate.
  • uSCG and SCG: SCG revisits stochastic conditional-gradient optimization for nonconvex deep-learning objectives, while weight decay interpolates between uSCG and SCG.For positive weight decay below one, the method remains an SCG instance for a constrained problem with a larger radius.
  • Norm control: Both uSCG and SCG explicitly control parameter norms, with norm bounds useful for long runs and avoiding overfitting in multi-epoch training.Under norm constraints, the LMO output is bounded by the constraint radius.
  • Operator-norm design: Operator norms constrain layer weight matrices and biases so bounded inputs produce bounded hidden states and logits.The paper considers RMS-, spectral-, row-, column-, and sign-based choices, with layerwise LMOs summarized in Tables 2–4.
  • Operator-norm design: For 1-hot encoded language inputs, several input-layer operator norms are equivalent on the restricted domain, enabling simple LMO update rules such as column normalization or sign updates.The recommended SCION and UNCONSTRAINED SCION configurations use different operator norms across first, intermediary, and last layers.
  • Hyperparameter transfer: The maximal update learning rate γ∗ is independent of width under the paper’s simplified momentum setting, allowing a rate tuned on a smaller model to transfer to a wider model.The analysis covers activation functions including ReLU, GELU, and Tanh, with α_k = 1.

4. Related Works

The paper frames LMO-based optimization as a unifying perspective connecting norm-based steepest descent, momentum optimizers, normalization methods, and related constrained algorithms. It distinguishes its approach through scale-invariant LMOs, all-layer application, and direct computation without explicit preconditioners.

  • Unifying framework: uSCG and SCG are invariant to gradient magnitude and avoid computing dual norms independently across layers.The scale invariance of the LMO also supports convergence analysis without knowledge of the Lipschitz constant.
  • Limitations and scope: Sharp-operator methods face nontrivial extensions to constrained problems, while spectral-norm projections can require expensive SVDs for dense model weights.The paper notes that gradients used in LMOs are often low-rank, enabling more efficient SVD approximations.
  • Unifying framework: LMO-based updates unify algorithms including Muon, normalized SGD variants, continuous greedy, and spectral-norm methods through different norm choices.Muon can be written exactly as uSCG under non-Nesterov momentum, while other norm choices recover related update rules.
  • Comparisons: Unlike practical Muon, uSCG and SCG apply their updates to all layers and demonstrate stepsize transferability.Muon typically excludes the first and last layers, using Adam(W) or SGD there.
  • Comparisons: The paper relates SCION to Lion-K, but notes that Lion-K’s continuous-time analysis does not apply when the chosen set is nonsmooth.This distinction arises because SCION uses the LMO as ∂K and the resulting K is not smooth.
  • Comparisons: SCION computes an LMO at each iteration rather than maintaining an explicit preconditioner, connecting it to—but distinguishing it from—PSGD methods.Both approaches relate to whitening the update, but their computational mechanisms differ.

5. Analysis

The analysis establishes nonasymptotic convergence guarantees for unconstrained and constrained stochastic conditional-gradient methods under non-Euclidean smoothness and unbiased bounded-variance gradients. The results characterize how momentum and stepsize choices trade off convergence speed, noise sensitivity, and criticality.

  • Assumptions: The analysis assumes gradient Lipschitz continuity in the dual norm, a lower-bounded objective, and an unbiased stochastic-gradient oracle with bounded variance.The smoothness norm need not be Euclidean, allowing geometries relevant to neural-network training.
  • Proof strategy: The convergence proofs use a descent inequality, second-moment error analysis, and averaging arguments under horizon-dependent constant stepsizes or vanishing momentum.The framework also allows time-based guarantees by replacing horizon-dependent stepsizes with iteration-dependent ones.
  • Interpretation: uSCG does not require knowledge of the Lipschitz constant L because it can be viewed as normalized steepest descent through the LMO.The LMO-based normalization supplies the relevant direction without explicitly using gradient magnitude.
  • Momentum effects: Constant momentum gives faster worst-case convergence but approaches a noise-dominated region with radius proportional to σ.Vanishing momentum instead guarantees a small expected dual-norm gradient at a slower rate.
  • Momentum effects: Smaller momentum reduces a σ-proportional constant term but can amplify O(1/√n) and O(σ/n) terms, requiring stepsize alignment.For vanishing momentum, the theory suggests a smaller constant stepsize such as γ = 3/(4n^3/4).

6. Experiments

Experiments evaluate SCION across language and vision settings, model sizes, and batch sizes. SCION transfers hyperparameters across widths, handles larger batches effectively, and achieves favorable validation-loss or training-time outcomes.

  • Implementation: SCION uses an efficient Newton-Schulz implementation for spectral-norm LMO computation.The implementation follows Jordan et al. and Bernstein & Newhouse.
  • NanoGPT: SCION’s optimal stepsize transfers across NanoGPT model widths, unlike Adam and Muon.Experiments span models between 64M and 1B parameters.
  • NanoGPT: 2.984 validation loss is achieved by UNCONSTRAINED SCION versus 2.988 for Muon on the 124M model.The comparison uses Muon tuned on the largest model size and SCION evaluated across three runs.
  • Memory and architecture: SCION completely removes Adam from the Muon implementation and stores only one half-precision set of weights and gradients across layers.The experiments also demonstrate compatibility with weight sharing.
  • 3B model: UNCONSTRAINED SCION outperforms all other methods on the 3B-parameter GPT experiment.The configuration is selected using the optimal 124M-parameter proxy model and increases model depth to 36.
  • Large batches: SCION maintains low validation loss as batch size increases and matches Muon’s validation loss with 25% smaller wallclock time at batch size 6144.The large-batch experiment fixes total tokens and rescales the number of steps.

B. Method

The method selects layerwise operator norms to preserve activation-scale guarantees and make LMO updates practical. Its scaling addresses input-dimension effects that otherwise prevent constant preactivations.

  • Algorithm: SCION and uSCG use layerwise norm choices while ignoring biases in the presented linear MLP treatment.Algorithm 3 specifies horizon, initialization, momentum, stepsize, and layer radii as inputs.
  • Input radius scaling: The proposed norm scaling resolves nonconstant preactivations when the input dimension exceeds the output dimension.A coordinate check shows the issue under spectral scaling, while max(1, √(d_out/d_in)) resolves it.
  • Norm design: The RMS→RMS operator norm bounds each layer’s output RMS norm, preserving a bounded RMS-norm input for the next layer.Its LMO can be computed efficiently, enabling direct use in the update rule.
  • Other norm choices: Alternative norms such as ℓ1→RMS and ∞→∞ are valid when output guarantees are converted into bounds on the next layer’s input norm.The resulting LMO choices for the three norm configurations are summarized in Table 6.
  • Initialization: Semi-orthogonal, column-wise normalized Gaussian, and row-wise normalized Gaussian initializations are scaled elementwise by the corresponding LMO scaling.These are presented as alternative initialization procedures.

C. Proofs for Section 3 (Our Methods)

The section specifies spectral-norm-based LMOs for neural-network layers and analyzes their update dynamics. Under stated data, initialization, and activation assumptions, the resulting maximal-update learning rate is invariant to hidden-layer widths.

  • Problem setting: The analysis defines an L-layer network through layer weight matrices, preactivations, postactivations, and a linear output transformation.The notation covers input dimension, hidden-layer widths, output dimension, and element-wise activations.
  • Assumptions: The analysis assumes bounded input second moments, a stability-preserving initialization scheme, and spectral-norm-based parameter updates.Both semi-orthogonal initialization and, with high probability, Gaussian initialization satisfy the stated initialization condition.
  • Spectral-norm LMOs: The method adopts layer-wise spectral-norm LMOs, with explicit constructions for the first and intermediate weight matrices.These choices correspond to the Spectral → Spectral → Spectral configuration and also cover the ColNorm → Spectral → Spectral configuration under 1-hot encoding.
  • Width-invariant updates: The maximal-update analysis considers general losses including MSE and logistic loss and shows that the optimal learning rate is independent of hidden-layer widths.This result is stated for Algorithm 1 with αk = 1, except when all activations are simultaneously zero.
  • LMO derivation: The spectral LMO for each layer is derived from the network gradient using the chain rule and the corresponding forward-pass representations.The derivation establishes the oracle used in the subsequent update-dynamics analysis.
  • Width-invariant updates: As widths increase, the unchanged learning rate preserves a stable but maximally adaptive regime, enabling direct transfer of a rate tuned on a smaller model.The convergence section analyzes uSCG using a descent inequality, second-moment error behavior, and a final rate derivation.

D.1. Convergence analysis of uSCG

The uSCG convergence analysis builds a descent template, controls stochastic momentum error recursively, and combines the bounds into convergence guarantees for constant stepsizes and momentum.

  • Template inequality: The analysis first establishes a generic dual-norm gradient inequality valid for constant or vanishing momentum αk ∈ (0, 1].This template applies throughout the uSCG analysis under the smoothness assumption.
  • Descent analysis: The descent proof relates the LMO step to the dual norm of the gradient and uses the norm-ball radius to bound the update terms.The argument applies Cauchy–Schwarz, the definition of the dual norm, and the radius of D.
  • Error control: The stochastic error is controlled through a linear recursive inequality using unbiasedness, bounded variance, Lipschitz continuity, and the update definition.The recursion is then combined with the descent template to derive the final rate.
  • Constant momentum: With constant stepsize γ = 1/√n and constant momentum α ∈ (0, 1), Theorem 5.4 provides the stated uSCG convergence rate.The proof combines the template inequality with the error bounds under Assumptions 5.1 and 5.3.

D.1.2. VANISHING αk

This section analyzes uSCG with vanishing momentum by bounding the gradient-error recursion and combining it with the descent template to obtain an accuracy guarantee.

  • Gradient-error bound: Lemma D.4 bounds the gradient error for uSCG when the momentum parameter vanishes over iterations.The bound is derived from the recursive inequality under the paper’s smoothness and stochasticity assumptions.
  • Convergence guarantee: Theorem 5.5 combines the descent template and vanishing-momentum error bound to state the convergence rate for uSCG.The proof estimates the resulting summation using the integral test and the decreasing function x ↦ 1/x^1/4.
  • Parameter schedule: The analysis uses a constant stepsize satisfying 1/(2n^3/4) < γ < 1/n^3/4 together with vanishing momentum αk = 1/√k.These parameter choices control the accumulated error terms over n iterations.

D.2. Convergence analysis of SCG

The SCG convergence analysis establishes a nonconvex descent inequality over a norm-ball, controls the stochastic momentum error, and combines these ingredients into a Frank–Wolfe-gap guarantee.

  • Constrained stationarity: The analysis measures constrained progress through the Frank–Wolfe gap and introduces the normal-cone stationarity condition for the norm-ball D.The normal cone appears in the constrained first-order condition.
  • Descent inequality: Lemma D.5 gives a descent characterization for Algorithm 2 for arbitrary stepsize γ and momentum αk ∈ (0, 1].The proof uses LMO optimality and bounds the distance between the oracle output and the current iterate by the set diameter.
  • Error control: The stochastic momentum error is bounded recursively using bounded variance and the Euclidean diameter D2 of the feasible set.The resulting recursion is applied to the descent inequality before summing over iterations.
  • Convergence guarantee: With constant stepsize γ = 1/√n and constant momentum α ∈ (0, 1), Theorem 5.6 states the SCG convergence guarantee for all u ∈ D.The proof sums the one-step inequality and controls the momentum terms using subadditivity and a geometric series.

D.2.2. SCG WITH VANISHING α

This section analyzes convergence for stochastic LMO methods with constant and vanishing momentum, including a direct-LMO variant. It establishes bounds through recursive estimates and a stepsize choice γ = 1/√n.

  • SCG with vanishing α: Lemma D.7 bounds the gradient error for Algorithm 2 when the momentum parameter vanishes.
  • Proof strategy: The analysis derives the gradient-error estimate by combining a recursion, auxiliary lemmas, expectation bounds, and summation over k = 1 to n.
  • SCG with vanishing α: Theorem 5.7 gives a convergence-rate result for SCG with vanishing α_k under Assumptions 5.1 and 5.3.
  • Averaged LMO directional descent: Algorithm 4 computes the LMO direction directly from each stochastic gradient, averages these directions with constant momentum, and returns a uniformly sampled iterate.
  • Averaged LMO directional descent: For Algorithm 4, Lemma D.8 analyzes iterates using stepsize γ = 1/√n under Assumptions 5.1 and 5.3.
  • Auxiliary lemmas: The appendix introduces linear-recursion lemmas covering constant momentum, constant coefficients, and vanishing coefficients for stochastic-gradient error analysis.

E.1. Additional experiments

Additional experiments examine hyperparameter transfer, memory-efficient implementation, and performance across MLP, GPT, NanoGPT, CNN, and DeiT settings. The reported results support width-transfer behavior and benefits from SCION’s norm control in several experiments.

  • MLP: On a 3-layer ReLU MLP, the optimal learning rate transfers across model width for all three output-layer configurations.
  • Shallow GPT: Shallow GPT experiments sweep input-layer LMOs and compare SCION with unconstrained variants under multiple layer configurations.
  • Memory-efficient implementation: SCG and uSCG can be implemented with one set of parameters and one set of gradients, with gradients possibly stored in half-precision.
  • ViT: In DeiT-base experiments, replacing LayerNorm with nonlearned RMSNorm significantly improves SCION performance, while the same improvement is not observed for AdamW.
  • NanoGPT: On NanoGPT, SCION produces smaller weight norms than unconstrained SCION, and norm control appears particularly important during long runs.
Loading 2502.07529v2…