PaperlayerView exact revision on arXiv

Source-linked AI summary

Understanding Gradient Orthogonalization for Deep Learning via Non-Euclidean Trust-Region Optimization

Dmitry Kovalev

arXiv:2503.12645v2cs.LGmath.OCstat.ML

TL;DR

Orthogonalized-gradient methods have shown promise for deep-network training, but their theoretical foundations remain limited. The paper interprets them as non-Euclidean spectral-norm trust-region methods and develops a stochastic momentum framework encompassing Muon and related optimizers. It proves convergence guarantees across non-convex, star-convex, and second-order-smooth settings, including improved rates and explanations for practical observations about Muon and weight decay.

  • Problem

    The paper addresses the limited theoretical understanding of orthogonalized-gradient optimization despite its empirical progress in training deep neural networks and language models.

  • Method

    The paper develops stochastic non-Euclidean trust-region gradient methods with momentum, weight decay, and extrapolation for composite optimization under arbitrary non-Euclidean norms.

  • Results

    The framework exactly recovers orthogonalized gradient updates under the matrix spectral norm, includes Muon and related optimizers, and proves 1/ε4, 1/ε3, and 1/ε3.5 complexity results in supported settings.

  • Takeaways & Limitations

    The theory supports explanations for Muon’s practical advantage over Orthogonal-SGDM and for the importance of weight decay in large-scale language-model training.

  • Takeaways & Limitations

    For star-convex functions, the extrapolation method does not substantially improve the purely stochastic complexity term, which remains (ρσD/ε)3.

Abstract

from arXiv · show

Optimization with matrix gradient orthogonalization has recently demonstrated impressive results in the training of deep neural networks (Jordan et al., 2024; Liu et al., 2025). In this paper, we provide a theoretical analysis of this approach. In particular, we show that the orthogonalized gradient method can be seen as a first-order trust-region optimization method, where the trust-region is defined in terms of the matrix spectral norm. Motivated by this observation, we develop the stochastic non-Euclidean trust-region gradient method with momentum, which recovers the Muon optimizer (Jordan et al., 2024) as a special case, along with normalized SGD and signSGD with momentum (Cutkosky and Mehta, 2020; Sun et al., 2023). In addition, we prove state-of-the-art convergence results for the proposed algorithm in a range of scenarios, which involve arbitrary non-Euclidean norms, constrained and composite problems, and non-convex, star-convex, first- and second-order smooth functions. Finally, our theoretical findings provide an explanation for several practical observations, including the practical superiority of Muon compared to the Orthogonal-SGDM algorithm of Tuddenham et al. (2022) and the importance of weight decay in the training of large-scale language models.

1 Introduction

The paper develops theoretical foundations for orthogonalized-gradient optimization by framing it as non-Euclidean trust-region optimization and analyzing stochastic composite problems. It derives convergence guarantees across multiple geometric and function classes while connecting the framework to Muon and related optimizers.

  • Motivation: The work is motivated by empirical progress using orthogonalized gradients to surpass AdamW in training large language models.Muon has been empirically reported to outperform AdamW, but the paper identifies limited theoretical understanding of its effectiveness.
  • Problem setting: The analysis targets stochastic composite optimization with a differentiable objective and a proper, closed, convex regularizer under non-Euclidean smoothness assumptions.The stochastic gradient estimator is assumed unbiased with bounded variance, while gradient Lipschitzness is measured using a potentially non-Euclidean norm and its dual.
  • Gradient orthogonalization as non-Euclidean trust-region optimization: Orthogonalized gradient updates are interpreted exactly as trust-region gradient steps under the matrix spectral norm.For matrix optimization, the spectral and nuclear norms make the non-Euclidean trust-region method equivalent to orthogonalized gradient descent.
  • Stochastic non-Euclidean trust-region gradient method with momentum: The proposed stochastic non-Euclidean trust-region method with momentum supports arbitrary non-Euclidean norms and recovers Muon, normalized SGD, and signSGD with momentum.For non-convex functions, the method has 1/ε4 iteration complexity and provides the first convergence result for Muon.
  • Convergence analysis for star-convex functions: Weight decay yields 1/ε3 iteration complexity under star-convexity, supporting its importance for Muon training of large-scale language models.Without weight decay, the same complexity requires boundedness of dom R; the paper also notes possible lesser importance for small-scale language models.
  • Convergence analysis for second-order smooth functions: With second-order smoothness, an extrapolation method achieves 1/ε3.5 complexity for non-convex functions, improving on the standard 1/ε4 complexity for SGD-type methods.The method also improves the star-convex guarantees obtained for the preceding algorithms.

2 Non-Euclidean Trust-Region Gradient Method

The paper contrasts majorization-minimization with a first-order trust-region method that minimizes a local gradient approximation inside a non-Euclidean norm ball. For matrix optimization, choosing spectral and nuclear norms makes this method exactly equivalent to gradient orthogonalization and supports convergence guarantees.

  • Non-Euclidean Trust-Region Gradient Method: The first-order trust-region method minimizes a gradient-based approximation within a radius-η ball under a chosen non-Euclidean norm.Unlike the majorization-minimization approximation, the trust-region approximation need not majorize the objective because it lacks squared-norm regularization.
  • Non-Euclidean Trust-Region Gradient Method: With the Euclidean norm and zero regularization, the trust-region iterations reduce to normalized gradient descent.The Euclidean specialization uses ∥·∥=∥·∥*=∥·∥2 and R≡0.
  • Non-Euclidean Trust-Region Gradient Method: For general non-Euclidean norms and arbitrary proper closed convex regularizers, prior convergence results were unavailable for both convex and non-convex objectives.The paper provides results for non-convex functions and separately for star-convex functions.
  • Gradient Orthogonalization for Matrix Function Optimization: Using the spectral norm for updates and the nuclear norm as its dual makes the trust-region method exactly equivalent to the orthogonalized gradient method.This establishes the paper’s central matrix-optimization connection and differs from interpreting orthogonalization as ordinary gradient descent under the same norm.
  • Convergence Analysis: The resulting deterministic method has convergence guarantees for non-convex composite objectives under the stated gradient-Lipschitz assumptions.The guarantees are expressed using generalized stationarity, measured by ∥∇f(xk)+∇̂Rk∥*.

3 Stochastic Non-Euclidean Trust-Region Gradient Method with Momentum

The paper extends the non-Euclidean trust-region method to stochastic optimization by replacing the gradient with a momentum estimate. Different norm choices recover normalized SGD with momentum, signSGD with momentum, and Muon, while the analysis establishes state-of-the-art non-convex complexity.

  • The Algorithm: Algorithm 1 replaces the gradient in the trust-region update with a momentum term and analyzes the resulting stochastic method for non-convex and star-convex functions.The method addresses non-Euclidean norms and nonzero regularizers beyond the scope of prior normalized-SGD analysis.
  • The Algorithm: Euclidean, infinity, and spectral norm choices recover normalized SGD with momentum, signSGD with momentum, and Muon, respectively.The spectral-norm case is the Muon specialization.
  • Convergence Analysis for Non-Convex Functions: Theorem 2 provides a stochastic convergence bound for Algorithm 1 under Assumptions (A1)–(A3).The bound uses a generalized stationarity measure for the composite objective.
  • Convergence Analysis for Non-Convex Functions: 1/ε4 iteration complexity matches state-of-the-art SGD-type results and cannot be improved under Assumptions (A1)–(A3).The result is stated for the generalized expected stationarity guarantee of Corollary 2.
  • The Algorithm: Muon is exactly recovered by using the spectral norm with zero regularization, whereas Orthogonal-SGDM was reported to perform worse than Muon or well-tuned SGD.The comparison concerns the ordering of momentum and orthogonalization in the two optimizers.

4 Algorithms with Weight Decay for Star-Convex Functions

For star-convex problems, the paper shifts the trust-region center toward zero to obtain bounded iterates while preserving convergence, yielding a stochastic method with weight decay. With suitable norms, the method recovers normalized SGD, signSGD, and Muon, and achieves improved stochastic complexity.

  • Method with weight decay: The shifted method moves the trust-region center toward zero, ensuring bounded iterations without materially changing the original trust-region convergence properties.The original and shifted trust-regions have a non-zero intersection.
  • Method with weight decay: Algorithm 2 applies the same shift to the stochastic method and is called the stochastic non-Euclidean trust-region gradient method with weight decay.It uses momentum through the stochastic gradient update defined in the algorithm.
  • Connections to existing optimizers: Choosing different norms reduces Algorithm 2 to normalized SGD, signSGD, and Muon with momentum and weight decay.With the spectral norm and zero regularizer, it exactly matches Muon with weight decay.
  • Convergence guarantees: 1/ε3 stochastic iteration complexity improves over the 1/ε3.5 result previously obtained for normalized SGD and signSGD with momentum.The improvement is established for the star-convex setting under the stated theorem assumptions.
  • Convergence guarantees: Both deterministic and stochastic corollaries recover the standard deterministic convex complexity O(LD2/ε), up to logarithmic factors.This complexity is stated as unimprovable in the general non-Euclidean norm setting.
  • Practical interpretation: For Muon, the theory links weight decay to bounded iterates and the improved 1/ε3 complexity.The paper reports that without weight decay, Muon iterates grow too large and its gains over AdamW diminish.

5 Algorithms with Extrapolation for Second-Order Smooth Functions

The section develops an extrapolated stochastic non-Euclidean trust-region method for second-order smooth objectives. Under second-order smoothness, it improves non-convex iteration complexity to 1/ε3.5, while star-convex improvements remain limited.

  • Algorithm 3: The extrapolated method replaces the stochastic gradient used in momentum with a gradient computed at an extrapolated point.This modification supports convergence analysis for arbitrary non-Euclidean norms and regularizers in non-convex and star-convex settings.
  • Muon connection: Algorithm 3 with the matrix spectral norm becomes a variant of Muon with extrapolation.For star-convex functions, the method uses weight decay together with extrapolation.
  • Non-convex convergence: 1/ε3.5 iteration complexity is obtained for non-convex functions under an additional second-order smoothness assumption.This matches normalized SGD and signSGD results and improves the standard 1/ε4 complexity for SGD-type methods.
  • Star-convex convergence: Star-convex complexity retains the term (ρσD/ε)3 shared with Algorithm 2, limiting major improvement over that method.The mixed term ρ4σ4HD7/ε5 may nevertheless improve over ρ2σ2LD3/ε3, and whether extrapolation improves Muon remains open.
  • Spectral-norm setting: The spectral norm is motivated because matrix-vector products satisfy ∥Xai∥2 ≤ ∥X∥op∥ai∥2, supporting the required first- and second-order smoothness bounds.The paper proves the Hessian-Lipschitz condition for the matrix spectral norm when component losses have Lipschitz Hessians.

A.3 Complexities for Non-Convex and Star-Convex Functions

This appendix clarifies which accuracy notions underpin the reported complexity comparisons. It also notes that star-convex stationarity-to-suboptimality implications require care, especially with zero regularization.

  • Accuracy notions: Non-convex comparisons use approximate first-order stationarity, whereas star-convex comparisons use approximate functional suboptimality.The two criteria are ∥∇f(xk) + ˆ∇Rk∥* ≤ ε and F(xk) − F(x*) ≤ ε, respectively.
  • Complexity comparison: 1/ε3 complexity for ε-stationarity is available under star-convexity and bounded-domain assumptions.This makes the result directly comparable with non-convex stationarity complexities.
  • Stationarity implication: Generalized (ε/D)-stationarity implies ε-approximate functional suboptimality under the stated conditions.The text cautions that this consideration is not valid, strictly speaking, for zero regularization R ≡ 0.

B Muon vs Orthogonal-SGDM

Muon and Orthogonal-SGDM differ in whether momentum or orthogonalization is applied first. The theory favors Muon because it preserves a variance-reducing momentum interpretation and feeds momentum into the trust-region step.

  • Trust-region interpretation: The paper identifies Muon with Algorithm 1 in the spectral-norm, zero-regularization setting.This connects the optimizer comparison to the non-Euclidean trust-region framework.
  • Update rules: Muon updates momentum from the gradient, orthogonalizes the momentum, and applies the orthogonalized update.Orthogonal-SGDM instead orthogonalizes the gradient before updating momentum and uses the resulting momentum for the parameter update.
  • Update rules: The main distinction between the methods is the order of momentum accumulation and orthogonalization.Both algorithms are presented in the comparison pseudocode.
  • Variance reduction: Muon’s momentum accumulates weighted gradient-estimation errors whose total variance is reduced by a batching effect.The argument uses unbiased, independent, centered errors with bounded variance.
  • Trust-region interpretation: The trust-region interpretation naturally feeds momentum into the optimization step, yielding orthogonalized momentum rather than momentum of orthogonals.This is the paper’s theoretical reason for preferring Muon’s update order.

C Algorithms with Weight Clipping for Star-Convex Functions

For star-convex problems, bounded domains can replace weight decay through weight clipping. This yields matching complexity guarantees, but spectral-norm clipping is difficult to implement, supporting weight decay for Muon.

  • Bounded-domain assumption: Boundedness of dom R supplies an alternative route to convergence guarantees for star-convex objectives without weight decay.The domain diameter D enters the resulting bounds.
  • Complexity guarantees: The deterministic and stochastic weight-clipping complexities match the corresponding weight-decay complexities.This comparison is stated for the iterations and Algorithm 1 versus their weight-decay counterparts.
  • Scope boundary: The bounded-domain assumption does not hold when the regularizer is identically zero.The paper proposes an indicator regularizer as one way to impose boundedness.
  • Algorithmic special case: Under the infinity norm, the clipped algorithm reduces to signSGD with momentum and weight clipping.Weight clipping is described as an alternative to weight decay.
  • Practical implications: Spectral-norm clipping is not easily implementable, which supports using weight decay in Muon.For small-scale or short-training language-model settings, clipping may be skipped if iterates remain within a reasonable radius.

D.1 Proof of Theorem 1

The proof of Theorem 1 applies the regularized trust-region lemma to bound the next objective value, using smoothness, norm inequalities, and the trust-region step constraint.

  • Trust-region bound: Lemma 3 incorporates an arbitrary proper, closed, convex regularizer into the non-Euclidean trust-region step.The theorem proof applies it with Ψ(x) = ⟨∇f(xk), x⟩ and center/update pair (xk, xk+1).
  • Objective analysis: The proof upper-bounds F(xk+1) using the definition of F, smoothness assumption (A2), Lemma 3, and the triangle inequality.The resulting argument combines the linearized objective with regularization and distance-control terms.
  • Trust-region control: The update satisfies ∥xk+1 − xk∥ ≤ η, as implied by eq. (9), enabling the trust-region radius to control the proof terms.This bound is explicitly used when applying the smoothness and regularizer inequalities.

D.1.1 Proof of Lemma 3

Lemma 3 is established by reformulating the constrained regularized problem with an indicator function, proving existence, and deriving the relevant subgradient optimality condition.

  • Existence: A solution z+ exists because the objective is lower semi-continuous and the constrained set is compact.Proper, closed, convex regularity of R supports lower semi-continuity.
  • Problem reformulation: The constrained problem is rewritten using the indicator Iη and combined regularizer Rη, which remains proper, closed, and convex.This converts the trust-region constraint into a composite convex optimization problem.
  • Optimality condition: The proof establishes −∇Ψ(z+) ∈ ∂Rη(z+) by contradiction using differentiability and convexity.Assuming the opposite yields an inequality contradicting δ > 0 and the derived δ ≤ 0.
  • Subgradient decomposition: The subdifferential decomposes as ∂Rη(z+) = ∂R(z+) + ∂Iη(z+) under the relative-interior intersection condition.The proof constructs a sequence in ri dom R approaching the trust-region center to verify that condition.
  • Final inequality: The resulting inequality follows from convexity of R, the Fenchel–Young equality, and the indicator definition.These steps complete the bound required by Lemma 3.

D.2 Proof of Theorem 2

Theorem 2 combines a trust-region descent inequality with momentum-error bounds to analyze stochastic non-Euclidean optimization with momentum.

  • Descent control: Lemma 4 provides a descent-lemma adaptation for the trust-region iterations under smoothness assumption (A2).It supplies the objective-progress inequality used in the theorem proof.
  • Momentum control: Lemma 5 upper-bounds the expected dual-norm distance between the momentum mk and the true gradient ∇f(xk).Its proof adapts the momentum analysis of Cutkosky and Mehta to arbitrary non-Euclidean norms.
  • Combination: Combining Lemmas 4 and 5 yields the theorem’s convergence inequality for the stochastic trust-region method.The proof first applies the descent bound and then substitutes the expected momentum-error estimate.
  • Objective bound: The theorem bounds F(xk+1) using smoothness, the trust-region step bound, the dual norm, and Lemma 3 applied to mk+1.The update ensures ∥xk+1 − xk∥ ≤ η.
  • Stochastic error analysis: The momentum discrepancy mk+1 − ∇f(xk) is expanded through the momentum recursion and bounded using triangle, smoothness, Jensen, variance, and independence arguments.The analysis applies assumptions (A1)–(A3) to control the expected error.

D.3 Proof of Theorem 3

Theorem 3 and related results are proved by controlling trust-region distances, momentum errors, and regularized objective recursions under star-convex and higher-order smoothness assumptions.

  • D.3 Proof of Theorem 3: Lemma 6 establishes the distance bounds ∥x − (1 − β)xk∥ ≤ η, ∥x − xk∥ ≤ 2η, ∥x − xk+1∥ ≤ 2η, and ∥xk+1 − xk∥ ≤ 2η.These bounds support subsequent analyses of Theorems 3 and 4, including Algorithm 2 and Algorithm 3.
  • D.3 Proof of Theorem 3: Theorem 3’s objective bound uses Lemma 6, smoothness, eq. (19), assumption (A4), and convexity of R.After rearrangement, the proof obtains the stated recursive inequality.
  • D.3 Proof of Theorem 3: The proof derives β∥xk∥ ≤ η by induction and then uses it with triangle inequalities to establish the auxiliary distance inequalities.The induction starts from eq. (22) and propagates the bound across iterations.
  • D.3 Proof of Theorem 3: For Algorithm 2, Lemma 7 provides a momentum-error inequality based on the same argument as Lemma 5, replacing the step bound η with 2η.Theorem 4 combines this estimate with Lemma 6 and star-convexity-related assumptions.
  • D.3 Proof of Theorem 3: Theorem 4’s recursion bounds F(xk+1) through smoothness, dual-norm control, Lemma 6, and convexity of R.The proof rewrites the inner-product term and applies the dual norm before using assumption (A4).
Loading 2503.12645v2…