PaperlayerView exact revision on arXiv

Source-linked AI summary

CAST: History-Adaptive Spectral Clipping for LLM Pretraining

Jiacheng Li, Jianchao Tan, Kefeng Zhang, Yuchen Xie, Xunliang Cai

arXiv:2610.19761v1cs.LG

TL;DR

Gradient clipping changes both current updates and AdamW’s later moment estimates, while fixed thresholds can apply different clipping strengths as gradient scales change. CAST uses AdamW’s previous second moment for matrix scales and clipped-output history for thresholds. In 5,000-step MoE pretraining, it achieved lower losses than nine baselines and reached AdaGC’s final smoothed loss in 7.2% fewer steps.

  • Problem

    Clipping affects both the current gradient update and the moment estimates used by later AdamW updates, while changing gradient scales can make fixed thresholds apply different clipping strengths.

    Source: §1 · p. 1

  • Method

    CAST standardizes gradient matrices with row and column summaries of AdamW’s previous second moment, applies soft spectral clipping with thresholds from an EMA of clipped-output spectral norms, and restores the original coordinates before AdamW updates.

    Source: §1 · p. 2

  • Results

    In 5,000-step MoE pretraining, CAST had lower training and mean validation losses than nine baselines and three component variants, the fewest relative gradient spikes, and reached AdaGC’s final smoothed loss in 7.2% fewer steps.

    Source: §6 · p. 8

  • Takeaways & Limitations

    These results support using both histories to improve pretraining quality and efficiency in the reported MoE study.

    Source: §6 · p. 8

    The exact timing and speedup can vary with tensor shape, hardware, and training settings, although the reported comparison uses a shared setup.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Soft spectral clipping (SSC) reduces gradient singular values, but its effect depends on the gradient scale and clipping threshold. We introduce CAST (Clipping with Adaptive Scales and Thresholds), a history-adaptive extension of SSC. AdamW's previous second moment sets row and column scales. An exponential moving average of the clipped output's estimated spectral norm sets each matrix's threshold. CAST restores the original coordinates before the AdamW update. We train a mixture-of-experts (MoE) model with 19.08B total parameters for 83.89B tokens. CAST achieves lower training and mean validation losses than nine gradient-processing baselines, with the fewest relative gradient spikes. It reaches the best baseline's final smoothed training loss in 7.2\% fewer steps.

1 Introduction

Gradient clipping affects both the current update and AdamW’s later moment estimates, while fixed thresholds may apply different clipping strengths as matrix scales change. CAST uses AdamW’s moment history to set matrix scales and clipped-output history to adapt thresholds; in the MoE study, it achieved the lowest training and mean validation losses among nine baselines.

Motivation

  • A fixed threshold may impose different clipping strengths as gradient scales change, motivating matrix-specific historical scales and thresholds.Persistent row and column magnitude differences also affect the spectrum on which clipping acts.Source: §1 · p. 1

CAST

  • CAST adds historical scaling and output feedback around soft spectral clipping.The two histories control the standardized input coordinates and the clipping threshold.Source: Figure 2 · p. 2, §1 · p. 2

Output feedback

  • An exponential moving average of each matrix’s clipped-output spectral norm sets its current clipping threshold from the previous history.The threshold follows changes in the clipped spectrum.Source: §1 · p. 2

Historical scaling

  • AdamW’s previous second moment supplies row and column summaries for standardizing each gradient matrix, which CAST restores before the AdamW update.AdamW updates both moments from the restored gradient; its updated second moment supplies the scale at the next step.Source: §1 · p. 2

Pretraining results

  • CAST achieved the lowest training and mean validation losses among nine baselines in the MoE study.Figure 1 shows validation loss after 5,000 steps as an unweighted mean over 20 subsets, with one run per method and seed 48.Source: Figure 1 · p. 1, §1 · p. 2

2 Related Work

The related work situates CAST within smooth spectral clipping and other gradient transformations. CAST combines historical coordinate scaling with output-based threshold adaptation for spectral clipping.

2.1 Spectral Clipping

  • SPECTRA smoothly reduces singular values through a matrix inverse square root computed with Newton–Schulz iterations, and studies changing the pre-clipping threshold during training (Jiang et al., 2026).Its decaying threshold schedule performed better than a fixed threshold in the cited comparison; CAST instead uses the previous second moment for scale and clipped output for threshold updates.Source: §2.1 · p. 2

2.2 Gradient Transformations

  • Gradient transformations differ in what they constrain: AdaGC adapts per-tensor thresholds from previous clipped norms, while ZClip and GradientStabilizer use global norm statistics (Wang et al., 2026a; Kumar et al., 2025; Huang et al., 2026).CAST’s spectral history uses an estimated spectral norm in standardized coordinates, alongside AdamW’s coordinatewise second moment.Source: §2.2 · p. 2
  • Other compared approaches bound unitwise or coordinate magnitudes, or apply a signed elementwise power before the optimizer (Brock et al., 2021; Wang et al., 2026b).Together, the methods cover global, tensor, coordinate, and spectral gradient processing.Source: §2.2 · pp. 2–3

3 Soft Spectral Clipping

Soft spectral clipping applies a smooth transformation to singular values, reducing small ones slightly and larger ones more strongly. Its implementation uses a full Gram matrix and a bypass gate before approximating the transformation.

Soft spectral clipping

  • The SPECTRA soft map slightly reduces small singular values and reduces larger singular values more strongly.For a matrix, its inverse square root acts on the full q × q Gram matrix and is approximated with Newton–Schulz iterations rather than an explicit singular value decomposition.Source: §3 · p. 3

SSC routine

  • The SSC routine returns the input unchanged when the estimated largest squared singular value is at or below the squared threshold.Otherwise it approximates the soft map with K Newton–Schulz iterations and returns a flag indicating whether clipping was applied.Source: §3 · p. 3

Implementation

  • CAST and both SSC baselines share the bypass gate and Newton–Schulz approximation, while the soft map describes only the active clipping branch.The paper separately rescales inputs to evaluate the routine safely in FP32.Source: §3 · p. 3

4 Method: CAST

CAST standardizes selected gradient matrices using AdamW’s previous second moment, applies soft spectral clipping with a threshold from output history, and restores the original coordinates before AdamW updates. Its implementation also initializes histories, handles other tensors with an output-norm controller, and distributes the full-matrix transform.

4 Method: CAST

  • CAST transforms accumulated, reduced, and unscaled gradients before either AdamW moment is updated, then AdamW applies its usual bias correction and decoupled weight decay.The transformation is applied to selected two-dimensional tensors in the dense optimizer group; embeddings and output matrices are excluded by name, while other tensors use an output-norm controller.Source: §4 · p. 3

4.1 Historical Scaling

  • Historical scaling uses the bias-corrected previous second moment to set row and column coordinate scales, then restores the clipped gradient with Gout = S ⊙ Zout.Positive floors provide a reference for sparsely active coordinates, and the row and column factors account for persistent magnitude differences.Source: §4.1 · pp. 3–4

4.2 Adaptive Thresholds

  • Each selected matrix’s previous clipped-output spectral-norm history sets its current threshold, and the complete standardized output supplies the next observation.The method uses three power iterations for the estimate; positive observations update the history, while zero outputs leave it unchanged.Source: §4.2 · p. 4

4.3 Implementation

  • The first 40 successful optimizer updates use global clipping at norm 1 to initialize moments and histories; other tensors, including MoE experts and embeddings, use an output-norm threshold controller.The distributed implementation all-gathers gradient fragments into complete logical matrices, transforms them, then writes each rank’s local optimizer slice; it uses FP32 within BF16 training.Source: §4.3 · p. 5

5 Experiments

The experiments evaluate CAST in 5,000-step MoE pretraining, examining pretraining quality, component effects, gradient stability, and training cost. CAST achieves the lowest training and mean validation losses among nine baselines, with fewest relative spikes and fewer steps to reach the best baseline’s final smoothed loss.

5.1 Experimental Setup

  • The evaluation used 5,000 steps (83.89B nominal tokens) on a 19.08B-parameter MoE; each condition had one seed-48 run and validation averaged 20 subsets.The subsets included domain corpora, question–answer text, and code-formatted text.Source: §5 · p. 5, §5.1 · p. 5

5.3 Ablation Study

  • Removing HS raised training and mean validation losses by 2.42×10−3 and 8.96×10−4, while cosine rather than fixed thresholds lowered them by 1.70×10−2 and 1.82×10−2.CAST’s mean validation loss was 2.96×10−2 lower than the cosine-threshold variant’s; CAST also had lower loss on 18 subsets than the variant without HS.Source: §5.3 · pp. 5–6, Table 2 · p. 6

5.2 Pretraining Results

  • At step 5,000, CAST had the lowest training and mean validation losses among nine baselines; mean validation loss was 1.26×10−2 below SSC-constant and 2.10×10−2 below SSC-cosine.CAST also had lower loss on 20 and 19 validation subsets than SSC-constant and SSC-cosine, respectively.Source: §5.2 · p. 5, Table 1 · p. 6

5.4 Training Stability

  • CAST recorded one relative spike and no input norm above 1; SSC-constant and SSC-cosine recorded 22 and 21 relative spikes, respectively.CAST and SSC-constant had similar median norms, while CAST had a smaller upper tail; AdaGC had six relative spikes and a lower 99th-percentile norm than CAST.Source: §5.4 · p. 7, Table 3 · p. 8

5.5 Training Efficiency

  • CAST reached AdaGC’s final smoothed loss at step 4,639, using 7.2% fewer steps; estimated time was 5.5% less than 5,000 ClipByValue steps.CAST took 2.0% longer per step than SSC-constant and 1.9% longer than ClipByValue; timings varied with cluster load.Source: Figure 7 · p. 8, Table 4 · p. 8, §5.5 · p. 8

6 Conclusion

CAST uses AdamW’s previous second moment to set gradient scales and an EMA of the clipped output’s estimated spectral norm to set thresholds. In 5,000-step MoE pretraining, it achieved lower losses than nine baselines, the fewest relative gradient spikes, and AdaGC’s final smoothed loss in 7.2% fewer steps.

6 Conclusion

  • CAST uses AdamW’s previous second moment for gradient scale and an EMA of the clipped output’s estimated spectral norm for each matrix’s threshold.CAST restores the original coordinates before the AdamW update.Source: §6 · p. 8
  • In 5,000-step MoE pretraining, CAST achieved lower training and mean validation losses than nine baselines and three component variants, with the fewest relative gradient spikes.It reached AdaGC’s final smoothed loss in 7.2% fewer steps; estimated time saving was 5.5% relative to 5,000 ClipByValue steps.Source: §6 · p. 8
  • These results support using both histories to improve pretraining quality and efficiency.Source: §6 · p. 8

Limitations

The timing comparison is a reference for the matched setup, since tensor shape, hardware, and training settings can change the measured speedup. The implementation also specifies how matrices are oriented, selected, calibrated, and transformed across updates.

Limitations

  • Tensor shape, hardware, and training settings can change the exact speedup, so the reported timing comparison is a reference for its matched setup.The comparison uses the same model, training settings, hardware type, and device count.Source: Limitations · p. 9

A.1 Orientation and Normalization

  • For standardized gradients, the routine orients the matrix so its left Gram matrix is the smaller Gram, then transposes the output back before restoring scale.Square matrices use the transposed orientation.Source: §A.1 · p. 10
  • The routine returns an unchanged gradient when its bound is at or below the squared threshold; an unrepresentable positive threshold or nonfinite intermediate raises an error before state commit.Its input is scaled before forming the Gram matrix.Source: §A.1 · p. 10

A.2 Spectral Norm Estimation

  • PowerTop estimates the clipped output’s spectral norm by alternating normalized products with X and X⊤ for three rounds, then measuring ∥Xx∥2.Its probe is seeded from parameter identity, successful update index, and probe purpose.Source: §A.2 · p. 10

B.1 Parameter Selection

  • The spectral branch selects two-dimensional logical parameters in the dense optimizer group, excluding embedding/output-name markers and position embeddings; other tensors use the norm branch.MoE projections and routers follow the same selection rule, while routing tensors assigned to the expert optimizer use the norm branch.Source: §B.1 · p. 11

B.2 Calibration and State Updates

  • The first 40 successful updates use global norm clipping at 1; a tensor enters the spectral branch only after its prior second moment and histories are ready and calibration is complete.Histories have a 10^-12 floor; missing gradients leave state unchanged, and controller states commit only after a successful AdamW update.Source: §B.2 · p. 11

B State and Communication

  • Gradient fragments are all-gathered into complete logical matrices for transformation, while factored moment summaries avoid repeatedly gathering the full previous second moment.Each rank writes its local optimizer slice after reconstructing and transforming the matrix.Source: §B.2 · p. 11

C.1 Parameter Counts

The architecture description gives dimensions for the shortcut blocks and MoE modules, and the parameter-counting caption clarifies how replicated and activated parameters are counted.

C.1 Parameter Counts

  • Each of the 10 shortcut blocks contains two MLA modules, two dense SwiGLU modules, and one MoE module; each MLA has 5,734,912 parameters.The hidden width is 1536, and the dense feed-forward width is 3072.Source: §C.1 · p. 11
  • Replicated parameters are counted once, and selecting 48 learned experts per MoE module gives an approximately 0.742B activated-backbone upper bound.The activated-backbone count uses the target of 32 learned experts per MoE module, whose actual count varies across tokens.Source: Table 6 · p. 12
  • Each MoE module contains 3,072 learned experts and 512 identity experts; routing targets 32 learned-expert activations per token, with the actual count varying by token.The module selects 48 candidates in total.Source: §C.1 · p. 12

C.2 Baseline Settings

The baselines and CAST variants differ in clipping rules, thresholds, and controller components. The training plot reports trailing 100-step means across the primary conditions and CAST ablations.

C.2 Baseline Settings

  • SSC-constant applies soft spectral clipping to selected dense matrices in raw coordinates with threshold 0.1 and five Newton–Schulz iterations, while unselected tensors pass through.It starts immediately.Source: §C.2 · p. 12
  • SSC-cosine replaces the fixed threshold with a schedule that starts at 10 and reaches 0.1 at iteration 5000, remaining there afterward.The iteration count includes skipped updates.Source: §C.2 · p. 13
  • GradientStabilizer, GradPower, ClipByValue, AGC, and ZClip use distinct norm, coordinate, or z-score clipping rules and the listed calibration settings.ZClip uses decay 0.97, z-score threshold 2.5, 25 calibration steps, and an additional norm cap of 1.Source: §C.2 · p. 13
  • The three CAST component variants remove historical scaling or replace the adaptive spectral threshold with a fixed or cosine threshold while retaining other controller components.Figure 9 plots trailing 100-step means from the first complete averaging window through iteration 5000 for ten primary conditions and three CAST ablations.Source: §C.2 · p. 13, Figure 9 · p. 13

D Validation Results

The validation-results section describes the labels, token-loss scope, and baseline comparisons represented in Table 7 and Figure 8.

D Validation Results

E Full Training Curves

Figure 9 presents the full training overview, while Figure 6 and Table 3 provide input-gradient traces and statistics for the conditions. Table 7 reports validation losses at iteration 5000 across all 20 subsets.

E Full Training Curves

  • Figure 9 shows the full training overview.Source: E Full Training Curves · p. 13
  • Figure 6 and Table 3 contain input-gradient traces and statistics for these conditions.Source: E Full Training Curves · p. 13
  • Table 7 reports token LM losses on all 20 validation subsets at iteration 5000 across ten conditions.Bold and underline mark the best and second-best values across both panels; GS denotes GradientStabilizer and Value denotes ClipByValue.Source: Table 7 · p. 14
Loading 2610.19761v1…