PaperlayerView exact revision on arXiv

Source-linked AI summary

SPAM: Spike-Aware Adam with Momentum Reset for Stable LLM Training

Tianjin Huang, Ziquan Zhu, Gaojie Jin, Lu Liu, Zhangyang Wang, Shiwei Liu

arXiv:2501.06842v2cs.LGcs.AIcs.CL

TL;DR

Gradient and loss spikes disrupt LLM training, while common recovery interventions can be resource-intensive. The paper investigates these spikes and proposes SPAM, combining momentum reset with spike-aware clipping. It reports that SPAM outperforms Adam and other state-of-the-art optimizers, including memory-efficient alternatives under memory constraints.

  • Problem

    Gradient and loss spikes disrupt LLM training, and checkpoint recovery can require frequent saves, monitoring, and repeated runs.

    Source: §1 · p. 1

  • Method

    SPAM addresses spikes through periodic momentum resets and spike-aware clipping, with sparse momentum for memory-efficient training.

    Source: §1 · p. 2, §3 · pp. 6–7

  • Results

    SPAM consistently outperformed Adam and other state-of-the-art optimizers, including GaLore and Adam-Mini under memory constraints.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

    The reported results support spike-aware optimization and sparse momentum as approaches for improving training performance while reducing memory costs.

    Source: §1 · p. 2, §7 · p. 10

    The theoretical analysis is explored in a limited context, assumes convexity, and depends significantly on the gradient bound G.

    Source: §2.3 · p. 6

Abstract

from arXiv · show

Large Language Models (LLMs) have demonstrated exceptional performance across diverse tasks, yet their training remains highly resource-intensive and susceptible to critical challenges such as training instability. A predominant source of this instability stems from gradient and loss spikes, which disrupt the learning process, often leading to costly interventions like checkpoint recovery and experiment restarts, further amplifying inefficiencies. This paper presents a comprehensive investigation into gradient spikes observed during LLM training, revealing their prevalence across multiple architectures and datasets. Our analysis shows that these spikes can be up to $1000\times$ larger than typical gradients, substantially deteriorating model performance. To address this issue, we propose Spike-Aware Adam with Momentum Reset SPAM, a novel optimizer designed to counteract gradient spikes through momentum reset and spike-aware gradient clipping. Extensive experiments, including both pre-training and fine-tuning, demonstrate that SPAM consistently surpasses Adam and its variants across various tasks, including (1) LLM pre-training from 60M to 1B, (2) 4-bit LLM pre-training,(3) reinforcement learning, and (4) Time Series Forecasting. Additionally, SPAM facilitates memory-efficient training by enabling sparse momentum, where only a subset of momentum terms are maintained and updated. When operating under memory constraints, SPAM outperforms state-of-the-art memory-efficient optimizers such as GaLore and Adam-Mini. Our work underscores the importance of mitigating gradient spikes in LLM training and introduces an effective optimization strategy that enhances both training stability and resource efficiency at scale. Code is available at https://github.com/TianjinYellow/SPAM-Optimizer.git

1 INTRODUCTION

Gradient and loss spikes disrupt LLM training and can make recovery resource-intensive. The paper investigates these spikes and introduces SPAM, which combines momentum resets and spike-aware clipping, with a sparse-momentum variant for memory-constrained training.

1 INTRODUCTION

  • Manual spike recovery can require checkpoint restarts, monitoring, and repeated runs, adding resource costs to LLM training.The paper motivates optimization-based mitigation against this intervention burden.Source: §1 · pp. 1–2
  • Gradient spikes can reach 1000× typical magnitudes, accompany frequent small loss bumps, and persist across layers, architectures, and datasets.The paper reports these patterns across models ranging from 60M to 1B parameters.Source: §1 · p. 2
  • Nullifying spiked gradients improves training performance, while Adam-like optimizers can accumulate their effects in momentum terms.The authors report empirical and theoretical demonstrations of this accumulation.Source: §1 · p. 2
  • Under memory constraints, SPAM uses sparse momentum to reduce memory costs and outperforms GaLore and Adam-Mini.Sparse momentum maintains and updates only a selected subset of momentum terms.Source: §1 · p. 2
  • SPAM combines periodic momentum resets and spike-aware clipping, and is reported to outperform Adam and Adafactor.The paper also reports superior performance across LLM sizes in pre-training and fine-tuning.Source: §1 · p. 2

2 GRADIENT SPIKES

The section defines gradient spikes by unusually large historical-relative gradient magnitudes, measured with the Gradient Spike Score, and investigates their prevalence, effects, and implications for Adam-like optimization. Spikes accompany loss bumps across models and settings, and experiments and preliminary analysis indicate harmful effects on training and regret bounds.

2.1 PRESENCE OF GRADIENT SPIKES DURING LLM TRAINING

  • Across LLaMA models from 60M to 1B, loss bumps coincided with gradient spikes, which appeared across layers, architectures, and datasets.The investigation also reports fewer spikes on cleaner SlimPajama than on C4, and especially frequent spikes in LayerNorm layers.Source: §2.1 · p. 3, Figure 2 · p. 4
  • Scaled Embed and Embed LN significantly reduced spike counts, but the evaluated architectural approaches did not eliminate gradient spikes.Other evaluated approaches offered little to no improvement.Source: §2.1 · pp. 3–4, Figure 4 · p. 4

2.2 EFFECTS OF GRADIENT SPIKES ON LLM TRAINING

  • Zeroing spiked gradients reduced final perplexity in LLaMA-60M on C4, indicating that spikes hindered effective training.The experiments detected spikes at multiple thresholds before setting them to zero.Source: Figure 5 · p. 5, §2.2 · p. 5
  • In an Adam simulation, a gradient spike’s effect persisted about 50 steps in the first moment and beyond 200 steps in the second moment.The authors attribute the second moment’s greater sensitivity plausibly to its larger decay rate and dependence on squared gradients.Source: §2.2 · p. 5

2.3 PRELINMINARY ANALYSIS WITH THEORY IMPLICATIONS

  • Under its limited convex analysis, spikes enlarge accumulated-gradient terms and loosen Adam-like regret bounds; the authors caution that assumptions may not fit non-convex settings.The authors describe the analysis as preliminary and note its bounded-gradient assumption may conflict with spikes.Source: §2.3 · p. 6

3 SPIKE-AWARE ADAM WITH WITH MOMENTUM RESET (SPAM)

SPAM combines periodic moment resets and spike-aware gradient clipping to limit the lasting effects of gradient spikes, while sparse momentum reduces optimizer memory use.

Momentum Reset

  • SPAM periodically resets Adam’s first and second moments every ∆T iterations, limiting the lingering influence of anomalously large gradients.It follows each reset with N steps of cosine warmup to mitigate potential instability.Source: §3 · p. 6

Spike-Aware Clipping

  • Spike-aware clipping detects gradients above threshold θ and scales them to manageable magnitudes while preserving their direction.For each detected spike, SPAM applies gi = sign(gi) · √θVi, using an on-the-fly approach based on Adam components.Source: §3 · p. 7
  • Unlike Adafactor’s update clipping, SPAM’s spike-aware clipping applies directly to gradients identified as spikes.Adafactor clips the whole weight-update matrix when its RMS exceeds 1; SPAM clips gradients whose magnitudes are significantly larger than √vi.Source: §3 · p. 7

Sparse Momentum

  • Sparse momentum updates and stores only a dynamically selected subset of moments, reducing memory use and computation during LLM training.The empirical analysis favors random parameter sampling and matching its frequency to the momentum-reset frequency.Source: §3 · p. 7

4 EXPERIMENTS

Across LLaMA pre-training and fine-tuning evaluations, SPAM outperforms the compared optimizers, including memory-efficient baselines; the experiments use matched memory costs for the fine-tuning comparison.

4 EXPERIMENTS

  • Comparisons include Adam, clipping variants, Adafactor, spike-mitigation methods, and memory-efficient optimizers such as Adam-Mini, GaLore, LoRA, and ReLoRA.The paper also evaluates supervised fine-tuning alongside pre-training.Source: §4 · pp. 7–8
  • The experiments compare optimizers on LLaMA models from 60M to 1B parameters using BF16, with learning rates swept from 1e−4 to 1e−3.SPAM uses reset interval ∆T = 500, warmup N = 150, and GSS threshold θ = 5000.Source: §4 · p. 8

4.1 PERFORMANCE OF LLM PRE-TRAINING

  • SPAM consistently outperforms the pre-training baselines across LLaMA model sizes and surpasses memory-efficient baselines at matched memory usage.The memory-efficient comparison reports validation perplexity and estimates total parameter-plus-optimizer-state memory in BF16.Source: §4.1 · p. 8, Table 3 · p. 8

4.2 PERFORMANCE OF LLM FINE-TUNING

  • On LLaMA2-7B fine-tuning across eight downstream tasks, SPAM substantially outperforms other memory-efficient methods and exceeds full fine-tuning.SPAM uses 0.25% density to match the memory cost of rank-8 low-rank baselines; Table 4 reports running GPU memory and means and standard deviations over 10 experiments.Source: §4.2 · p. 8, Table 4 · p. 9

5 ABLATION STUDY

The ablations favor random subset selection, a moderate reset interval and spike threshold, and warmup after each reset; overly frequent resets or an excessively low threshold can hurt performance.

Selection strategy for sparse momentum

  • Random subset selection performs significantly better than weight- or gradient-based selection for sparse momentum.The authors suggest random sampling may explore parameters more rapidly, while the alternatives may repeatedly select the same subset.Source: §5 · p. 9, Figure 6 · p. 9

Momentum reset interval ∆T

  • Performance improves as the reset interval decreases from 2500 to 500, but further shortening the interval degrades results.The experiment uses LLaMA-130M on C4, with 150 warmup steps and threshold θ = 5000.Source: §5 · p. 9, Figure 6 · p. 9

GSS threshold θ

  • Reducing threshold θ from extremely large values toward 1000 improves performance, but excessively small thresholds interfere with true gradients and degrade results.The authors report that spike-gradient clipping and momentum reset reinforce one another.Source: §5 · p. 9, Figure 6 · p. 9

Warmup steps N

  • No warmup after a momentum reset causes a significant performance drop, while approximately 150 warmup steps yields optimal performance.The ablation evaluates LLaMA-130M across warmup durations from 0 to 200 steps.Source: §5 · p. 9, Figure 6 · p. 9

6 RELATED WORK

Related work addresses training instability through architectural and initialization changes, while this paper positions momentum reset and memory-efficient optimization within distinct prior research lines.

Instability of Training Large Language Models

  • The related-work overview acknowledges that it may not cover all significant contributions in the field.This qualification applies to the review of stabilization techniques.Source: §6 · p. 10

Architectural approaches

  • Architectural approaches include Pre-LN, Embed LN, Embed Detach, DeepNorm, and αReparam, which target gradient behavior or training stability in different ways.The review notes that Embed LN may cause performance degradation (Scao et al., 2022).Source: §6 · p. 10

Initialization approaches

  • Initialization approaches include Scaled Embed, Scaled Initialization, and Fixup, which propose scaling or initialization changes to stabilize training.Fixup claims proper initialization can eliminate the need for LayerNorm (Zhang et al., 2019; Huang et al., 2020).Source: §6 · p. 10

Recent instability accounts

  • Sun et al. (2025) attribute deep-layer underperformance to Pre-LN and propose scaling down LayerNorm output as a remedy.Their account concerns the “Curse of Depth” in modern LLMs.Source: §6 · p. 10

Momentum reset

  • Earlier momentum-reset work addressed Nesterov acceleration or function-value increases, whereas SPAM applies resets to gradient spikes in large-scale non-convex LLM training.The cited prior work includes Gu et al. (2013), Nesterov (2013), and O’donoghue & Candes (2015).Source: §6 · p. 10

Memory-efficient optimizers

7 CONCLUSION

The paper studies gradient and loss spikes in LLM training and proposes SPAM, which uses momentum reset and spike-aware gradient clipping. It reports that SPAM outperforms Adam and other state-of-the-art optimizers, including memory-efficient alternatives when training with sparse momentum.

7 CONCLUSION

  • SPAM counters gradient spikes through momentum reset and spike-aware gradient clipping.The paper presents it as a novel optimizer for improving training stability and performance.Source: §7 · p. 10
  • Across various LLM model sizes, SPAM consistently outperformed Adam and other state-of-the-art optimizers.The conclusion reports a broad comparison but gives no numerical margin here.Source: §7 · p. 10
  • Under memory constraints, SPAM supports sparse momentum and outperforms GaLore and Adam-Mini.Sparse momentum maintains and updates only a subset of momentum terms.Source: §7 · p. 10

A STATISTICS ANALYSIS OF GRADIENT SPIKES ACROSS VARIOUS TYPES OF LAYERS

The analysis finds that spike frequency differs across LLaMA layer types and that spikes typically precede loss bumps. The pseudocode describes a masked, spike-aware Adam update with momentum resets and gradient clipping.

A STATISTICS ANALYSIS OF GRADIENT SPIKES ACROSS VARIOUS TYPES OF LAYERS

B LOCATIONS OF LOSS BUMPS AND GRADIENT SPIKES

C PSEUDOCODE

  • The pseudocode masks a parameter subset, detects and clips spiked gradients, resets moments, and applies a warmup-scaled Adam update.Its inputs include a mask density, spike threshold, momentum reset interval, and warmup steps.Source: C PSEUDOCODE · p. 16

D ARCHITECTURE AND HYPERPARAMETERS

The paper specifies LLaMA pre-training configurations and learning-rate tuning across model sizes. It reports that the momentum reset interval and warmup steps remain stable across those sizes.

D ARCHITECTURE AND HYPERPARAMETERS

  • All LLaMA models use a maximum sequence length of 256, batch size 512, and 131K tokens per batch.The training schedule uses 1,000 warmup steps followed by cosine annealing to 10% of the initial learning rate.Source: D ARCHITECTURE AND HYPERPARAMETERS · p. 16
  • For pre-training, learning rates are tuned from 1e-4 to 1e-3 in increments of 2×10-4, selecting the best by validation perplexity.The tuning is conducted for each method and model size from 60M to 1B.Source: D ARCHITECTURE AND HYPERPARAMETERS · p. 16
  • The momentum reset interval and warmup steps are reported to be insensitive to model size and stable across sizes.The paper also points to separate tables for SPAM pre-training and fine-tuning hyperparameters.Source: D ARCHITECTURE AND HYPERPARAMETERS · p. 16

E VISION TASKS

On vision tasks, SPAM performs on par with or better than AdamW. In a separate LLaMA-60M subset-selection study, random selection performs best among the tested strategies, and resetting all moments outperforms retaining overlapping moments.

E VISION TASKS

  • On ImageNet-1K, SPAM achieves performance on par with or better than AdamW on vision tasks.Experiments use ConvNeXt-Tiny and ViT-Tiny, trained for 120 epochs with the default ConvNeXt recipe.Source: Table 10 · p. 17, E VISION TASKS · p. 17

F MORE ABLATION STUDY OF SUBSET SELECTIOIN STRATEGIES

G EXPERIMENTS ON TIME SERIES DATA

On Weather time-series forecasting with PatchTST, SPAM’s advantage over Adam grows as the severity of introduced anomalies increases.

G EXPERIMENTS ON TIME SERIES DATA

H PROLONGED DETRIMENTAL EFFECTS OF GRADIENT SPIKES DURING REAL TRAINING

These analyses find that spikes disrupt moments for an extended period during LLaMA-60M training, while SPAM’s clipping performs favorably across tested thresholds and clipping comparisons, with slightly higher per-iteration runtime.

H PROLONGED DETRIMENTAL EFFECTS OF GRADIENT SPIKES DURING REAL TRAINING

I SENSITIVITY ANALYSIS OF HYPERPARAMETER θ ON LLM ARCHITECTURES

J GSS VS. DISTRIBUTION BASED CLIPPING

  • Gaussian-based clipping falls short of GSS-based clipping on LLaMA-60M and LLaMA-130M.The comparator assumes Gaussian gradient distributions; it estimates the full-dataset gradient with a batch-gradient moving average and flags elementwise deviations greater than 3δ for clipping.Source: J GSS VS. DISTRIBUTION BASED CLIPPING · pp. 19–20, Table 12 · p. 20

K GSS BASED CLIPPING VS. NULLIFYING

L COMPUTATIONAL ANALYSIS

  • SPAM incurs slightly higher per-iteration computational overhead than Adam, Adam-mini, and Adafactor, primarily from spike detection and sparse-mask gradient selection.The authors consider the overhead negligible relative to pre-training that can take dozens or hundreds of hours.Source: L COMPUTATIONAL ANALYSIS · p. 20, Table 15 · p. 21

Related papers

Loading 2501.06842v2…