Constrained Stochastic Spectral Preconditioning Converges for Nonconvex Objectives

Abstract: In this work, we develop proximal preconditioned gradient methods with a focus on spectral gradient methods providing a proximal extension to the Muon and Scion optimizers. We introduce a family of stochastic algorithms that can handle a wide variety of convex and nonconvex constraints and study its convergence under heavy-tailed noise, through a novel analysis tailored to the geometry of the proposed methods. We further propose a variance-reduced version, which achieves faster convergence under standard noise assumptions. Finally, we show that the polynomial iterations used in Muon are more accurately captured by a nonlinear preconditioner than by the ideal matrix sign, leading to a convergence analysis that more faithfully reflects practical implementations.

• The paper offers a novel proximal preconditioning framework that unifies and extends spectral, normalized, and sign gradient methods to nonconvex and constrained settings.
• It establishes convergence guarantees under heavy-tailed noise with an O(K^(-1/4)) rate for stochastic updates and O(K^(-1/3)) for variance-reduced variants, without requiring prior smoothness or noise constants.
• Empirical results on benchmarks from CIFAR10 to NanoGPT demonstrate that enforcing spectral and Euclidean constraints improves training stability and generalization in deep learning models.

Constrained Stochastic Spectral Preconditioning for Nonconvex Objectives: A Technical Synthesis

Problem Context and Algorithmic Foundations

The paper develops a new class of proximal preconditioned gradient methods targeting composite stochastic optimization with nonconvex, potentially nonsmooth regularization. The principal innovation lies in extending the anisotropic proximal gradient framework—originally limited to unconstrained or convex-constrained settings—to problems involving nonconvex constraints, prominent in the context of machine learning regularization (e.g., spectral and Euclidean norm constraints). The update rule analytically generalizes normalized gradient and spectral preconditioned methods (e.g., Muon, Scion), employing a nonlinear dual-space preconditioner derived from spectral reference functions. This provides a geometric grounding for weight-normalization updates observed in state-of-the-art optimizers.

For an objective $F(x) = f(x) + g(x)$, where $f$ is generally nonconvex and $g$ possibly nonsmooth and nonconvex, the algorithm iterates:

• Forward step: $y_k = x_k - \gamma_k \nabla \varphi^*(d_k)$,
• Proximal step: $$x_{k+1} = \arg\min_{x} g(x) + \gamma_k \varphi^*(x - y_k)$$, where $\varphi$ is a strongly convex, even reference function and $d_k$ is a (potentially stochastic) momentum term.

This construction yields a unifying lens: classic SGD, normalized gradient, sign gradient, and modern spectral methods are embedded for particular choices of $\varphi$. The extension to matrix-valued updates (spectral isotropic/anisotropic reference functions) allows the inclusion of layerwise and spectral-constraint optimizers previously only analyzed in special cases.

Convergence Analysis under Generalized Noise Models

A critical technical achievement is establishing convergence guarantees for the stochastic variant of these preconditioned proximal algorithms in settings with heavy-tailed gradient noise—relaxing the standard bounded-variance assumption to finite $p$-th central moments for $p \in (1,2]$. Both momentum-based and variance-reduced (STORM) directions are considered. The authors introduce a novel geometrically-tailored suboptimality gap,

$f$0

with $f$1, and establish that the minimal expected gap over $f$2 iterations decays as $f$3 (Algorithm 1, Theorem 4.3), with an improved rate $f$4 for the variance-reduced variant (Theorem 4.5). Notably, these rates do not require prior knowledge of problem smoothness or noise constants. The analysis departs from typical majorization-minimization arguments, instead leveraging convex analysis and a new regularized gap function to accommodate generic nonconvex $f$5 and anisotropic preconditioning.

The authors explicitly address the bias introduced by nonlinear preconditioning in stochastic updates, circumventing the need for large-batch asymptotics and showing resilience in the presence of heavy-tailed noise—a notable advantage over classical stochastic algorithms.

Nonlinear Spectral Preconditioning and Empirical Alignment

A pivotal theoretical insight is that polynomial spectral iterations as used in practical implementations (e.g., Polar Express updates in Muon) are more faithfully modeled as nonlinear preconditioners rather than as idealized matrix sign operations. This aligns theoretical analysis with empirical practice, capturing discrepancies—especially the pre-normalization steps necessary for polynomial convergence. The deterministic and stochastic cases are rigorously treated, demonstrating that practical convergence behavior (rate and stationarity measures) is accurately reflected by the nonlinear spectral preconditioner framework.

Matrix and Layerwise Extensions

By exploiting the theory of orthogonal invariance and absolute symmetry, the matrix extension reduces the high-dimensional update and projection operations to the singular value domain. This enables efficient computation of backward steps for constraints such as the spectral ball, Stiefel manifold, Frobenius ball, and spectral sphere—commonly encountered in deep learning architectures. The layerwise setting naturally decomposes for network architectures with block-separable parameters, integrating seamlessly with the current design of neural network optimizers.

Numerical Results

The empirical section corroborates theoretical claims using benchmarks ranging from modular arithmetic and CIFAR10 to large-scale transformer training (NanoGPT). Experiments illustrate:

• Incorporating spectral or Frobenius norm constraints in the backward step reliably improves training stability and generalization, with consistent acceleration of phenomena such as grokking and improved test accuracy compared to unconstrained baselines.
• The predicted tighter control of iterates and robustness to noise, as suggested by the analysis, directly translates into improved empirical performance especially on non-convex, overparameterized models.

Theoretical and Practical Implications

The proposed unified framework addresses several open theoretical questions. It rigorously extends spectral preconditioning, normalized, and sign-based methods to composite, nonconvex, and non-Euclidean constraint settings previously lacking convergence proofs—especially under non-standard noise. Practically, this brings strong theoretical support for a range of modern training heuristics used in LLM pretraining and robust model design. The framework's geometric adaptivity obviates tuning for noise or smoothness parameters and naturally encompasses variant architectures and regularizations.

The explicit connection between polynomial spectral updates and nonlinear preconditioners paves the way for more faithful modeling and analysis of future adaptive optimizers, especially as optimization moves toward exploiting geometry and operator structure beyond the Euclidean paradigm.

Future Directions

Looking forward, several avenues are directly enabled:

• Extension of nonlinear spectral preconditioning to distributed, federated, or large-batch settings, leveraging the demonstrated robustness to heavy-tailed stochasticity.
• Design of new variance reduction or adaptive preconditioning schemes exploiting the flexibility of reference functions and their matrix-analytic counterparts.
• Exploration of proximal spectral methods in other nonconvex compositional settings such as generative models, control, and scientific machine learning, where geometry and constraints play essential roles.

Conclusion

This work closes a significant analytical gap in the theory of preconditioned, normalized, and spectral gradient methods for constrained, composite, nonconvex stochastic optimization. It provides a unifying, mathematically rigorous, and practically motivated framework that aligns with—and advances—the state of the art in large-scale nonconvex optimization and deep learning. The results decisively broaden the scope of spectral preconditioning and provide an extensible toolkit for both theory and applications in high-dimensional statistical learning and optimization.