alphaFeedback

WorksComputer Science

Finite Spectral Depth Controls Gaussian Entrance Rank and Avoidance of the First Gradient-Flow Saddle

AI contributions
Not recorded for this paper
View details

Abstract

Matrix factorization provides a simple model for studying the training dynamics of neural networks. Recent studies of gradient flow under small initialization have analyzed how trajectories escape from a neighborhood of the origin and approach their first nonzero stationary point. In this paper, we study continuous-time dynamics that applies a finite number of Newton-Schulz iterations to the normalized matrix gradient (Muon). We show that the number of Newton-Schulz steps can change the rank selected during the early stages of training and allow the resulting trajectory to bypass the first nonzero stationary point approached by gradient flow. Under an additional residual-inertia condition, we prove that bypass occurs for all sufficiently large finite Newton-Schulz depths and provide an explicit sufficient threshold. We also construct an example in which the bypass threshold is four.

Introduction

Muon and its variants have emerged as practical optimizers for large-scale language-model training [23] (Kimi Team 2025), using Newton–Schulz iterations to approximately orthogonalize matrix-valued updates [13]. Their design is closely related to steepest descent under the spectral norm, where the polar factor provides an idealized update direction [2]. In practice, however, orthogonalization is performed by a fixed, small number of polynomial iterations. This distinction is particularly relevant near rank deficiency: exact polar orthogonalization maps every nonzero singular value to one, whereas finite Newton–Schulz iterations have a bounded amplification factor near zero at each fixed depth. Analyses of idealized spectral updates have identified learning dynamics distinct from gradient descent in concrete matrix models [29] (Ma et al. 2026; Li et al. 2026), while finite Newton–Schulz updates have been studied through convergence guarantees and their smoothing effect [14] (Li and Tsuchiya 2026). These developments motivate a direct study of how the finite number of iterations shapes the optimization trajectory.

Polynomial iterations for matrix orthogonalization have a long history in numerical linear algebra [16, 4, 8, 9]. More recently, optimization analyses that explicitly retain finitely many Newton–Schulz iterations have investigated convergence to stationary points, the dependence on iteration depth, and the effects of stabilization and momentum [14] (Li and Tsuchiya 2026).

©2026 Yanlin Long.

License: CC-BY 4.0, see https://creativecommons.org/licenses/by/4.0/. 2026; Peng et al., [28]; Li and Tan (2026). These results provide a useful foundation for studying the dynamical consequences of a fixed-depth spectral transformation.

Gradient-based training of factorized models from small initialization can exhibit incremental learning, with different components becoming appreciable at successive stages [30, 10, 3, 12] (Li et al. 2021). In matrix factorization, this behavior can take the form of dominant directions being learned first, followed by weaker directions that increase the effective rank of the predictor. The trajectory may consequently pass near intermediate low-rank stationary points. Related analyses of homogeneous networks describe directional convergence near the origin and saddle points, as well as the subsequent trajectory toward the first saddle after escape [17, 18]. A complementary approach studies training dynamics in high-dimensional limits, extracting effective structure from the evolving parameters or deriving dynamical mean-field descriptions of gradient flow [25, 26]. Related work on feature learning and the distinction from lazy training provides another perspective on how different directions become accessible during optimization [6, 7].

Since Newton–Schulz iterations amplify small singular values of the normalized gradient, a natural question is whether a finite increase in the number of iterations can make weaker directions emerge early enough to change which stationary points the trajectory approaches. Answering this question requires tracking both singular-value amplification and the coupled evolution of different directions, as the gradient and its singular subspaces evolve with the factorized parameters.

We consider a two-layer neural network with quadratic activation and all second-layer weights fixed to one. This model admits an equivalent representation as a positive-semidefinite matrix factorization. Such quadratic networks provide a tractable setting for analyzing nonlinear training dynamics, and have also been studied through gradient flow and spectral gradient methods [25, 26, 5]. We study continuous-time optimization dynamics in which KK cubic Newton–Schulz iterations are applied to the normalized factor gradient, and refer to KK as the spectral depth.

For the analysis, we use a two-scale Gaussian initialization: one hidden row is initialized at scale ε\varepsilon, while the remaining rows have scale εb\varepsilon^{b}, where b>1b > 1. The gradient normalization includes an additive stabilization proportional to ε\varepsilon. This choice of initialization scales allows us to identify the leading directional dynamics first and then derive the coupled equations governing the weaker components in the remaining independent directions. We compare trajectories on the same task and with the same underlying Gaussian seed as ε\varepsilon tends to zero, while keeping each spectral depth fixed. Removing the row-dependent initialization scales would instead lead to an isotropic Gaussian initialization, for which the competing directions may exhibit different and potentially richer dynamics. The analysis of this regime remains open.

Another question concerns the extent to which small-initialization asymptotics can determine the subsequent training trajectory. Under suitable structural conditions, gradient flow admits a precise description in which the initialization scale can be used to reverse-engineer the full training dynamics, leading also to a characterization of time-dependent implicit regularization [3]. For finite-depth Newton–Schulz dynamics, the nonlinear matrix iteration prevents this reconstruction argument from carrying over directly. Although the dynamics can be analyzed in a small neighborhood of the origin, the resulting local description does not by itself determine the trajectory after leaving that neighborhood. This creates two distinct dynamical scales: the entrance dynamics near the origin and the subsequent evolution away from it. Establishing how the entrance dynamics are transmitted to the later trajectory is one of the main results of this paper.

Following the framework of Li et al. (2021), which characterizes the first nonzero stationary point selected by gradient flow under suitable assumptions, we study the evolution from small initialization through the first stationary point encountered after escape. This allows us to compare the stationary points approached by gradient descent and by finite-depth Newton–Schulz dynamics. Under an additional structural assumption, we establish a sufficient large-depth guarantee: for all sufficiently large spectral depths, the selected Newton–Schulz continuation approaches a first stationary point whose loss is strictly smaller than that of the first stationary point selected by gradient descent. For smaller depths, our results do not provide a general criterion determining whether the trajectory visits or bypasses the gradient-descent stationary point. Moreover, without the additional assumption, the first stationary points reached by finite-depth Newton–Schulz dynamics and gradient descent may have the same loss. Such equal-loss behavior may reveal further dynamical structures, which we leave for future work.

The rest of this paper is organized as follows. § 2 introduces the model and pathwise comparison criteria. § 3–§ 5 establish the Gaussian entrance law, continuation criteria, and eventual depth guarantee. § 6 presents an exact empirical task, and § 7 discusses the scope of the results and open questions. Detailed proofs are given in the appendices.

Notations. We write Sr\mathbb{S}^{r} for the real symmetric r×rr \times r matrices and S+r\mathbb{S}_{+}^{r} for the positive-semidefinite cone. The matrix inner product is ⟨A,B⟩=tr⁡(ATB)\langle A, B\rangle= \operatorname{tr}(A^{\mathsf{T}}B). Subscripts F and op denote the Frobenius and operator norms, respectively; vector norms are Euclidean. For a symmetric matrix AA, n+(A)n_{+}(A) counts its strictly positive eigenvalues with multiplicity. We write O(m)={Q∈Rm×m:QTQ=Im}O(m) = \{Q \in\mathbb{R}^{m \times m} : Q^{\mathsf{T}}Q = I_{m}\}.

Main Results

In this section, we introduce the empirical quadratic model and finite-depth Newton–Schulz dynamics, and formulate the pathwise comparison with gradient flow from small initialization.

Empirical quadratic learning

Fix m≥2m \ge2 and r≥3r \ge3, and data {(xl,yl)}l=1n\{(x_l, y_l)\}_{l=1}^{n} with xl∈Rrx_l \in\mathbb{R}^{r} spanning Rr\mathbb{R}^{r} and yl∈Ry_l \in\mathbb{R}. Put q=min⁡{m,r}q = \min\{m,r\}. For W∈Rm×rW \in\mathbb{R}^{m \times r}, set

M=WTW,Φ(M)=12n∑l=1n(yl−xlTMxl)2,f(W)=Φ(WTW).(1)M = W^{\mathsf{T}}W,\qquad\Phi(M) = \frac{1}{2n}\sum_{l=1}^{n}(y_l - x_l^{\mathsf{T}} M x_l)^2,\qquad f(W) = \Phi(W^{\mathsf{T}}W). \tag*{(1)}

This is the predictor of the quadratic network x↦∥Wx∥22x \mapsto\lVert W x\rVert_{2}^{2} with fixed unit readouts. Define the task matrix, measurement operator, and residual by

C=1n∑lylxlxlT,L(H)=1n∑l(xlTHxl)xlxlT,R(M)=C−L(M).(2)C = \frac{1}{n}\sum_l y_l x_l x_l^{\mathsf{T}},\qquad\mathcal{L}(H) = \frac{1}{n}\sum_l (x_l^{\mathsf{T}} H x_l)x_l x_l^{\mathsf{T}},\qquad R(M) = C - \mathcal{L}(M). \tag*{(2)}

The gradient is

G(W):=−∇f(W)=2WR(WTW).G(W) := -\nabla f(W) = 2W R(W^{\mathsf{T}}W).

Finite spectral depth and two-scale initialization

We use the cubic Newton–Schulz polynomial iteration (Björck and Bowie, [4]; Higham, [8])

P0(X)=X,Pk+1(X)=32Pk(X)−12Pk(X)Pk(X)TPk(X).(3)P_0(X) = X,\qquad P_{k+1}(X) = \frac{3}{2}P_k(X) - \frac{1}{2}P_k(X)P_k(X)^{\mathsf{T}}P_k(X). \tag*{(3)}

On the unit Frobenius ball, PkP_k preserves singular vectors and applies the scalar polynomial pkp_k to each singular value, where

p0(s)=s,pk+1(s)=32pk(s)−12pk(s)3,aK=pK′(0)=(3/2)K.(4)p_0(s) = s,\qquad p_{k+1}(s) = \frac{3}{2}p_k(s) - \frac{1}{2}p_k(s)^3,\qquad a_K = p_K'(0) = (3/2)^K. \tag*{(4)}

Fix b>1b > 1, cη>0c_\eta> 0, and finite K∈N0={0,1,2,…}K \in\mathbb{N}_0 = \{0,1,2,\ldots\}. For a matrix U∈Rm×rU \in\mathbb{R}^{m \times r} with independent standard normal entries, initialize

WKε,b(0;U)=Dε,bU,Dε,b=diag⁡(ε,εb,…,εb),ηε=cηε.(5)W_K^{\varepsilon,b}(0;U) = D_{\varepsilon,b}U,\qquad D_{\varepsilon,b} = \operatorname{diag}(\varepsilon,\varepsilon^b,\ldots,\varepsilon^b),\qquad\eta_\varepsilon= c_\eta\varepsilon. \tag*{(5)}

The actual trajectory, in physical time tt, solves

dWdt=PK(G(W)∥G(W)∥F+ηε).(6)\frac{\mathrm{d}W}{\mathrm{d}t} = P_K\left(\frac{G(W)}{\lVert G(W)\rVert_F + \eta_\varepsilon}\right). \tag*{(6)}

The additive stabilization in the denominator follows the original Muon implementation, which uses a fixed numerical constant of 10−710^{-7} in gradient normalization (Jordan et al., [13]). The choice ηε=cηε\eta_\varepsilon= c_\eta\varepsilon introduces an additional structural assumption: the stabilization term is taken to be of the same order as the small-initialization scale. Different scaling choices may lead to different balances between the gradient norm and the stabilization term, and hence different rescaled dynamics. For simplicity, we adopt this linear scaling throughout the paper.

Throughout the paper, we fix a finite spectral depth KK before taking the limit ε↓0\varepsilon\downarrow0, with the same Gaussian seed UU.

Since pKp_K is odd, there is a polynomial qKq_K satisfying

pK(s)=sqK(s2),qK(0)=aK.p_K(s) = s q_K(s^2),\qquad q_K(0) = a_K.

For G∈Rm×rG \in\mathbb{R}^{m \times r} and η≥0\eta\ge0, define

HK,η(G)={(∥G∥F+η)PK(G/(∥G∥F+η)),∥G∥F+η>0,0,G=0, η=0.(7)H_{K,\eta}(G) = \begin{cases} (\lVert G\rVert_F + \eta)P_K\left(G/(\lVert G\rVert_F + \eta)\right), & \lVert G\rVert_F + \eta> 0,\\ 0, & G = 0,\ \eta= 0. \end{cases} \tag*{(7)}

We use τ\tau for the time variable in dW/dτ=HK,η(G(W))\mathrm{d}W/\mathrm{d}\tau= H_{K,\eta}(G(W)).

For the actual predictor, write

MKε,b(t;U)=WKε,b(t;U)TWKε,b(t;U).M_K^{\varepsilon,b}(t;U) = W_K^{\varepsilon,b}(t;U)^{\mathsf{T}}W_K^{\varepsilon,b}(t;U).

Continuation and the gradient-flow comparison

Multiplying the physical vector field by its positive denominator gives a convenient time parameter for continuation. The following bounds keep that change of time and all later finite-interval comparisons on the same ordered path.

Proposition 1 (Continuation of finite-depth dynamics). For a spanning empirical task, the following assertions hold in arbitrary positive factor dimensions. Fix finite KK and η≥0\eta\ge0. Every solution of

dWdτ=HK,η(G(W))\frac{\mathrm{d}W}{\mathrm{d}\tau}=H_{K,\eta}(G(W))

is uniquely defined for all τ≥0\tau\ge0 and converges to one stationary factor WsW_s. Every finite time preserves the initial rank. Every PSD sublevel of Φ\Phi, and every factor sublevel of ff, is compact. Along the solution,

−dfdτ≥∥∇f∥F2,∥dWdτ∥F≤aK∥∇f∥F,−dfdτ≥aK−1∥∇f∥F∥dWdτ∥F.(8)-\frac{\mathrm{d}f}{\mathrm{d}\tau}\ge\lVert\nabla f\rVert_F^2,\qquad\left\lVert\frac{\mathrm{d}W}{\mathrm{d}\tau}\right\rVert_F\le a_K\lVert\nabla f\rVert_F,\qquad-\frac{\mathrm{d}f}{\mathrm{d}\tau}\ge a_K^{-1}\lVert\nabla f\rVert_F\left\lVert\frac{\mathrm{d}W}{\mathrm{d}\tau}\right\rVert_F. \tag*{(8)}

For η>0\eta>0, the physical-time equation with denominator ∥G(W)∥F+η\lVert G(W)\rVert_F+\eta has the same ordered path, is defined for all t≥0t\ge0, converges to WsW_s, and satisfies

rank⁡W(t)=rank⁡W(0)(0≤t<∞).\operatorname{rank}W(t)=\operatorname{rank}W(0)\qquad(0\le t<\infty).

The change of time is dt/dτ=∥G(W)∥F+η\mathrm{d}t/\mathrm{d}\tau=\lVert G(W)\rVert_F+\eta. More generally, a bounded C1C^1 curve satisfying (8) for a real-analytic loss on a neighborhood of its closure converges to a stationary point. These conclusions also apply to one-row and projected empirical problems used below.

At depth zero, H0,η(G)=GH_{0,\eta}(G)=G. The corresponding predictor evolves by ordinary factor gradient flow. Its first small-initialization stage selects a rank-one trajectory whose stationary endpoint will remain fixed throughout the optimizer comparison. We further assume that the two largest eigenvalues of CC are distinct.

Assumption 2 (Task spectrum). The task matrix CC is nonsingular. Its signed eigenvalues and an orthonormal eigenbasis satisfy

Cvl=λlvl,λ1>λ2≥⋯≥λr,λ2>0.Cv_l=\lambda_l v_l,\qquad\lambda_1>\lambda_2\ge\cdots\ge\lambda_r,\qquad\lambda_2>0.

Spectral separation conditions are common in analyses of incremental learning in matrix factorization [11]. Here, strict gap between the two largest eigenvalues, allowing multiplicities among the remaining ones. Removing the leading spectral gap may give rise to other interesting dynamical structures, which we leave for future work.

References

  1. [1]P.-A. Absil, R. Mahony, and B. Andrews. Convergence of the iterates of descent methods for analytical cost functions. SIAM Journal on Optimization, 16(2):531–547, 2005. doi: 10.1137/040605266.DOI
  2. [2]J. Bernstein and L. Newhouse. Old optimizer, new norm: An anthology. In NeurIPS 2024 Workshop on Optimization for Machine Learning, 2024. URL https://arxiv.org/abs/2409.20325.
  3. [3]R. Berthier. Incremental learning in diagonal linear networks. Journal of Machine Learning Research, 24(171):1–26, 2023.arxiv.org/abs/2208.14673
  4. [4]Å. Björck and C. Bowie. An iterative algorithm for computing the best estimate of an orthogonal matrix. SIAM Journal on Numerical Analysis, 8(2):358–364, 1971. doi: 10.1137/0708036.DOI
  5. [5]G. Braun, H. Bao, W. Huang, and M. Imaizumi. Spectral gradient descent mitigates anisotropy-driven misalignment: A case study in phase retrieval. In Proceedings of the 43rd International Conference on Machine Learning, volume 306 of Proceedings of Machine Learning Research, pages 9687–9733, 2026.arxiv.org/abs/2601.22652
  6. [6]L. Chizat, E. Oyallon, and F. Bach. On lazy training in differentiable programming. In Advances in Neural Information Processing Systems, volume 32, 2019.arxiv.org/abs/1812.07956
  7. [7]Y. Dandi, F. Krzakala, B. Loureiro, L. Pesce, and L. Stephan. How two-layer neural networks learn, one (giant) step at a time. Journal of Machine Learning Research, 25(349):1–65, 2024. URL https://jmlr.org/papers/v25/23-1543.html.arxiv.org/abs/2305.18270
  8. [8]N. J. Higham. Computing the polar decomposition—with applications. SIAM Journal on Scientific and Statistical Computing, 7(4):1160–1174, 1986. doi: 10.1137/0907079.DOI
  9. [9]N. J. Higham and R. S. Schreiber. Fast polar decomposition of an arbitrary matrix. SIAM Journal on Scientific and Statistical Computing, 11(4):648–655, 1990. doi: 10.1137/0911038.DOI
  10. [10]A. Jacot, F. Ged, B. Şimşek, C. Hongler, and F. Gabriel. Saddle-to-saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity. arXiv preprint arXiv:2106.15933, 2021.arxiv.org/abs/2106.15933
  11. [11]L. Jiang, Y. Chen, and L. Ding. Algorithmic regularization in model-free overparameterized asymmetric matrix factorization. SIAM Journal on Mathematics of Data Science, 5(3):723–744, 2023. doi: 10.1137/22M1519833.DOI
  12. [12]J. Jin, Z. Li, K. Lyu, S. S. Du, and J. D. Lee. Understanding incremental learning of gradient descent: A fine-grained analysis of matrix sensing. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 15200–15238, 2023.arxiv.org/abs/2301.11500
  13. [13]K. Jordan, Y. Jin, V. Boza, J. You, F. Cesista, L. Newhouse, and J. Bernstein. Muon: An optimizer for hidden layers in neural networks. Technical blog, 2024. URL https://kellerjordan.github.io/posts/muon/.
  14. [14]G. Y. Kim and M.-h. Oh. Convergence of Muon with Newton–Schulz. In International Conference on Learning Representations, 2026.arxiv.org/abs/2601.19156
  15. [16]Z. Kovarik. Some iterative methods for improving orthonormality. SIAM Journal on Numerical Analysis, 7(3):386–389, 1970. doi: 10.1137/0707031.DOI
  16. [17]A. Kumar and J. Haupt. Early directional convergence in deep homogeneous neural networks for small initializations. arXiv preprint arXiv:2403.08121, 2024.arxiv.org/abs/2403.08121
  17. [18]A. Kumar and J. Haupt. Towards understanding gradient flow dynamics of homogeneous neural networks beyond the origin. Journal of Machine Learning Research, 26(269):1–83, 2025.arxiv.org/abs/2502.15952
  18. [23]J. Liu, J. Su, X. Yao, Z. Jiang, G. Lai, Y. Du, Y. Qin, W. Xu, E. Lu, J. Yan, Y. Chen, H. Zheng, Y. Liu, S. Liu, B. Yin, W. He, H. Zhu, Y. Wang, J. Wang, M. Dong, Z. Zhang, Y. Kang, H. Zhang, X. Xu, Y. Zhang, Y. Wu, X. Zhou, and Z. Yang. Muon is scalable for LLM training. arXiv preprint arXiv:2502.16982, 2025.arxiv.org/abs/2502.16982
  19. [25]S. Martin, F. Bach, and G. Biroli. On the impact of overparameterization on the training of a shallow neural network in high dimensions. In Proceedings of the 27th International Conference on Artificial Intelligence and Statistics, volume 238 of Proceedings of Machine Learning Research, pages 3655–3663, 2024. URL https://proceedings.mlr.press/v238/martin24a.html.arxiv.org/abs/2311.03794
  20. [26]S. Martin, G. Biroli, and F. Bach. High-dimensional analysis of gradient flow for extensive-width quadratic neural networks. Journal of Machine Learning Research, 27(136):1–182, 2026.arxiv.org/abs/2601.10483
  21. [27]B. S. Mityagin. The zero set of a real analytic function. Mathematical Notes, 107(3):529–530, 2020. Preprint arXiv:1512.07276.
  22. [28]H. Peng, H. Wang, and Y. Yu. Convergence of practical Muon with finite Newton–Schulz iterations and Nesterov momentum. arXiv preprint arXiv:2609.39595, 2026.arxiv.org/abs/2609.39595
  23. [29]M. Rhee, J. Simon, and D. Karkada. Muon learns balanced solutions in matrix factorization without slow saddle-to-saddle dynamics. arXiv preprint arXiv:2606.30509, 2026.arxiv.org/abs/2606.30509
  24. [30]A. M. Saxe, J. L. McClelland, and S. Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In International Conference on Learning Representations, 2014. URL https://arxiv.org/abs/1312.6120.
  25. [31]G. Teschl. Ordinary Differential Equations and Dynamical Systems, volume 140 of Graduate Studies in Mathematics. American Mathematical Society, 2012. URL https://www.mat.univie.ac.at/~gerald/ftp/book-ode/ode.pdf.DOI

Paper details

vlnc:W_sqfmc9mc
v1 · latestHistory (1)

published · Fresh — added to the Hub within the last 30 days

added ·

Usage

views · 0

downloads · 0

A view counts once per browser per day, a download each time the PDF is requested.

CC BY 4.0

Original source

  • Computer Science

Not recorded for this paper · View statement →

Looking for related papers…
Contents