Enhancing LLM Training via Spectral Clipping

Abstract

While spectral-based optimizers like Muon operate directly on the spectrum of updates, standard adaptive methods such as AdamW do not account for the global 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, recovering Frobenius and $\ell_{\infty}$-norm regularization with SGD-based and sign-based methods. 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, 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.

Figures (24)

• Figure 1: Final validation loss comparison for large Llama-style models trained for the Chinchilla optimal horizon. 'Thin and deep' models have smaller embedding dimensions but have more layers than 'wide and shallow' ones. The corresponding validation perplexities and running time comparisons w/wo using SPECTRA are provided in Figure \ref{['fig:780M-16k-perplexity']} and \ref{['fig:run_time_iter_820M']}. We use a total batchsize of 1012 and 992 for 780M and 820M model respectively with 1024 sequence length. The hyperparameters such as learning rate used for each method are reported in the tables in Section \ref{['sec:lLM-pretrain-appendix']}.
• Figure 2: Validation loss comparison under small batch size training. Spectra-Signum maintains stability at the large learning rate, whereas Signum exhibits training instability (Signum typically prefers smaller LR than Adam-ish methods semenov2025benchmarking). Spectra-Signum with both pre and post clipping almost overlap with Spectra-Signum with post clipping.
• Figure 3: Final validation loss comparison for a small size Llama model trained with both cos and wsd learning rate schedule. The run time comparisons w/wo using SPECTRA are provided in Figure \ref{['fig:run_time_iter_160M_wsd']}. We use a total batchsize of 256 with 512 sequence length. The hyperparameters such as learning rate used for each method are reported in the tables in Section \ref{['sec:lLM-pretrain-appendix']}.
• Figure 4: Evolution of the aggregated $\ell_{\infty}$ and Frobenius norms of all weight matrices during training. SPECTRA consistently maintains a lower $\ell_{\infty}$ norm compared to the base optimizer.
• Figure 5: Comparison of different methods with or without using SPECTRA during the warm-up phase. SPECTRA improves convergence across three methods. (AdamW exhibits training instability during warm-up phase due to the use of a large step size required to achieve a low final validation loss.)

Theorems & Definitions (46)

• Proposition 1
• proof
• Theorem 1
• Remark 2
• Definition 1
• Lemma 3
• Lemma 4
• Theorem 5
• Theorem 6
• Definition 2