OptMuon: Closed-Loop Orthogonalized Momentum Methods for Stochastic Optimization with Zero-Noise Optimality

Technical Summary of "OptMuon: Closed-Loop Orthogonalized Momentum Methods for Stochastic Optimization with Zero-Noise Optimality"

1. Problem Statement

The paper addresses the problem of stochastic first-order optimization for matrix-valued nonconvex objectives:
$$\min_{X \in \mathbb{R}^{m \times n}} f(X) := \mathbb{E}_{\xi \sim \mathcal{D}}[f(X; \xi)]$$
The goal is to find a random iterate $X_{out}$ such that the expected Frobenius norm of the gradient, $\mathbb{E}[\|\nabla f(X_{out})\|_F]$, is minimized within $T$ stochastic gradient calls.

While Muon-style optimizers have demonstrated empirical stability by transforming momentum matrices into approximately orthogonal directions (via polar factorization), existing methods typically pair these directions with open-loop or constant magnitude rules. These rules often rely on prescribed smoothness constants ($L$) or variance levels ($\sigma^2$), making them sensitive to hyperparameter tuning and unable to explicitly calibrate update magnitudes based on the observed optimization trajectory. Conversely, adaptive methods like STORM and AdaGrad-Norm offer "closed-loop" scalar adaptation and noise adaptivity but generally operate on raw gradients rather than orthogonalized momentum directions.

The central challenge addressed is whether closed-loop scalar adaptation can be successfully combined with Muon-style orthogonalized momentum to retain geometric stability while achieving noise adaptivity and zero-noise optimality without manual retuning.

2. Methodology: OptMuon

The authors propose OptMuon, a framework that decouples the update direction from the update magnitude.

• Direction (Orthogonalization): The update direction is the polar factor of the momentum matrix $M_t$, denoted as $\text{Orth}(M_t) = U W^\top$ (where $M_t = U \Sigma W^\top$ is the SVD). This ensures the direction is geometry-normalized and separates orientation from singular-value scaling.
• Magnitude (Closed-Loop Adaptation): The step size is determined by a scalar coefficient $\gamma_t$ derived from the observed trajectory (gradient and momentum history), rather than fixed problem constants.
  • Coefficient Schedule: The core of the magnitude control is a running-maximum corrected AdaGrad-Norm schedule. It accumulates stochastic gradient energy in the denominator to suppress steps in noisy regimes but includes a numerator term ($\max g_i^2$) to prevent the coefficient from collapsing excessively due to isolated gradient spikes.
  • Lipschitz-Free: The algorithm does not require knowledge of the smoothness constant $L$, variance $\sigma^2$, or bounded-gradient constant $\hat{G}$ for parameter selection, although these appear in the theoretical constants.

Two Variants:

1. OptMuon-A (Average Smoothness): Designed for the average smoothness assumption. It uses one stochastic gradient per iteration and a standard momentum accumulation ($M_t = G_t + (1-\alpha_t)M_{t-1}$). The exponent in the schedule is $q=1/2$.
2. OptMuon-I (Individual Smoothness): Designed for the individual smoothness assumption. It employs a STORM-type recursive momentum estimator ($M_t = (1-\alpha_t)(M_{t-1} - \nabla f(X_{t-1}; \xi_t)) + G_t$), utilizing two gradient evaluations on the same sample to reduce variance. The exponent is $q=2/3$.

Both variants update the parameters via:
$$X_{t+1} = X_t - \theta \gamma_t \|M_t\|_F \text{Orth}(M_t)$$

3. Key Contributions

The paper claims the following specific contributions:

• Closed-Loop Momentum Orthogonalization: A framework that cleanly separates matrix geometry (via the polar factor) from magnitude calibration (via a self-normalized scalar schedule).
• Running-Maximum Corrected Schedule: A novel coefficient design that preserves the self-normalized suppression of AdaGrad-Norm while mitigating the risk of coefficient collapse caused by single gradient spikes.
• Noise-Adaptive Convergence Rates: Theoretical guarantees showing that the stochastic terms in the convergence bounds improve automatically as the variance $\sigma$ decreases.
• Zero-Noise Optimality: The demonstration that when $\sigma = 0$, the bounds automatically reduce to the nearly optimal deterministic rate $\tilde{O}(T^{-1/2})$ without any hyperparameter retuning.

4. Theoretical Results

Under assumptions of lower-boundedness, unbiased stochastic gradients with bounded variance, smoothness (average or individual), and almost-sure bounded stochastic gradients, the paper proves:

• OptMuon-A (Average Smoothness): Achieves a convergence rate of:
  $$\tilde{O}\left(T^{-1/2} + \sigma^{1/2}T^{-1/4}\right)$$
• OptMuon-I (Individual Smoothness): Achieves a sharper stochastic rate of:
  $$\tilde{O}\left(T^{-1/2} + \sigma^{1/3}T^{-1/3}\right)$$

Significance of Results:

• Zero-Noise Recovery: In the zero-noise regime ($\sigma = 0$), both bounds reduce to $\tilde{O}(T^{-1/2})$, matching the optimal deterministic first-order rate. This contrasts with existing Muon analyses that often yield $\tilde{O}(T^{-1/4})$ or require fixed step sizes.
• Noise Adaptivity: The stochastic terms ($\sigma^{1/2}T^{-1/4}$ and $\sigma^{1/3}T^{-1/3}$) are adaptive; they tighten as the noise level decreases, unlike methods with fixed step sizes.
• Comparison: The paper notes that OptMuon-A improves upon the $\tilde{O}(T^{-1/4})$ rate established in recent variance-reduced Muon variants (e.g., Chang et al., 2025) in the low-noise regime, while maintaining the zero-noise optimality.

5. Significance and Limitations

Significance:
The paper establishes that closed-loop scalar adaptation can be coupled with Muon-style orthogonalized momentum without sacrificing geometric stability. It extends the concept of "zero-noise adaptivity"—previously found in scalar adaptive stepsize rules—to polar-factor momentum updates. This allows for robust optimization that automatically adjusts to noise levels and recovers optimal deterministic performance without manual intervention.

Limitations (as stated by the authors):

• Exact Orthogonalization: The theory assumes exact polar-factor computation. While Newton-Schulz iterations can approximate this, the current analysis does not quantify the error introduced by inexact orthogonalization.
• Bounded Gradients: The analysis relies on an almost-sure bounded stochastic gradient condition (Assumption 4). The paper does not claim a bounded-gradient-free theory.
• Empirical Validation: The paper focuses on theoretical convergence guarantees and does not include large-scale empirical evaluations or extensions to composite/constrained objectives, which are noted as directions for future work.