Source-linked AI summary

When do spectral gradient updates help in deep learning?

Damek Davis, Dmitriy Drusvyatskiy

arXiv:2512.04299v2cs.LGmath.OCstat.ML

TL;DR

The paper asks when spectral updates should outperform Euclidean gradient methods. It compares their one-step loss decreases blockwise using gradient nuclear rank and activation stable rank. It reports low-stable-rank activations and dimension-dependent predicted advantages in specified regimes, while also identifying a counterexample.

  • Problem

    The paper asks in what regimes spectral updates should outperform standard Euclidean gradient methods for deep networks and transformers.

    Source: §1 · p. 2
  • Method

    The paper compares one-step Euclidean and spectral loss-decrease bounds for a matrix block using its incoming activation matrix and gradient.

    Source: §1 · p. 2, §1.2 · p. 6
  • Results

    After short burn-in in spiked random-feature models, gradient nuclear rank grows with dimension, yielding a dimension-dependent predicted spectral-update advantage.

    Source: §1 · pp. 2–3
  • Takeaways & Limitations

    Spectral updates are predicted to be most advantageous for blocks with low-stable-rank incoming features and gradients with large nuclear rank.

    Source: §1 · p. 3

    In a random-feature model with centered activations, the nuclear-rank condition fails and spectral updates are slower than Euclidean ones.

    Source: §1.4 · p. 15

Abstract

from arXiv · show

Spectral gradient methods, such as the recently popularized Muon optimizer, are a promising alternative to standard Euclidean gradient descent for training deep neural networks and transformers, but it is still unclear in which regimes they are expected to perform better. We propose a simple layerwise condition that predicts when a spectral update yields a larger decrease in the loss than a Euclidean gradient step. This condition compares, for each parameter block, the squared nuclear-to-Frobenius ratio of the gradient to the stable rank of the incoming activations. To understand when this condition may be satisfied, we first prove that post-activation matrices have low stable rank at Gaussian initialization in random feature regression, feedforward networks, and transformer blocks. In spiked random feature models we then show that, after a short burn-in, the Euclidean gradient's nuclear-to-Frobenius ratio grows with the data dimension while the stable rank of the activations remains bounded, so the predicted advantage of spectral updates scales with dimension. We validate these predictions in synthetic regression experiments and in NanoGPT-scale language model training, where we find that intermediate activations have low-stable-rank throughout training and the corresponding gradients maintain large nuclear-to-Frobenius ratios. Together, these results identify conditions for spectral gradient methods, such as Muon, to be effective in training deep networks and transformers.

1 Introduction

The section develops a layerwise account of when spectral updates are favored, relating activation stable rank and gradient geometry across layered networks and transformers. It also considers how centered gated features, including SwiGLU, can weaken that advantage and positions the criterion alongside prior work on gradient geometry and activation rank.

1 Introduction

  • Spectral updates are favored on blocks whose incoming activations have low stable rank and whose gradients have large nuclear rank.The layerwise comparison applies to standard MLPs and extends to transformer projections, using their incoming normalized hidden states or MLP post-activations.Source: §1 · pp. 2–3

1.1 Post-activation matrices have low-stable rank.

  • At Gaussian initialization, non-centered activations yield bounded stable rank: for ReLU, the stated moment-ratio bound is π, and the result propagates through MLP hidden layers.The result applies to fixed inputs and Gaussian weights; the MLP bound holds for mean-spike-inducing activations.Source: §1.1 · pp. 3–4
  • Token-indicator structure and stable-rank propagation yield width- and sequence-length-independent bounds for transformer hidden states and MLP activations at Gaussian initialization.The token-indicator stable rank is the inverse frequency of the most common token, and Gaussian embeddings preserve control up to constants.Source: §1.1 · p. 4

1.2 Impact of post-activation degeneracy on training algorithms

  • The introduction reports that low-stable-rank activations persist in synthetic and NanoGPT-scale training, while NanoGPT gradients can have large nuclear rank.The observations cover MLP post-activations, RMS-normalized states, and embedding/token-indicator behavior.Source: §1.1 · p. 5, §1.2 · p. 8
  • In random-feature regression, SpecGD significantly outperforms GD, especially early in training when gradient nuclear ranks are large.The comparison also shows that the condition favoring SpecGD holds along both GD and SpecGD trajectories.Source: §1.2 · p. 7, Figure 3 · p. 8

1.2.1 Multiple layers and transformer blocks

  • For coupled multilayer architectures, the proposed analysis controls the composite-loss Hessian by a feature-based term plus a smaller parameter-dependent term.The framework fixes preactivation matrices at each layer and obtains activations through twice-differentiable maps.Source: §1.2.1 · p. 9
  • Internal transformer blocks pair low-stable-rank activations with typically large gradient nuclear ranks, favoring spectral updates; embeddings and the language-model head have smaller predicted advantages.The embedding gradient-to-feature ratio is only moderately above one, while the head is intermediate; these comparisons are geometric heuristics, not precise hyperparameter recommendations.Source: §1.2.1 · pp. 12–13

1.3 When do we expect the nuclear rank of the gradient to be large?

  • After a short O(log d) burn-in, gradients have nuclear rank Ω(d) for Θ(d) iterations and at least d^(1−ε) for Θ(d log d) iterations.These windows form a constant fraction of the Θ(d log(d/δ)) steps Euclidean gradient descent needs to reach relative error δ, so restarting SpecGD there gives a dimension-dependent one-step improvement.Source: §1.3 · p. 14

A “ SwiGLUpW1X, W2Xq :“ SiLUpW1Xq ˝ pW2Xq,

  • Centered SwiGLU random features lack the mean-induced spike, have stable rank growing with width and modest gradient nuclear rank, and yield faster objective reduction for GD than SpecGD.The nuclear-rank condition is violated in this example; sparse-regression SpecGD can still work at small batch sizes, but its advantage fades as batch size and stable rank increase.Source: §1.4 · p. 15, A “ SwiGLUpW1X, W2Xq :“ SiLUpW1Xq ˝… · pp. 16–17
  • With a SwiGLU MLP in transformers, internal activations remain low-stable-rank, though typically higher than in the squared-ReLU run, while corresponding gradient nuclear ranks remain large.The authors attribute this to low-rank token-indicator structure at the input being preserved through attention and MLP blocks, with width- and sequence-length-independent stable-rank control at initialization.Source: A “ SwiGLUpW1X, W2Xq :“ SiLUpW1Xq ˝… · p. 17
  • For non-gated MLPs, validation loss was 3.9242 (All Adam), 3.7023 (Muon on Wout only), and 3.5654 (full Muon); gated losses were 3.8837, 3.7843, and 3.5125, respectively.In the gated case, those values correspond to All Adam, Muon on Wout only, and Muon on all attention and MLP weights; Wout alone captured about 27% of the full gain, versus about 62% for non-gated MLPs.Source: Table 3 · p. 17, A “ SwiGLUpW1X, W2Xq :“ SiLUpW1Xq ˝… · p. 17
  • In SwiGLU random-feature regression, SpecGD is slower than GD because centered post-activations lack a mean-induced spike and have stable rank that grows with width.The resulting gradient nuclear-rank-to-activation-stable-rank condition is more easily violated; the proposed explanation for gated transformers remains unverified and requires additional ablations.Source: Figure 11 · p. 18, A “ SwiGLUpW1X, W2Xq :“ SiLUpW1Xq ˝… · p. 18

1.5 Related literature

  • The layerwise comparison complements Shen et al. [47] by replacing their Hessian ratio with post-activation stable rank, while retaining the gradient nuclear-rank ratio.The paper does not claim novelty for gradient nuclear rank controlling spectral-update advantages; it characterizes regimes where low activation stable rank and dimension-dependent gradient nuclear rank make the condition hold.Source: §1.5 · p. 19
  • Relative to Huh et al., the paper bounds layerwise activation stable rank theoretically and verifies low stable rank across intermediate MLP and NanoGPT layers throughout training.Huh et al. studied low effective rank in final-layer embeddings; the paper’s results extend the scope to intermediate layers and the optimization trajectory.Source: §1.5 · p. 20

2 Low-stable-rank activations: mean spikes and token-embeddings

This section examines how mean-induced spikes and frequent-token embeddings can produce low-stable-rank activations, developing concentration results and bounds for random-data polynomial features. It also gives stable-rank bounds for Gaussian random-weight post-activations and token embeddings, with the token-indicator stable rank equal to the inverse of the most common token’s frequency.

2 Low-stable-rank activations: mean spikes and token-embeddings

  • For transformers, token-indicator matrices and initial token embeddings have low stable rank when a single token constitutes at least a constant fraction of the sequence.The section presents mean-induced spikes and frequent-token embeddings as base cases for later arguments on low-stable-rank post-activation matrices.Source: §2 · p. 21

2.1 Concentration of the stable rank.

  • Under a small-ball condition, the sample stable rank is bounded by a constant multiple of population effective rank with constant probability, and it converges almost surely to that rank as sample size grows.With light-tail assumptions, a separate high-probability bound removes the constant-probability theorem’s multiplicative failure-probability limitation.Source: §2.1 · pp. 22–23

2.2 Mean spikes

  • A nonzero mean creates a rank-one spike in the second-moment matrix, so when squared mean norm is comparable to the second moment, effective rank is small and empirical stable rank is bounded by empirical noise-to-signal ratio.Under mild light-tail assumptions, empirical and population noise-to-signal ratios are close.Source: §2.2 · p. 24

2.2.1 Bounding the effective rank for random data: polynomial maps

  • For squared polynomial features of Gaussian data, stable rank admits a dimension-free bound with constant probability, while its dependence grows exponentially with polynomial degree.For quadratic feedforward networks, the resulting depth dependence is doubly exponential; the theorem has only constant-probability guarantees.Source: §2.2.1 · pp. 25–26

2.2.2 Bounding the stable rank with random weights: polynomial-growth activations

  • With random weights and fixed inputs, post-activation matrices have iid rows, so the analysis applies iid-column bounds to their transposes and expresses the population noise-to-signal ratio through one-dimensional Gaussian moments.The supplied passage introduces this approach but does not include the resulting random-weights bounds.Source: §2.2.2 · p. 26
  • For q-homogeneous activations, sσ(a,b) and κσ,q(a,b) are scale-invariant, and the low-stable-rank result has a simpler specialization when the activation mean is nonzero.The corollary assumes q ≥ 1 and σ(0) = 0.Source: §2.2.2 · pp. 27, 29
  • Gaussian weights yield low-stable-rank post-activation matrices for polynomial-growth activations under the theorem’s stated conditions.The bound is expressed in terms of sσ(a,b) and holds with high probability.Source: §2.2.2 · p. 27

2.3 Token indicator matrices and Gaussian embeddings

  • Under the Gaussian embedding model, the embedding matrix’s stable rank is controlled by pmax up to a constant factor, so it is O(1) when pmax is bounded below by a numerical constant.The stated bound st(X) ≤ 1/pmax is generally loose for the embedding matrix.Source: §2.3 · pp. 30–32
  • The token-indicator matrix has stable rank exactly 1/pmax; in the NanoGPT experiment it stayed near 26 throughout training, corresponding to pmax near 1/26.The passage relates this observed frequency to a Zipf-type distribution over a vocabulary of 50,257 tokens, with exponent about 0.86.Source: §2.3 · p. 30

3 Propagation of the stable rank

This section develops Gaussian-initialization bounds for propagating stable-rank control through transformer building blocks. It establishes operation-specific conditions and supporting properties of residual updates and gated layers.

3 Propagation of the stable rank

  • Propagation bounds for transformer building blocks relate output stable rank to the stable ranks of the matrices entering each operation.The section covers pointwise nonlinearities, residual connections, gating, and RMSNorm, then packages atomic-operation bounds for attention and MLP sublayers.Source: §3 · p. 32

3.1 Atomic operations

  • Atomic-operation results provide rules for tracking how low-stable-rank structure propagates through transformer components.The section presents these results first for individual operations, then applies propagation rules to transformer data matrices; RMSNorm iteration additionally requires column-norm regularity.Source: §3.1 · p. 32, §3 · p. 32

3.1.1 Linear layer

  • At Gaussian initialization, a Gaussian linear layer yields a high-probability stable-rank bound for its output in terms of the input stable rank.Theorem 3.1 assumes iid Gaussian weights and gives the bound with probability at least 1 − 4 exp(−cε²k).Source: §3.1.1 · p. 32

3.1.2 RMSNorm layer

  • RMSNorm has a stable-rank propagation bound when every input column norm lies between positive lower and finite upper bounds.The theorem’s column-norm regularity condition supports applying the bound iteratively.Source: §3.1.2 · p. 33

3.1.3 Pointwise nonlinear layer

  • For Gaussian weights, pointwise activations satisfying |σ(t)| ≤ |t| and a nonzero first Hermite coefficient admit a high-probability stable-rank propagation bound.Theorem 3.4 states the activation conditions and provides the bound with probability at least 1 − 8 exp(−cε²p²k).Source: §3.1.3 · p. 35

3.1.4 Residual connections

  • A Gaussian residual skip connection also satisfies a high-probability stable-rank propagation bound for fixed input and residual matrices.Theorem 3.5 gives probability at least 1 − 10 exp(−cε²d).Source: §3.1.4 · p. 37
  • At Gaussian initialization, residual updates preserve column-norm regularity by concentrating each updated column’s squared norm, scaled by 1/d, around 1.This regularity is needed for subsequent stable-rank bounds, including the RMSNorm bound.Source: §3.1.4 · p. 38

3.1.5 Gated nonlinear layer

  • The gated-layer theorem bounds how the sum’s trace and operator norm compare with the largest corresponding individual summand under stated conditions.The trace bound assumes bounded diagonal entries of H and a nonzero first Hermite coefficient; the operator-norm bound has conditions on the columns of Z.Source: §3.1.5 · pp. 39–40

A lower bound on ETi via the mean outer product. Let v „ Np0, 1

This section develops layerwise bounds and constructions for controlling activation stable-rank propagation through transformer components. It covers gated activations, attention and MLP sublayers, residual connections, and mixture-of-experts routing.

A lower bound on ETi via the mean outer product. Let v „ Np0, 1

  • The mean outer product gives K = Cov(a) + μμ^T ⪰ μμ^T, yielding a lower bound on E[T] for simplex weights.The argument identifies T as having the same distribution as (w^T a)^2 and uses the mean-vector lower bound.Source: A lower bound on ETi via the mean… · pp. 41–42
  • A concentration argument combines the lower bound on E[T] with Bernstein’s inequality to obtain a high-probability lower bound on the sum of the T_i.The supplied passages give the concentration conclusion but do not preserve its full probability expression legibly.Source: A lower bound on ETi via the mean… · pp. 42–43
  • Theorem 3.8 establishes stable-rank propagation for a gating block induced by a nonlinear activation with nonzero mean, under proportionate-width and stated input-norm and activation assumptions.The setup uses Gaussian matrices and a 1-Lipschitz activation with σ(0)=0; the supplied passages do not expose the theorem’s full bound.Source: A lower bound on ETi via the mean… · p. 43
  • Lemma 3.9 bounds gated-activation column norms to help control their evolution through the transformer architecture.The lemma assumes normalized input-vector norm, independent Gaussian gate and input weights, and an activation with bounded positive second moment.Source: A lower bound on ETi via the mean… · p. 45
  • The supplied proof bounds S + T by (1 + ε)S, and hence by (1 + ε)^2 m^2 k, establishing the claim after adjusting constants in the last term.Source: A lower bound on ETi via the mean… · p. 46

3.2.1 Attention layer propagation

  • Under column-stochastic head mixing, the attention-layer bounds are uniform in the number of heads and do not depend on masks or positional encodings.The lemma assumes bounded input column norms and d_head ≥ Cδ^-2 log(n n_head), and provides high-probability column-norm and stable-rank bounds.Source: §3.2 · p. 47, §3.2.1 · pp. 47–48

3.2 Compound operations: attention and MLP layers

  • The propagation framework distinguishes MSI activations, which induce a rank-one second-moment contribution and uniformly low stable rank, from activations whose stable rank tracks the RMS-normalized input.The latter regime also tracks uniform column-norm bounds; gated MLPs and pointwise activations on normalized inputs are examples.Source: §3.2 · p. 47
  • Under arbitrary routing, the one-off MoE calculation shows stable rank propagation assuming only σ̂_1(1) > 0 from Theorem 3.4.The paper says it does not return to MoE layers elsewhere.Source: §3.2 · p. 47

3.2.2 Mean-spike-inducing (MSI) MLP layer propagation

  • With min(d, k) ≥ Cδ^-2 log n, the MSI MLP lemma bounds output column norms and stable ranks with high probability.For ReLU, the stated activation parameters are st_σ = 2π and p_σ = 1/2, with success probability 1 − c exp(−Ck).Source: §3.2.2 · pp. 50–51

3.2.3 MLP layer propagation without mean induced spikes

  • Gated and pointwise activations with positive first Hermite coefficient satisfy the MLP propagation condition at Gaussian initialization.The condition tracks both stable-rank propagation from the RMS-normalized input and a uniform upper bound on intermediate column norms.Source: §3.2.3 · pp. 53–54
  • Under Assumption B, an independent Gaussian output map preserves bounded column norms and controls stable ranks through the MLP residual update when d ≥ Cδ^-2 log n.The stable rank of B is controlled by that of the RMS-normalized input; the residual output is also bounded in stable rank on the same high-probability event.Source: §3.2.3 · p. 55

3.2.4 Mixture of Experts layer propagation

  • For arbitrary routing, the MoE calculation controls stable ranks for every expert and gives a stable-rank bound for the output, with routing weights entering only through the simplex bound.Source: §3.2.4 · pp. 56–58

4 MLPs with Gaussian initialization

At Gaussian initialization, the stable rank of post-activation matrices in bias-free MLPs follows different depth behavior depending on the activation. Mean-spike-inducing activations reset stable rank at each layer, while a nondegenerate first Hermite coefficient permits propagation with at most exponential inflation.

4 MLPs with Gaussian initialization

  • Mean-spike-inducing activations keep every post-activation matrix at stable rank bounded by a constant; otherwise, a uniformly nondegenerate first Hermite coefficient allows propagation through depth with at most exponential inflation.The bound in the mean-spike case depends only on the activation’s one-dimensional Gaussian moments and is independent of input stable rank.Source: §4 · p. 58
  • For homogeneous activations with nonzero Gaussian mean, the stable-rank bound is input-independent under the stated growth and moment assumptions.The theorem’s bound depends on activation-specific quantities, including its homogeneity degree and Gaussian moments.Source: §4 · p. 59
  • When the Gaussian mean is zero, stable-rank propagation instead uses first-Hermite-coefficient nondegeneracy; the stated corollary assumes a 1-Lipschitz activation with σ(0)=0.The propagation argument applies this control conditionally layer by layer.Source: §4 · pp. 59–60

5 Transformers at Gaussian initialization

At Gaussian initialization, the section derives stable-rank bounds for transformer activations by propagating bounds through attention and MLP sublayers. The bounds depend on whether the MLP induces mean spikes: MSI yields at-most-quadratic growth with depth, while nonMSI propagation yields exponential-in-depth bounds.

5 Transformers at Gaussian initialization

  • The analysis bounds stable ranks of RMS-normalized activations and MLP post-activations throughout a Gaussian-initialized transformer using layerwise propagation.The bounds assume column-stochastic attention mixing and Gaussian initialization of value and output projections; no particular query/key initialization or attention-generation method is required.Source: §5 · p. 60

5.1 MSI transformer

  • In the MSI regime, MLP post-activation stable rank is bounded by stσ, while RMS-activation stable-rank bounds grow at most quadratically with depth.The bound depends explicitly on stσ and pσ.Source: §5.1 · pp. 61–62

5.2 Transformer without mean induced spikes

  • Without mean-induced spikes, the MLP does not reset stable rank, so hidden-state propagation is multiplicative and the resulting activation bounds grow exponentially with depth.The stated regime includes gated MLPs and pointwise activations satisfying the cited condition on σ̂1(1).Source: §5.2 · p. 64

6 Nuclear Rank in random feature models

In spiked random-feature regression, gradient nuclear rank can grow from a constant at initialization to dimension-dependent values after a short training period. The section establishes this behavior for realizable and teacher–student models, including multi-spike settings and more realistic step sizes.

6 Nuclear Rank in random feature models

  • At initialization the gradient nuclear rank is O(1), but after a few steps it becomes Θ(d) in the spiked realizable and teacher–student models.The rank remains Ω(d) over a window of at least Θ(d) gradient-descent iterates.Source: §6 · p. 66
  • Figure 16 shows nuclear rank rising rapidly in the first few steps in both models, with a stronger increase at larger feature dimension k=2d.The figure plots gradient nuclear rank against gradient step for several feature dimensions.Source: Figure 16 · p. 67

6.1 Random feature model I

This section studies how the gradient’s nuclear rank evolves in random feature regression, from a bounded value at initialization to levels proportional to the feature dimension after a burn-in. It connects this growth to spiked post-activation matrices and analyzes how long the large-rank regime persists under the stated model assumptions.

6.1 Random feature model I

  • At initialization the gradient’s nuclear rank is bounded, but after one gradient step it grows with dimension, supporting an advantage for spectral descent.The section’s opening result frames these claims under mild assumptions for random feature regression.Source: §6.1 · p. 67

6.1.1 ReLU post-activation matrices follow the spiked model

  • Under row incoherence, ReLU post-activation matrices follow a spike-plus-bulk model; for Gaussian V with k proportional to d and n at least C0d, the mean norm squared scales with d while covariance eigenvalues stay bounded.The theorem states these properties hold with probability tending to one as d grows.Source: §6.1.1 · pp. 70–71

6.1.2 The nuclear rank at initialization is constant

  • At the origin, the gradient’s nuclear rank is bounded by a constant under the stated condition that the ground-truth matrix does not eliminate the spike.The passage notes this condition holds, for example, when the singular values of the ground-truth matrix are bounded above and below by constants.Source: §6.1.2 · p. 72

6.1.3 The nuclear rank after one gradient step scales with dimension.

  • After one gradient step, the gradient’s nuclear rank grows with dimension under the spiked model and the specified stepsize; typical ground-truth matrices have dimension-growing nuclear rank.The lemma establishes bounded conditioning of Tη in the proportional regime, leading to the growth result.Source: Figure 17 · p. 68, §6.1.3 · p. 72

6.1.4 Multi-step and Multi-spike

  • With multiple spikes and stepsize η = 1/(c‖B‖op), the same qualitative growth holds after a short burn-in, with large gradient nuclear rank over a long time window.This removes the single-spike and stepsize simplifications of the preceding analysis.Source: §6.1.4 · p. 73
  • For constant accuracy, this large-nuclear-rank window is shorter than gradient descent’s order-d^u log(d/ε) full run by only a log(d) factor.Source: §6.1.4 · pp. 74–75
  • With high probability, the gradient nuclear rank is at least C3d for iterations from C1d^{u−ℓ}log(d) through C2d^u under Assumption E and Gaussian ground truth.The burn-in is logarithmic in d when ℓ=u.Source: §6.1.4 · pp. 73–74
  • The bulk-versus-spike analysis extends the window of large nuclear ranks by a logarithmic factor, with the proof’s horizon set to c d^u log(d).The supplied passage says the corresponding lower bound is d^-ε but omits the quantity being bounded, so its value is not specified here.Source: §6.1.4 · p. 75

6.2 Random feature model II

This section examines the teacher–student random feature model under Gaussian initialization and spiked-model assumptions. It tracks how gradient nuclear rank grows with dimension—from O(1) initially to order d after one step and, under multi-spike assumptions, after a short burn-in—preserving the spectral-versus-Euclidean separation over a long training window.

6.2 Random feature model II

  • Under the spiked-model assumptions, the rank-one spike separates from the bulk, whose non-leading singular values remain bounded within fixed constants.The setup assumes the two matrices share the same right singular vector for their spikes.Source: §6.2 · pp. 75–76

6.2.2 The nuclear rank at initialization is constant.

  • At initialization, the gradient’s nuclear rank is bounded by a constant when the ground-truth matrix does not eliminate the spike and satisfies the stated norm conditions.The conditions include bounded operator norm and a nontrivial action on the spike vector.Source: §6.2.2 · p. 79

6.2.3 The nuclear rank after one gradient step scales with dimension.

  • After one gradient step, the gradient’s nuclear rank grows with dimension under Assumption F and the stated step size; the Gaussian case also gives this growth when n ≥ Cd.The theorem considers k proportional to d; its Gaussian specialization additionally assumes m proportional to d and k/d → γ ∈ (1,∞).Source: §6.2.3 · p. 80

6.2.3 The nuclear rank after one gradient step scales with dimension

  • In the rank-one spiked teacher–student model, the gradient nuclear rank rises from O(1) at initialization to order d after one gradient step.The section notes that a similar growth persists over multiple steps for multi-spiked Gram matrices in the realizable model.Source: §6.2.4 · p. 82

6.2.4 Multi-step and multi-spike for the teacher–student model

  • The multi-spike model assumes low-rank signal components with aligned right singular directions embedded in well-conditioned bulk matrices.The signal rank is p=o(d), while the bulk has order d singular values of constant size.Source: §6.2.4 · pp. 82–83
  • Gradient descent gives the time-t gradient as the initial gradient multiplied on the right by (I−ηB)^t, making the spectrum of this matrix family control its evolution.The recursion is ∇L(W_t)=−W̄B(I−ηB)^t.Source: §6.2.4 · p. 83
  • Under the multi-spike assumptions, Theorem 6.11 establishes a long window of large gradient nuclear rank, beginning after burn-in d^(u−ℓ) log d and extending through d^u.The section concludes that the spectral-versus-Euclidean separation persists in the teacher–student setting.Source: §6.2.4 · pp. 83, 85, 86
  • With iid Gaussian initialization and the stated multi-spike scaling assumptions, the gradient nuclear rank is at least order d throughout a window from burn-in d^(u−ℓ) log d to C3d^u, with probability tending to one.The result is conditioned on the data matrices; k/d tends to a positive constant, and u−ℓ < 1/2.Source: §6.2.4 · p. 86
  • The large-nuclear-rank window can be extended by a logarithmic factor, to a horizon c d^u log d for a sufficiently small constant c.For any ε>0, the result holds with probability tending to one as d grows.Source: §6.2.4 · p. 86

7 A General Layered Model and Spectral Descent

The section extends the Euclidean-versus-spectral descent comparison to general layered models under smoothness assumptions, deriving layerwise curvature bounds in terms of feature and parameter terms. It gives a sufficient condition favoring spectral updates on blocks whose gradients have high nuclear-to-Frobenius ratio relative to the stable rank of incoming core activations.

7 A General Layered Model and Spectral Descent

  • Under mild smoothness assumptions, the layered model preserves the spectral-versus-Euclidean descent tradeoff, reducing comparison to gradient nuclear and stable ranks and incoming-feature stable rank.The model covers linear parameter blocks with smooth activation maps, including MLPs, convolutional networks, and transformer blocks.Source: §7 · pp. 86–87

7.1 Model and smoothness assumptions

  • The model fixes input features and uses parameter-linear preactivations, C2-smooth activation maps, and a C2-smooth outer loss.Each layer’s activation depends on previous preactivations, while each preactivation is linear in that layer’s weights and incoming activations.Source: §7.1 · p. 87

7.2 Layerwise Hessian bound

  • The layerwise Hessian bound separates a feature-dependent term from a parameter-only term, with constants depending on local derivative bounds, weight norms, and depth, but not layer widths.For averaged squared loss or softmax cross-entropy, the derivative constants scale as Jf = O(1/sqrt n) and Hf = O(1/n); the bound’s constants do not grow with sample count or layer widths.Source: §7.2 · pp. 88, 90, 91

7.2.1 Derivative constants for MLPs and transformer blocks

  • For fixed smooth MLP activations, the activation-side derivative constants depend on scalar bounds for σ′ and σ″, not network width; the cited averaged losses add only 1/n factors.The squared-loss and softmax cross-entropy examples inherit their scaling from averaging over n examples.Source: §7.2.1 · p. 92
  • A decoder block is represented in the layered model through attention projections, value aggregation, residual connections, RMS normalization, and MLP maps.The construction enumerates trainable matrices for query, key, value, output, and MLP projections, with smooth maps collecting the nonlinear operations feeding each linear map.Source: §7.2.1 · pp. 92–93
  • Derivative constants for MLP and transformer blocks remain bounded as model width and sequence length vary.The bounds combine primitive estimates for RMSNorm, softmax, and the MLP nonlinearity, with constants independent of model width and sequence length.Source: §7.2.1 · p. 94

7.3 From feature metric to blockwise Lipschitz constants

  • The feature-based curvature ratio between spectral and Frobenius geometries equals the stable rank of the core activation, because surrounding linear maps cancel.For transformer blocks, the relevant core activations include RMS-normalized hidden states and MLP post-activations; curvature is governed by their stable rank up to fixed operator-norm factors.Source: §7.3 · p. 95

7.4 Descent bounds and spectral-vs.-Euclidean comparison

  • Under the quadratic model, spectral updates on selected blocks guarantee at least as much decrease as Euclidean updates when each selected block’s gradient nuclear rank is at least its incoming core activation’s stable rank.The guarantee is up to the common Taylor remainder; other blocks retain Euclidean updates.Source: §7.4 · p. 98

7.5 Extension: diagonal RMS-normalization weights

  • For diagonal RMSNorm weights, spectral descent reduces to a signSGD-type update and is favored when the gradient’s ℓ1/ℓ2 nuclear-rank analogue exceeds the incoming activation’s stable rank.This is the diagonal analogue of the matrix-block condition.Source: §7.5 · pp. 98, 100

A Sparse Regression and SwiGLU activations

With SwiGLU, hidden-layer stable ranks at initialization grow with batch size, unlike the ReLU and squared-ReLU cases. SpecGD shows an early loss advantage at the smaller batch, while its trajectory essentially matches GD at the larger batch.

A Sparse Regression and SwiGLU activations

  • SwiGLU post-activation stable ranks grow noticeably with batch size, unlike ReLU and squared-ReLU, whose mean-induced spike keeps them essentially constant.SwiGLU is approximately centered and does not induce a large mean spike in hidden-layer Gram matrices.Source: A Sparse Regression and SwiGLU… · p. 105, Figure 19 · p. 105
  • At n = 512, SpecGD shows a visible early loss advantage; at n = 8196, GD and SpecGD trajectories essentially coincide.The larger-batch hidden layers have substantially higher initialization stable ranks, reducing the condition favoring spectral updates.Source: A Sparse Regression and SwiGLU… · p. 105, Figure 20 · p. 106
  • During both runs, the first two hidden-layer stable ranks stay near initialization, while the scalar-output third layer rapidly drops to about 3.SpecGD applies spectral updates only to hidden layers, not the final vector-weight layer.Source: Figure 20 · p. 106

B NanoGPT with SwiGLU activations

The NanoGPT-scale run replaces the squared-ReLU MLP nonlinearity with SwiGLU and reports the same diagnostics as pointwise analogues of the main-text plots.

B NanoGPT with SwiGLU activations

  • The NanoGPT-scale run replaces squared-ReLU with SwiGLU and tracks the same stable-rank and gradient nuclear-rank diagnostics as the corresponding main-text plots.The figures are intended as direct pointwise analogues of the modded-NanoGPT results.Source: B NanoGPT with SwiGLU activations · p. 106

C Trace and operator-norm bounds for the sample second moment under sub-Gaussian tails

The appendix introduces sample-versus-population second-moment deviation bounds under sub-Gaussian assumptions. It also states a high-probability lower bound relating an empirical second-moment operator norm to its population counterpart.

C Trace and operator-norm bounds for the sample second moment under sub-Gaussian tails

  • The appendix states separate trace-bound and non-iid operator-norm lower-tail results for sample second moments.The supplied theorem titles identify these two bounds, while their full conclusions are not given in the passages.Source: C Trace and operator-norm bounds for… · pp. 106–107
  • The setup assumes directional sub-Gaussian control of centered observations, expressed through their ψ2 norms.The passage introduces ψ2 as the sub-Gaussian norm and ψ1 as the sub-exponential norm.Source: C Trace and operator-norm bounds for… · p. 106
  • Lemma C.3 states that the empirical second-moment operator norm is at least (1 − ε) times the population operator norm with high probability.The lemma concerns independent samples and a matrix whose rows are the sampled vectors.Source: C Trace and operator-norm bounds for… · p. 108

D Auxiliary linear algebraic results.

The appendix states auxiliary bounds on blocked matrices and nuclear rank, alongside SwiGLU results showing large gradient nuclear ranks in MLP and attention blocks.

E Shardwise SpecGD and distributed orthogonalization

The section compares global and shardwise ways to orthogonalize sharded gradients for spectral updates. Shardwise SpecGD avoids incremental optimizer communication and reduces per-device work, but its polar step is local rather than global; a blockwise majorization gives a condition for its one-step decrease to match or exceed Euclidean GD.

  • Muon applies or approximates the gradient’s polar factor, commonly computed with an inverse square root and approximated using Newton–Schulz iterations.When Q ≥ P, each iteration costs on the order of Θ(P^2Q) flops.Source: E Shardwise SpecGD and distributed… · p. 113
  • Global alternatives are all-gather followed by local orthogonalization, distributed Newton–Schulz with repeated synchronization, or resharding with two all-to-alls.Source: E Shardwise SpecGD and distributed… · p. 113
  • The factor-S per-device flop reduction comes with a structural tradeoff: shardwise orthogonalization enforces orthogonality only within each shard and does not equal the global polar factor.Source: E Shardwise SpecGD and distributed… · pp. 112, 114
  • Shardwise SpecGD requires no incremental optimizer communication and Θ(P^2Q/S^2) per-device orthogonalization flops.The comparison table contrasts this with Θ(P^2Q) for all-gather then local orthogonalization and Θ(P^2Q/S) for distributed Newton–Schulz and resharding.Source: E Shardwise SpecGD and distributed… · p. 114
  • For row or column partitions, the blockwise polar step guarantees at least as much one-step decrease as Euclidean GD when the gradient nuclear-rank proxy exceeds the incoming activation stable rank.Source: E Shardwise SpecGD and distributed… · p. 116
  • In the SwiGLU run, WO gradient nuclear rank is typically larger than WV, and both exceed the single-digit stable ranks of their corresponding RMS inputs.Source: Figure 27 · p. 117
Loading 2512.04299v2…