PaperlayerView exact revision on arXiv

Source-linked AI summary

Drop-Muon: Update Less, Converge Faster

Kaja Gruntkowska, Yassine Maziane, Zheng Qu, Peter Richtárik

arXiv:2510.02239v1cs.LGmath.OCstat.ML

TL;DR

The paper asks whether updating every layer at every step is always optimal. It introduces Drop-Muon, which updates randomized subsets of layers, and reports faster wall-clock convergence than full-network Muon in controlled CNN experiments. Its cost analysis finds full-network updates are not optimal except under a specific relationship between layer smoothness constants.

  • Problem

    The paper questions the default of updating every layer at every step, asking whether full-network updates are universally optimal.

    Source: §1 · p. 2

  • Method

    Drop-Muon samples a randomized subset of layers per iteration, updates those layers, and freezes the others.

    Source: §2 · pp. 2–3

  • Results

    Roughly 1.2× faster convergence on average across accuracy thresholds was reported for Drop-Muon on Fashion-MNIST.

    Source: Results on Fashion-MNIST · p. 9

  • Takeaways & Limitations

    Full-network updates are not optimal unless a specific relationship between layer smoothness constants holds.

    Source: §3 · p. 4

    The convergence analysis uses supp(D)-layer-wise smoothness assumptions, which provide separate constants for layer subsets without restricting the function class beyond standard smoothness.

    Source: §4.1 · p. 5

Abstract

from arXiv · show

Conventional wisdom in deep learning optimization dictates updating all layers at every step-a principle followed by all recent state-of-the-art optimizers such as Muon. In this work, we challenge this assumption, showing that full-network updates can be fundamentally suboptimal, both in theory and in practice. We introduce a non-Euclidean Randomized Progressive Training method-Drop-Muon-a simple yet powerful framework that updates only a subset of layers per step according to a randomized schedule, combining the efficiency of progressive training with layer-specific non-Euclidean updates for top-tier performance. We provide rigorous convergence guarantees under both layer-wise smoothness and layer-wise $(L^0, L^1)$-smoothness, covering deterministic and stochastic gradient settings, marking the first such results for progressive training in the stochastic and non-smooth regime. Our cost analysis further reveals that full-network updates are not optimal unless a very specific relationship between layer smoothness constants holds. Through controlled CNN experiments, we empirically demonstrate that Drop-Muon consistently outperforms full-network Muon, achieving the same accuracy up to $1.4\times$ faster in wall-clock time. Together, our results suggest a shift in how large-scale models can be efficiently trained, challenging the status quo and offering a highly efficient, theoretically grounded alternative to full-network updates.

1 Introduction

Recent geometry-oriented optimizers use layer-specific non-Euclidean linear minimization oracles, but still update every layer at each step. The paper questions whether full-network updates are optimal and frames the network parameters as layer-specific blocks.

1 Introduction

  • Muon and its successors are emerging alternatives to Adam and related methods, motivated by geometry-aware optimization and results in large-language-model training.The cited successors are Scion and Gluon.Source: §1 · p. 1
  • Layer-specific linear minimization oracles optimize over non-Euclidean norm balls, while the shared update rule updates all layers at every iteration.The oracle selects a point in a layer's ball by minimizing its inner product with the momentum term.Source: §1 · pp. 1–2
  • Full-network updates may not be optimal, challenging a design choice shared by the layer-wise optimizers discussed in the introduction.The paper examines this question from theoretical and practical perspectives.Source: §1 · p. 2

1.1 Background

  • The model parameters are represented as layer-specific blocks in a product space, with each block corresponding to one layer.Each block is a matrix-valued parameter variable.Source: §1.1 · p. 2

2 The Algorithm

Drop-Muon replaces full-network updates with randomized subset updates, freezing unselected layers. Its non-Euclidean layer-wise update also has a normalized steepest-descent form, which reduces to layer-wise momentum SGD under Euclidean norms.

2 The Algorithm

  • Drop-Muon samples a layer subset each iteration, updates only those parameters, and freezes the rest.The subset is drawn from a user-defined distribution.Source: §2 · pp. 2–3
  • The main LMO update has an equivalent sharp-operator formulation that gives a layer-wise normalized steepest-descent step.The step size is scaled by the dual norm of the momentum term.Source: §2 · p. 3
  • With Euclidean norms, the sharp operator is the identity and the update becomes layer-wise stochastic gradient descent with momentum.The paper notes this equivalence is used in its proofs.Source: §2 · p. 3

3 Cost Model

The cost analysis challenges full-network updates as a general default, while the paper presents Drop-Muon as a randomized layer-subsampling framework with convergence guarantees. Controlled CNN experiments report faster wall-clock convergence than standard full-network Muon.

3 Cost Model

  • Randomized Progressive Training updates layers from the last layer down to the smallest sampled index, matching backpropagation's gradient computation.Activations before that index can be cached and reused, while gradients for frozen layers before it are not computed.Source: §3 · p. 3
  • The cost model combines per-layer forward/backward costs with the costs of sharp-operator computation and parameter updates.The expected total cost is evaluated for a fixed target accuracy using the number of iterations to reach an ε-stationary point in expectation.Source: §3 · pp. 3–4
  • Full-network updates are not optimal unless a specific relationship between layer smoothness constants holds.The cost analysis accounts for per-layer computation and update costs; the paper describes the required condition as unlikely in practice.Source: §3 · p. 4
  • Drop-Muon generalizes LMO-type methods by allowing randomized layer subsets, with full-network training as a special case.The contribution summary focuses on Randomized Progressive Training, which is aligned with backpropagation mechanics and avoids redundant gradient computations.Source: §3 · p. 4
  • The paper establishes convergence guarantees under layer-wise smoothness and layer-wise (L0, L1)-smoothness, including a deterministic gradient variant.It identifies the latter as the first convergence guarantees for progressive-training-style methods in the non-smooth setting.Source: §3 · pp. 3–4
  • Drop-Muon reaches the same accuracy as standard full-network Muon up to 1.4× faster in wall-clock time across controlled CNN experiments.The experiments cover MNIST, Fashion-MNIST, and CIFAR-10.Source: §3 · p. 4

4 Randomized Progressive Training

Drop-Muon uses randomized progressive layer updates and establishes convergence guarantees across deterministic and stochastic settings. Its cost analysis finds that full-network updates are optimal only under restrictive smoothness conditions.

4 Randomized Progressive Training

  • At each iteration, randomized progressive training samples a smallest active layer and updates it and all later layers, freezing earlier layers.The sampling distribution must assign positive probability to the first layer so every layer can be updated.Source: §4 · pp. 4–5, §4.1 · p. 5

4.1 Iteration Complexity – Deterministic Gradient Setting

  • Under layer-wise smoothness, the deterministic method achieves an O(K^-1) rate for a weighted sum of squared gradient-component norms.Every gradient component must contribute to the weighted average, requiring p1 > 0.Source: §4.1 · p. 5
  • Under layer-wise (L0, L1)-smoothness, the deterministic method achieves an O(K^-1/2) rate for a weighted sum of gradient-component norms.Full-network training recovers the deterministic Gluon rate in this regime.Source: §4.1 · p. 6

4.1.1 Cost Optimization

  • The cost analysis models per-step computation using layer costs and fixed overhead, accounting for backpropagation from the last layer to the smallest active index.Forward activations before that index can be cached and reused.Source: §4.1.1 · p. 7
  • Full-network updates minimize the analyzed cost only under a specific condition on layer smoothness constants, independent of cost parameters.The layer-wise (L0, L1)-smooth analysis likewise makes full-network updates optimal only if the first layer has the largest smoothness constant.Source: §4.1.1 · p. 7

4.2 Iteration Complexity – Stochastic Gradient Setting

  • The stochastic method achieves an O(K^-1/4) convergence rate for a weighted sum of gradient-component norms.The guarantee assumes a stronger smoothness variant than the deterministic analysis and an unbiased stochastic gradient estimator with bounded variance.Source: §4.2 · pp. 7–8

5 Prior Work on Progressive Training

Progressive Training began as a way to grow GANs over time, while later randomized variants supplied convergence analysis but retained practical limitations. Drop-Muon shares the randomized-subnetwork idea while adding layer-wise, non-Euclidean updates and stochastic-gradient support.

5 Prior Work on Progressive Training

  • Progressive Training began in GANs by gradually increasing input resolution and network depth from a shallow model.This progressive growth was reported to improve computational efficiency and stabilize training.Source: §5 · p. 8
  • Randomized Progressive Training established convergence rates for general smooth objectives, but its sketched-GD formulation flattens layer structure and requires non-stochastic gradients.These limitations restrict its practical utility for large-scale architectures and prevent recovering layer-specific methods such as Muon as special cases.Source: §5 · p. 8
  • Drop-Muon shares randomized subnetwork updates with prior RPT while supporting layer-specific non-Euclidean rules and stochastic gradients.The cited RPT approach is described as a randomized proxy for progressive training.Source: §5 · p. 8

6 Experiments

Controlled CNN experiments show Drop-Muon can reach accuracy targets sooner in wall-clock time, despite slightly slower per-epoch convergence. The accompanying background explains its optimizer context and the smoothness and coordinate-descent ideas relevant to its analysis.

6 Experiments

  • 27s versus 39s per epoch on MNIST lets Drop-Muon train faster overall despite slightly slower per-epoch convergence.Experiments use three-layer CNNs and compare against standard Muon under identical hyperparameter configurations; accuracy is tracked by epochs and wall-clock time.Source: §6 · p. 8, Figure 1 · p. 9
  • Up to 1.4× speed-up on MNIST and roughly 1.2× faster average convergence on Fashion-MNIST; on CIFAR-10, Drop-Muon reaches 90% training accuracy earlier in wall-clock time.The MNIST and Fashion-MNIST comparisons aggregate time-to-target results over multiple seeds; the CIFAR-10 comparison is from a representative run.Source: Figure 2 · p. 9, §6 · p. 9, Results on Fashion-MNIST · p. 9

A.1 Muon (and Friends)

  • Muon uses spectral-norm steepest-descent updates, while Scion extends LMO updates across layer types and Gluon provides a convergence framework for arbitrary norm balls.Muon approximates orthogonalization with Newton-Schulz and incorporates momentum.Source: §A.1 · p. 15

A.2 Generalized Smoothness

  • The paper’s layer-wise (L0, L1)-smoothness analysis allows subsets of layers to have distinct effective constants, yielding tighter bounds without restricting the function class.This contrasts with the global smoothness condition discussed in the comparison to Riabinin et al. (2025).Source: §A.2 · pp. 15–16

A.3 Randomized Block Coordinate Descent

  • Randomized block coordinate descent updates one randomly selected block while leaving others unchanged; blockwise smoothness can permit stepsizes scaling with 1/L_i rather than 1/L.The global smoothness constant must cover worst-case variation across directions, including cross-block interactions, and can exceed individual block constants.Source: §A.3 · p. 16

B More on Smoothness Assumptions

The appendix shows why layer-specific smoothness constants matter, extends the framework to arbitrary layer sampling, and gives convergence results for the layer-wise smooth case.

B More on Smoothness Assumptions

  • For nested subsets, the layer-wise smoothness constants can be chosen so the smaller subset’s constants are bounded by those for the larger subset.This follows from Assumption 4.2 and is stated in Lemma B.1.Source: B More on Smoothness Assumptions · p. 18
  • Without further assumptions, nested layer subsets need not admit finite smoothness constants related by a uniform function.The appendix demonstrates this with scalar-block constructions and shows that the obstruction extends beyond single-layer subsets.Source: B More on Smoothness Assumptions · pp. 18–19

C Arbitrary Sampling

  • Drop-Muon’s arbitrary-sampling framework updates only the randomly selected subset, with sampling distribution and layer-wise norms jointly parameterizing the algorithms.The general convergence results recover the RPT guarantees when specialized to that setting.Source: C Arbitrary Sampling · p. 20

D.1 Layer-Wise Smooth Case

  • The layer-wise smooth-case bound is meaningful only when every layer receives positive sampling weight.The appendix identifies this requirement as ensuring every layer is sampled with positive probability.Source: §D.1 · p. 21

D.2 Layer-Wise Generalized Smooth Case

This section establishes a layer-wise generalized-smooth convergence guarantee and analyzes deterministic-gradient costs for randomized layer sampling. It derives sampling choices for submodels and blocks, examines how costs and smoothness determine block size, and formulates training-cost optimization across two regimes.

D.2 Layer-Wise Generalized Smooth Case

  • Theorem D.4 gives a generalized-smooth convergence guarantee, meaningful only when every layer has positive sampling weight.The stated condition is equivalent to requiring positive inclusion probability for every layer.Source: §D.2 · p. 22

E Optimizing the Cost – Deterministic Gradient Case

  • The expected cost per iteration is expressed through two marginal probabilities, so evaluating a sampling scheme’s cost reduces to computing those marginals.The analysis considers example schemes under both layer-wise and generalized layer-wise smoothness.Source: E Optimizing the Cost – Deterministic… · p. 26

E.1 Smooth Case

  • Because optimizing over all distributions on the power set is intractable for large b, the smooth-case analysis restricts sampling to parametric families with simpler objectives.For certain families, the ratio objective becomes a linear-fractional program in cumulative marginals, with closed-form or efficient solutions.Source: §E.1 · p. 28

E.1.1 Randomized Progressive Training

  • For randomized progressive training, the analysis derives a recursive prescription for optimal sampling probabilities from a linear-fractional program.At an optimum, a positive coordinate requires its corresponding constraint to be tight; the recursion uses this property.Source: §E.1.1 · pp. 29–30

E.1.2 τ-nice Sampling

  • For τ-nice sampling, τ = b need not be optimal: increasing τ can raise the Lipschitz term while lowering another cost factor, and a stated linear-scaling example yields τ⋆ = 1.The analysis notes that the balance depends on the problem-specific generalized smoothness constants.Source: §E.1.2 · pp. 31–33

E.1.3 τ-Submodel Sampling

  • For τ-submodel sampling, the analysis specializes to sampled blocks that partition [b], with b divisible by τ, and reduces the total cost to a block-based expression.Each partition block is selected with positive probability.Source: §E.1.3 · p. 33
  • Optimal block probabilities are proportional to each block’s worst-case local smoothness, and strong duality certifies the solution.The analysis also derives the corresponding minimal expected cost.Source: §E.1.3 · p. 35
  • The cost-minimizing block size need not be τ = b: it is b when L0(τ) is constant, may be intermediate for sublinear growth, and is 1 for at least linear growth.These cases are derived under constant costs and block-index-independent worst-case local smoothness.Source: §E.1.3 · p. 36

E.1.4 Arbitrary Submodel Sampling

  • For arbitrary submodel sampling, optimal block probabilities follow from a linear-fractional reduction, and partitioned sampling beats full-model training under a derived cost condition.Grouping similar-cost layers or same-position layers across transformer blocks provides settings for interpreting that condition.Source: §E.1.4 · pp. 37–39

E.2 Generalized Smooth Case

  • Generalized-smooth cost optimization separates into O(1/ε^2) and O(1/ε) regimes; distribution optimization is generally intractable, but certain parametric families yield efficiently solvable linear-fractional programs.The section introduces specific examples of such families.Source: §E.2 · p. 40

E.2.1 Randomized Progressive Training

The appendix compares sampling designs through their expected costs. Its results show that full-layer or full-model sampling need not be optimal, depending on smoothness and cost trade-offs.

E.2.1 Randomized Progressive Training

  • The all-layers sampling distribution (p1, p2, …, pb) = (1, 0, …, 0) is not an optimal solution.In randomized progressive training, a starting index is sampled and the updated set consists of that index and all later layers.Source: §E.2.1 · pp. 40, 43

E.2.2 τ-Nice Sampling

  • A smaller τ can minimize either expected cost when the smoothness-related factor grows fast enough to offset the decreasing factor in τ.The appendix states that even τ = 1 can be optimal under this trade-off.Source: §E.2.2 · p. 45

E.2.3 τ-Submodel Sampling

  • For partitioned τ-submodel sampling, the best number of blocks depends on smoothness-constant growth and its interaction with update costs, and need not be b.The expected cost’s dependence on τ is more complex for one of the objectives.Source: §E.2.3 · p. 46

E.2.4 Arbitrary Submodel Sampling

  • With arbitrary submodel sampling, submodel training can improve on full-model training in certain regimes.For a fixed partition, the appendix also derives probabilities minimizing one expected-cost objective.Source: §E.2.4 · p. 47

F Convergence Results – Stochastic Gradient Case

The stochastic-gradient appendix gives a convergence result for Algorithm 1 under stated assumptions and connects its RPT specialization to Theorem 4.4. It also describes the benchmark and model setup used for experiments.

F Convergence Results – Stochastic Gradient Case

G Experiments

  • Experiments evaluate Drop-Muon on MNIST, Fashion-MNIST, and CIFAR-10 using 3-layer CNNs of varying capacity.MNIST and Fashion-MNIST runs used Tesla V100-SXM2-32GB GPUs; CIFAR-10 used a mix of A100, P100, and GTX 1080 Ti GPUs.Source: G Experiments · p. 61

G.1 Key Training Parameters

The experiments use a grid search to identify Drop-Muon configurations that reduce training time while maintaining final accuracy. They test a schedule that shifts layer-sampling emphasis from shallow to deep layers and evaluate results using two multi-run aggregation procedures.

G.1 Key Training Parameters

  • Grid search identified uniform and epoch-shift sampling configurations that consistently reduced training time while maintaining final accuracy.The study varied batch sizes, learning rates, and channel configurations; it fixed momentum at 0.5 and used five Newton-Schulz iterations.Source: §G.1 · p. 61

G.2 Epoch-Shift Distribution

  • Epoch-shift sampling favors shallow layers early, then shifts probability toward deeper layers as training progresses; α controls the shift's strength.The illustrated schedule uses α = 0.5.Source: §G.2 · pp. 61–62, Figure 4 · p. 62

G.3 Evaluation Metrics

  • Results are aggregated across random seeds using normalized curve averaging and wall-clock time-to-target ratios between Muon and Drop-Muon.Normalized averaging rescales each Muon run to end at t = 1; time-to-target measures time to fixed accuracy thresholds.Source: §G.3 · p. 62

G.4 Results on MNIST

On MNIST, Drop-Muon with uniform layer sampling can have lower accuracy per epoch yet higher accuracy at the same wall-clock time. Multi-seed time-to-target results show speedups of up to 1.4× across several accuracy thresholds.

G.4 Results on MNIST

  • For training times up to approximately 150 seconds, Drop-Muon consistently achieved higher accuracy than full-network Muon, despite lower accuracy per epoch initially.The reported wall-clock advantage applies to reaching 95% training accuracy or less.Source: §G.4 · p. 62
  • Up to 1.4× speedup was reported at 60%, 70%, 80%, and 99% accuracy thresholds, with a slightly lower speedup at 90%.The authors note variability across seeds but report practical acceleration across training.Source: §G.4 · p. 62
  • Figure 5 averages normalized curves from several MNIST runs using uniform index sampling.The caption specifies batch size 8192, learning rate 0.1, and channels [64, 128, 256].Source: Figure 5 · p. 63

G.5 Results on Fashion-MNIST

On Fashion-MNIST, Drop-Muon shows practical speedups while maintaining comparable final accuracy; on CIFAR-10, its speedup depends on the target accuracy. The discussion notes implementation trade-offs, seed variability, and online sampling as a future direction.

G.5 Results on Fashion-MNIST

  • On Fashion-MNIST, Drop-Muon progresses faster in wall-clock time early in training despite slightly lower per-epoch accuracy.Figure 7 presents normalized curve averaging across multiple seeds; smoother high-accuracy curves do not indicate reduced variance.Source: §G.5 · p. 63, Figure 7 · p. 64
  • Fashion-MNIST time-to-target results show approximately 1.2 speedups across considered accuracy thresholds, while neither method reaches 99% in the plotted runs.The authors report comparable final accuracy overall.Source: §G.5 · p. 63, Figure 7 · p. 64

G.6 Results on CIFAR-10

  • On CIFAR-10, Drop-Muon makes faster wall-clock progress despite slightly lower per-epoch accuracy; its normalized multi-seed curves end earlier.The experiments use epoch-shift sampling; the Figure 8 caption specifies batch size 8192, learning rate 0.1, and channels [128, 256, 512].Source: §G.6 · p. 64, Figure 8 · p. 64
  • CIFAR-10 time-to-target speedups are notable at 90% accuracy but minimal or absent at 60%, 70%, and 80%, attributed mainly to suboptimal hyperparameters.The stated results concern this configuration.Source: §G.6 · p. 64, Figure 9 · p. 65

G.7 Discussion and Practical Remarks

  • Drop-Muon is described as implementable in a few lines, but turning off gradient computations every iteration can be costly for large models.The discussion suggests sampling the cutoff layer less often to reduce overhead; when batch or model size is small, computing all gradients but orthogonalizing only a subset may be advantageous.Source: §G.7 · pp. 64–65
  • The method shows non-negligible seed variance, particularly on smaller datasets; online learning of sampling distributions is proposed as future work.Stabilization mechanisms are also identified as a possible way to mitigate variance.Source: §G.7 · p. 65
Loading 2510.02239v1…