PaperlayerView exact revision on arXiv

Source-linked AI summary

Enhancing LLM Training via Spectral Clipping

Xiaowen Jiang, Andrei Semenov, Sebastian U. Stich

arXiv:2603.14315v2cs.LGmath.OC

TL;DR

Coordinate-wise optimizers do not account for the global spectral structure of weights and gradients, leaving update norms and sparse gradient-noise spikes insufficiently controlled. SPECTRA adds post-spectral clipping, optionally clips raw gradients, and frames post-clipping as constrained optimization with implicit regularization. Across reported pretraining experiments, it improves validation loss for base optimizers and yields smaller weight norms.

  • Problem

    Coordinate-wise methods such as AdamW do not account for the global spectral structure of weights and gradients, while updates can have excessive spectral norms and stochastic-gradient noise can contain sparse spectral spikes.

    Source: §1 · p. 1

  • Method

    SPECTRA applies post-spectral clipping to optimizer updates and optionally pre-spectral clipping to raw gradients, with post-clipping analyzed as Composite Frank-Wolfe optimization under spectral-norm constraints and implicit weight regularization.

    Source: §1 · pp. 1–2

  • Results

    SPECTRA consistently improves base-optimizer validation loss across synthetic and LLM pretraining tasks from 124M to 1.5B parameters, and coordinate-wise variants achieve the best performance among state-of-the-art methods.

    Source: §1 · p. 2

  • Takeaways & Limitations

    SPECTRA-trained models have smaller weight norms, consistent with the paper’s link between spectral clipping and weight regularization.

    Source: §1 · p. 2

    The Composite Frank-Wolfe equivalence is established for the stated setting with Q=Q2 and exact solution of each subproblem for computing V_{k+1}.

    Source: §3 · p. 4

Abstract

from arXiv · show

While spectral-based optimizers like Muon operate directly on the spectrum of updates, standard adaptive methods such as AdamW do not account for the spectral structure of weights and gradients, leaving them vulnerable to two empirical issues in large language model (LLM) training: (i) the optimizer updates can have large spectral norms, potentially destabilizing training and degrading generalization; (ii) stochastic gradient noise can exhibit sparse spectral spikes, with a few dominant singular values much larger than the rest. We propose SPECTRA, a general framework addressing these by (i) post-spectral clipping of updates to enforce spectral-norm constraints (ii) optional pre-spectral clipping of gradients to suppress spectral noise spikes. We prove that post-clipping constitutes a Composite Frank-Wolfe method with spectral-norm constraints and weight regularization. We further analyze how pre-clipping mitigates sparse spectral spikes. We propose efficient soft spectral clipping via Newton-Schulz iterations, avoiding expensive SVD. Experiments on LLM pretraining show SPECTRA uniformly improves validation loss for various optimizers, including AdamW, Signum, Mars, and AdEMAMix, with the best-performing variants achieving state-of-the-art results. Models trained with SPECTRA exhibit smaller weight norms, confirming the link between spectral clipping and regularization.

1. Introduction

Coordinate-wise optimizers remain competitive, but they do not account for global spectral structure in weights and gradients. SPECTRA adds post-spectral clipping, with optional pre-spectral clipping, to control updates and gradient noise.

Contributions

  • SPECTRA is an optimizer-agnostic wrapper that clips update matrices spectrally and can optionally clip raw gradients to filter spectral noise spikes.The framework is designed for base optimizers with decoupled weight decay, including AdamW, Signum, Mars, and AdEMAMix.Source: §1 · pp. 1–2

Implicit Regularization & Spectral Norm Constraints

  • Post-spectral clipping is equivalent to Composite Frank-Wolfe optimization with a spectral-norm constraint and implicit weight regularization.The hyperparameters control the constraint-ball radius and regularization strength.Source: §1 · p. 2

Robustness to Spectral Spikes

  • The authors report spectral noise spikes in LLM stochastic gradients and analyze theoretically how spectral clipping can mitigate their effects, potentially outperforming global norm clipping.The analysis concerns low-rank noise components with singular values much larger than the signal.Source: §1 · p. 2

Empirical SOTA

  • SPECTRA consistently improves base-optimizer validation loss across synthetic and LLM pretraining tasks from 124M to 1.5B parameters.It also enables larger learning rates; coordinate-wise methods with SPECTRA achieve the best performance among state-of-the-art methods.Source: §1 · p. 2
  • SPECTRA-trained models exhibit smaller weight norms, consistent with the paper’s theoretical link between spectral clipping and regularization.The framework’s clipping approximation uses Newton-Schulz iterations and fast square matrix–matrix multiplications.Source: §1 · p. 2

2. Preliminaries and Setup

The setup focuses on transformer weight matrices and standard decoupled-weight-decay optimizer updates. It motivates spectral clipping by describing excessive update norms and sparse spectral spikes in stochastic gradients, then introduces clipping as a remedy.

Preliminaries and Setup

  • The optimizer update rule applies decoupled weight decay to the current matrix and subtracts a learning-rate-scaled optimizer update.The optimizer-specific computation of the update matrix distinguishes methods including AdamW, AdEMAMix, Signum, Mars, and Muon.Source: §2 · p. 2

Phenomenon I: Excessive Spectral Norm of Update Matrices

  • Uncontrolled update spectral norms can let weight matrices grow rapidly, while small learning rates or extended warm-up may be needed to reduce instability risk.The setup notes that AdamW may produce excessive update magnitudes, particularly early in training or before loss spikes.Source: §2 · p. 2

Phenomenon II: Sparse Spectral Spikes in Stochastic Gradients

  • Stochastic-gradient noise can contain low-rank spectral spikes whose singular values greatly exceed the signal and whose subspace is typically nearly orthogonal to principal signal directions.The signal is approximated using large-batch gradients, with the residual treated as noise.Source: §2 · p. 3

Remedy: Spectral Clipping Operator

  • Spectral clipping applies a threshold to singular values of update matrices, strictly controlling their spectral norm.The operator can also be applied to raw stochastic gradients; in the presence of sparse spikes, it can remove large noise components while preserving the remaining signal’s scale and structure.Source: §2 · p. 3
  • Coordinate-wise clipping ignores entry correlations, while global clipping suppresses useful signal along with noise; spectral clipping targets the spectral structure instead.The figure caption compares final validation loss for large Llama-style models trained to the Chinchilla-optimal horizon, including thin/deep and wide/shallow configurations.Source: §2 · p. 3, Figure 1 · p. 3

3. Post-Spectral Clipping: Spectral Norm Constrained Optimization with Implicit Weight Regularization

Post-spectral clipping can be reformulated as composite Frank-Wolfe optimization over a spectral-norm ball with weight regularization. Its parameters control the constraint radius and regularization strength, and the framework also describes other spectral optimizers.

3. Post-Spectral Clipping: Spectral Norm Constrained Optimization with Implicit Weight Regularization

  • Post-spectral clipping is equivalent to a stochastic composite Frank-Wolfe method with a spectral-norm constraint and Frobenius-norm regularization.The equivalence holds when the Frank-Wolfe subproblems are solved exactly.Source: §3 · pp. 3–4
  • The analysis provides standard convergence guarantees for the clipped update under the stated convexity and stochastic-gradient assumptions.The theorem also assumes a smooth objective for its additional guarantee.Source: §3 · p. 4
  • The clipping and decay parameters determine the spectral-ball radius and regularization level, respectively.With α fixed, c/λ controls the constraint radius, while λ determines regularization strength.Source: §3 · p. 4
  • As α increases and c is set to 1/α, post-spectral clipping recovers Muon as a special case without regularization.Source: §3 · p. 4

3.1. Extensions: A Family of Spectral Optimizers

  • Choosing different regularizers yields a family of spectral optimizers, including nuclear-norm and matrix-entropy variants, while clipped Signum approximates a regularized Frank-Wolfe step.Source: §3.1 · pp. 4–5

4. Robustness to Sparse Spectral Spikes

The section examines sparse spectral spikes in stochastic gradients and how spectral clipping can mitigate their effects. Its empirical observations and theoretical analysis contrast spectral clipping with global clipping, especially in preserving signal under low-rank spike noise.

4.1. Empirical Observations

  • In LLM training, stochastic-gradient noise has dominant singular-value spikes that can exceed the signal norm and are nearly orthogonal to signal directions.The observations include AdamW/Signum training and report heavy-tailed raw-gradient spectra across layers.Source: §4.1 · p. 5
  • Because high-magnitude noise components are geometrically separated from the signal, spectral clipping can truncate spikes without significantly distorting the signal direction.Source: §4.1 · p. 5

4.2. Theoretical Insights

  • The theoretical model treats the stochastic gradient as a true signal plus low-rank, zero-mean noise with spike magnitude ℓ.The analyzed noise directions are assumed to satisfy bounded-anisotropy conditions.Source: §4.2 · pp. 5–6
  • Global clipping either weakens signal preservation as spike magnitude grows or retains variance scaling as ℓ², because it scales the entire matrix, including signal.Source: §4.2 · p. 5
  • Spectral clipping improves the noise dependence from rℓ² to r min(√rM, ℓ)² and is robust to any spike level ℓ with c ≃√rM.The rate comparison is stated for the theorem's low-rank spike-noise setting and its assumptions.Source: §4.2 · p. 6

5. SPECTRA: SPEctral Clipping for LLM TRaining Acceleration

SPECTRA applies spectral clipping within a broad class of optimizers. Soft Spectral Clipping approximates exact clipping using square matrix multiplications, while the algorithm optionally clips raw gradients before computing optimizer updates.

5.1. Soft Spectral Clipping

  • Soft Spectral Clipping approximates exact spectral clipping, which otherwise requires a costly full SVD.The approximation applies a smooth scalar clipping function to singular values and requires square matrix–matrix multiplications.Source: §5.1 · p. 6
  • Newton-Schulz iterations implement the approximation using an inverse-square-root computation, with reported overhead becoming a smaller fraction of step time at larger scale.The method uses the well-conditioned matrix I + XXᵀ/c².Source: §5.1 · p. 6

5.2. Algorithm Description

  • SPECTRA accepts a base optimizer, optionally spectrally clips raw gradients before its update computation, then spectrally clips the resulting update.The optional pre-clipping step is described as useful for removing sparse-spike noise before it affects momentum buffers, typically with small batches.Source: §5.2 · pp. 6–7
  • The recommended configuration uses λ ≈0.1, c ≈10, and Ns = 10 across the reported experiments.The c/λ ratio controls the spectral-constraint radius, while λ sets regularization strength.Source: §5.2 · p. 7

3. Dynamic Clipping Schedules.

During warm-up, the clipping threshold is adjusted inversely with the learning rate to maintain effective update capacity; afterward, the default schedule keeps it constant. In the small-batch comparison, Spectra-Signum remains stable at a large learning rate, unlike Signum.

3. Dynamic Clipping Schedules.

  • During warm-up, the clipping threshold is scaled inversely with the learning rate to maintain effective update capacity.The threshold is held at c after warm-up by default.Source: §3 · p. 7
  • Under small-batch training, Spectra-Signum remains stable at a large learning rate, whereas Signum exhibits training instability.Source: Figure 2 · p. 7
  • With post-clipping already applied, adding pre-clipping produces a validation-loss curve that almost overlaps the post-only variant.Source: Figure 2 · p. 7

6. Experiments

Experiments test SPECTRA on synthetic optimization and LLM pretraining, where it improves optimizer performance and supports stable training under challenging settings. Across model sizes and training durations, SPECTRA variants achieve strong validation-loss and downstream results, with smaller weight norms also observed.

6.1. Synthetic Experiments.

  • On synthetic matrix logistic regression, the update rule minimizes the composite objective within the spectral-norm ball, and appropriate regularization improves generalization.The experiments use hyperparameters from Theorem 3.4.Source: §6.1 · p. 7
  • Under rank-r spectral-spike noise, spectral clipping maintains fast convergence across noise scales, while vanilla and globally clipped SGD require lower learning rates to avoid divergence.The reported benefit also extends to momentum-based updates.Source: §6.1 · p. 7

6.2. LLM Pretraining Tasks.

  • At a 124M scale with small batches, Spectra-Signum remains stable at a high learning rate and improves validation perplexity by 2.5 over the best Signum baseline.Standard Signum becomes unstable and has poor validation loss as the learning rate increases.Source: §6.2 · p. 8
  • Spectra-AdEMAMix consistently outperforms other methods by a large margin across the evaluated schedules; the clipping-threshold ablation shows a wide plateau, indicating practical insensitivity to c.Ns values of 5, 10, and 15 yield almost no difference in final loss.Source: §6.2 · p. 8
  • At 820M scale, SPECTRA improves on base optimizers, and all three variants exhibit smaller aggregated ℓ∞ weight norms during training.The experiments use 1K-token sequences, batch size 992, and 16K iterations.Source: §6.2 · p. 8, Figure 4 · p. 8
  • On 10 in-context learning tasks, all three SPECTRA variants improve average accuracy over the baselines, with Spectra-Signum achieving the best result.Table 1 reports 5-shot accuracy for the 820M model pretrained on FineWeb-Edu with 16.8B tokens.Source: Table 1 · p. 9, §6.2 · p. 9
  • At 780M scale, SPECTRA variants rank among the top four, while baselines given up to 20% more steps remain on par with or worse than original-step SPECTRA runs.SPECTRA's per-step overhead is around 2% at this scale.Source: Table 2 · p. 9, §6.2 · p. 9
  • SPECTRA-AdEMAMix achieves the best validation losses across varying training durations and model sizes, and leads all considered methods at 1.5B parameters.The paper recommends Spectra-Signum for standard pretraining when batch size is sufficiently large, citing its memory efficiency.Source: §6.2 · p. 9

7. Conclusion.

The conclusion presents SPECTRA as an efficient LLM-training framework that constrains update spectral norms and improves stability and validation performance relative to standard baselines. It identifies larger-scale and longer-duration experiments, broader downstream evaluation, and further theory and implementation work as future directions.

7. Conclusion.

  • SPECTRA explicitly constrains update-matrix spectral norms, improving convergence stability and final validation performance compared with standard baselines.The conclusion says this is achieved without additional memory overhead.Source: §7 · p. 9
  • The conclusion identifies scaling to larger models, extending training, and evaluating broader downstream tasks as future experimental directions.Other directions include broader noise assumptions, alternative clipping implementations, distributed implementations, explicit weight regularization, and reinforcement learning.Source: §7 · p. 9
  • The paper also calls for further work on distributed spectral clipping, explicit weight regularization, and applications beyond LLM optimization.The listed application includes reinforcement learning.Source: §7 · p. 9

Impact Statement

The section discusses potential environmental and hardware-access benefits of SPECTRA’s efficiency contribution and notes general risks from advancing LLM capabilities. The supplied algebraic-derivation fragments add no substantive impact claim.

Impact Statement

  • SPECTRA may reduce the environmental footprint of large-scale model training and lower hardware barriers to research; the authors foresee no specific negative societal consequences beyond general LLM-capability risks.The impact statement frames these as potential positive impacts.Source: Impact Statement · p. 10

A. More Related Work

D. Analysis of SGD under Sparse Noise with Spectral Spikes

  • Lemma D.1 establishes a rotation-alignment result for a matrix with rank-r spectral-spike noise when the spike magnitude ℓ exceeds the signal spectral norm ∥G∥2.The noise has form N = ℓ U_N V_N^T, with U_N and V_N having orthonormal columns; the supplied passage does not show the resulting bound.Source: D. Analysis of SGD under Sparse Noise… · p. 20
  • The analysis says subtracting two equations yields an expression, but the supplied formula is incomplete.A following fragment gives a norm inequality, also incomplete in the supplied text.Source: D. Analysis of SGD under Sparse Noise… · p. 20

D.1. Proof of Lemma 4.2

The proof bounds the residual left by spectral clipping under low-rank noise by controlling both singular-value deviations and changes in the leading singular subspaces.

D.1. Proof of Lemma 4.2

  • The clipping residual is supported on the top-r singular directions, with diagonal entries max(σi(g), c) − c.It vanishes when c ≥ ℓ + ∥G∥2.Source: §D.1 · p. 20
  • The deviations max(σi(g), c) − ℓ are bounded by ∥G∥2, yielding ∥D̃∥F ≤ √r∥G∥2.This holds for each i ∈ [r] across the cases considered in the proof.Source: §D.1 · p. 20
  • The proof also controls the leading left and right singular-subspace distances using the signal norm and the noise spectral gap.It relates these distances to ∥G∥2 and σ1(N) − σr+1(N) − ∥G∥2.Source: §D.1 · p. 21

D.2. Proof of Theorem 4.6

The proof treats separate regimes for spike magnitude and clipping threshold, using bounds on the clipped stochastic gradient to establish the convergence claims.

D.2. Proof of Theorem 4.6

  • When ℓ ≥ 9√rM, the proof applies Lemma 4.2 with G = ∇f(Xk); for smaller ℓ, it bounds σ1(g(Xk)) and obtains clipsp_c(g(Xk)) = g(Xk).The latter case uses c ≥ M and the stated bound on the gradient spectral norm.Source: §D.2 · pp. 22–23

D.3. Proof of Lemma 4.3

  • Lemma 4.3 separates active-clipping and inactive-clipping regimes, with clipping never active when c ≥ ℓ + ∥G∥2.For c ≤ √r(ℓ − ∥G∥2), the proof states that global clipping is always active.Source: §D.3 · pp. 23–24

D.4. Proof of Theorem 4.7

  • For the Theorem 4.7 proof, one considered regime has bounds no better than vanilla SGD without clipping.The proof then proceeds to a case with c ≤ 2/3 √rℓ.Source: §D.4 · p. 24

E.1. Proof of Proposition 3.1

The appendices establish convergence tools for composite Frank-Wolfe, analyze clipping under sparse spectral noise, and report synthetic and transformer-spectrum findings.

E.1. Proof of Proposition 3.1

  • The convergence analysis uses a recurrence that accounts for curvature, gradient-estimation error, and the Frobenius diameter of the feasible set.The appendix also states stochastic CFW bounds under specified batch-size and parameter choices.Source: §E.1 · pp. 25–26

E.2. Discussion of CFW under sparse noise with spikes

  • For sparse spikes, pre-spectral clipping controls the clipped-gradient noise level by c rather than ℓ, but introduces estimator bias.The authors note that bias can be negligible in some cases, such as when the true gradient has rank much smaller than q; rigorous bias analysis is left for future work.Source: §E.2 · p. 27

E.3. Solution of the Subproblem with Other Regularizers

  • For separable regularizers acting on singular values, the subproblem reduces to q independent one-dimensional minimizations.The authors leave non-separable cases such as a squared sum of singular values for future exploration.Source: §E.3 · p. 28

F.1.1. VALIDATION OF CONSTRAINED OPTIMIZATION WITH WEIGHT REGULARIZATION

  • The synthetic logistic-regression experiment reports that update rule (4) minimizes the composite objective within the spectral ball and that b > 0 improves test loss.Larger b also yields smaller Frobenius norms; the experiment varies b and the spectral-ball radius D2.Source: §F.1.1 · pp. 29–30

F.1.2. VALIDATION OF ROBUSTNESS TO SPARSE SPECTRAL SPIKES

  • The synthetic noise experiment compares vanilla SGD, global clipping, and spectral clipping with low-rank noise spikes of magnitudes 10, 100, and 1000 against M ≈ 15.It fixes the spectral-clipping threshold at c = 15.Source: §F.1.2 · p. 31, Figure F.5 · p. 31
  • As noise spikes grow, spectral clipping maintains fast convergence, while vanilla SGD and global clipping need smaller learning rates to avoid divergence and converge more slowly.In the momentum comparison, clipping the raw stochastic gradient filters sparse spikes and lets the buffer accumulate valid signal directions.Source: §F.1.2 · p. 31, Figure F.5 · p. 31, Figure F.6 · p. 32

F.2.1. SINGULAR VALUE DISTRIBUTION OF STOCHASTIC GRADIENTS AND UPDATES. / F.2.2. NOISE STRUCTURE ANALYSIS.

  • The stochastic-gradient spectrum has a few singular values orders of magnitude larger than the rest across recorded layers and throughout training.Noise spikes frequently exceed the true-signal norm, and their principal directions are nearly orthogonal to those of the signal.Source: §F.2.1 · p. 32, §F.2.2 · p. 33

F.2.1. SINGULAR VALUE DISTRIBUTION OF STOCHASTIC GRADIENTS AND UPDATES.

  • SPECTRA-AdamW imposes a controlled update spectral-norm bound, whereas AdamW’s coordinate-wise normalization only partially mitigates spectral imbalance.The AdamW update lacks a hard ceiling; the corresponding Signum comparisons are referenced separately.Source: §F.2.1 · p. 32

F.3. LLM Pretraining Tasks.

This section reports LLM pretraining configurations and examines post-orthogonalization and pre-spectral clipping alternatives. It also compares their validation-loss outcomes and computational overhead.

F.3. LLM Pretraining Tasks.

  • Experiments use FineWeb-Edu, decoder-only transformer architectures, and a four-H100-GPU node.The section reports full validation loss and perplexity after training.Source: §F.3 · p. 33
  • At 780M and 820M scales, SPECTRA adds about 2% per-step overhead, down from about 17–19% at 160M.The forward/backward pass dominates at larger scale.Source: Table F.2 · p. 34
  • Spectra-AdamW outperforms the best-performing AdamW-Post-Orth baseline in both schedule configurations on the 160M model trained for 16.8B tokens.Post-orthogonalization normalizes all update singular values to 1, using Newton-Schulz iterations.Source: §F.3 · p. 34
  • Under noisy SlimPajama training, pre-SSC improves post-SSC-only SPECTRA-Signum by 0.014, versus 0.005 on cleaner FineWeb-Edu training.The comparison uses batch sizes 8 and 32, respectively.Source: §F.3 · p. 35
  • A cosine-decaying threshold from c = 10 to c = 0.01 outperforms a constant c = 0.1 threshold.The passage attributes the fixed threshold's weakness to over-clipping later, as gradient norms decrease and vary across layers.Source: §F.3 · p. 36

F.4. Evaluation of downstream Tasks.

This section describes five-shot downstream evaluation and reports spectral analyses of gradients and optimizer updates. It also compares optimizer resource costs and characterizes spectral noise across training stages.

F.4. Evaluation of downstream Tasks.

  • Downstream tasks use five-shot evaluation, with length-normalized accuracy for HellaSwag, ARC, PIQA, OpenBookQA, and SciQ.The remaining listed tasks use standard accuracy.Source: §F.4 · p. 36
  • SOAP is comparable to SPECTRA-Signum at 780M but worse at 820M, despite requiring more time per step and nearly twice the memory.SOAP is tuned over three learning rates; its time overhead is about 19% at 780M and 8% at 820M.Source: Table F.7 · p. 37
  • Raw AdamW gradients show a few singular values orders of magnitude larger than the rest, producing high condition numbers.The spectra are shown at four training stages across representative network layers.Source: Figure F.9 · p. 40
  • AdamW updates have a more uniform spectrum than raw gradients, but lack a hard ceiling; Signum update spectral norms range from 70 to 1300.The distributions are compared across four training stages for representative layers.Source: Figure F.10 · p. 41, Figure F.11 · p. 42
  • Spectra-AdamW and Spectra-Signum update matrices exhibit hard spectral-norm ceilings and small condition numbers.The reported spectra cover representative layers at four training stages.Source: §F.4 · pp. 43–44
  • In the 124M AdamW model, sparse noise spikes are nearly orthogonal to the large-batch signal's principal directions.The figure compares gradient singular values and spectral and chordal subspace distances across training stages.Source: Figure F.14 · p. 45, Figure F.15 · p. 46
  • With muP, spectral outliers are less pronounced in the 124M GPT than in the 124M LLaMA, and noise spikes appear mostly late in training.The muP noise analysis covers representative layers at four training stages.Source: Figure F.16 · p. 47, Figure F.17 · p. 48
Loading 2603.14315v2…