SPAM: Spike-Aware Adam with Momentum Reset for Stable LLM Training
Authors: Tianjin Huang
Ziquan Zhu
Gaojie Jin
Lu Liu
Zhangyang Wang
Shiwei Liu
Venue: ICLR 2025 Poster (2025)
Keywords: Gradient Spikes
Spike-Aware Adam
LLMs
Primary area: foundation or frontier models, including LLMs
TL;DR: We present SPAM , an optimizer that mitigates gradient spikes in LLM training, improving stability and efficiency.

Abstract: 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× 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 a range of model scales. 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 submitted.

Published as a conference paper at ICLR 2025
SPAM: SPIKE-AWARE ADAM WITH MOMENTUM RE- SET FOR STABLE LLM TRAINING
Tianjin Huang!., Ziquan Zhu Gaojie Jin', Lu Liu', Zhangyang Wang", and Shiwei Liu
Department of Computer Science, University of Exeter, Exeter, UK DDepartment of Computer Science, University of Leicester, Leicester, UK Mathematical Institute, University of Oxford, Oxford, UK DDepartment of Electrical Computer Engineering, The University Texax at Austin, Austin, US Depaatment of Mathematics and Computer Science, Eindhoven University of Technology, Eindhoven, NL
ABSTRACT
Large Language Models (LLMs) have demonstrated exceptional performance across diverse tasks, yet their training remains highly resource-intensive and sus- ceptible to critical challenges such as training instability. A predominant source of this instability stems from gradient and loss spikes, which disrupt the Iearn- ing process, often leading to costly interventions like checkpoint recovery and experiment restarts, further amplifying inefficiencies. This paper presents prehensive investigation into gradient spikes observed during LLM training, re- vealing their prevalence across multiple architectures and datasets. Our analysis shows that these spikes can be up to x000 larger than typical gradients, sub- stantially 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 clip- ping. Extensive experiments, including both pre-training and fine-tuning, demon- strate that SPAM consistently surpasses Adam and its variants across a range model scales. Additionally, SPAM facilitates memory-efficient training by en- abling sparse momentum, where only a subset of momentum terms are main- tained and updated. When operating under memory constraints, SPAM outper- forms of 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 scale. Code is available com htt.:raaannniiiieeea SPAM Opt imiieeiie... zer.eerreo git
INTRODUCTION
Large Language Models (LLMs) have become fundamental in advancing state-of-the-art Al sys- tems. Scaling LLMs, such as GPT-3 (Brown 2020) and LLaMA (Touvron al. 2023). has show- cased unprecedented capabilities. However, training these large-scale models is fraught with chal- lenges, particularly training instability. major factor contributing
this instability the occur- rence of gradient and loss spikes during training. which disrupt the learning process at unpredictable intervals (Chowdhery et al. 2023 Zhang et al. 2022 Le Scao al. 2023).
While architectural innovations have been proposed to mitigate these issues Nguyen & Salazar 2019 Shoeybi et al. 2019 Zeng er al. 2022 Ding et al. 2021 Wang et al. 2024 Dettmers al.,
2021 Scao et | 2002 Takase al. 2023))
completely prevent the occurrence of spikes. In practice, the most widely adopted solution is to manually intervene by restarting training from a previous checkpoint and skipping data affected by the spike (Chowdhery et al. (023). This method resource-intensive, requiring frequent checkpoint saves, manual monitoring, and repeated experi- ment runs all inefficient and undesirable.
Moreover, the sheer scale of LLMs necessitates vast computational resources. For example, training ILaMA required over 2048 A A100-80GB GPUs (Tourron al., 2023), posing significant environ-

---

Published as a conference paper at ICLR 2025
mental and financial costs (Rillig al. 2023
Patterson et al. 2021
These challenges highlight
the need for more efficient training paradigms that reduce resource consumption without sacrificing performance.
(1) LLaMA-60M
(2) LLaMA-130M
19 (3) LLaMA-1B
30
45
Perplexity
40
27
17-
35
25
15
30
23
3K
5K
7K
7K
11K
15K
19K
30K
50K
70K
90K
Update Steps
Update Steps
Update Steps
Adam
Adam-Cradient-Clip
Adafactor
Adam-mini
SPAM
Figure 1: Perplexity of LLaMA models on C4 trained with various optimizers.
In this paper, approach the issue from optimization perspective rather than architectural one. We first conduct an in-depth investigation of loss and gradient spikes during the training of various LLM architectures, spanning models from 60M 1B parameters. Our study reveals several key observations:
Small yel frequent loss bumps: Although catastrophic loss spikes are rare, we observe frequent small loss bumps that can easily be overlooked without close scrutiny.
Gradient spikes accompanying loss bumps: These loss bumps, depiste small by their own, consistently accompanied by significant gradient spikes, whose magnitudes can reach up to 1000 greater than typical gradients. These spikes persist across layers, architectures, and datasets, even with established techniques applied.
Harmfulness of gradient spikes: By nullifying the spiked gradients, we observe notable improve- ments in training performance, confirming that these spikes have detrimental effect. Momentum- based optimizers, like Adam (Kingma) 2014,) [oshchilov] (017), suffer particularly from the ac- cumulation of these spikes in their momentum terms, as we demonstrate both empirically and theoritically.
Inspired by these findings, we introduce Spike-Award Adam with Momentum Reset (SPAM), an op- timizer designed to counteract the negative effects gradient spikes. SPAM introduces two key novations: (1) periodic reset of the first and second moments to eliminate the harmful accumulation of spiked gradients, and (2) identification and adaptive re-scaling of spiked gradients to manageable levels, preserving their directional information while mitigating their magnitude. We validate SPAM through extensive experiments, demonstrating its superior performance across various LLM sizes in both pre-training and fine-tuning tasks.
Furthermore, momentum reset enables the development of sparse momentum, where only a selected subset of momentum terms is computed and stored during training, drastically reducing memory costs. Our results show that SPAM surpasses leading memory-efiiciee optimizers such as GaLore (Zhao al. 2024) and Adam-Mini (Zhang et al. 2024a) with good margins, even under memory constraints.
Summary of Contributions:
Comprehensive analysis of gradient spikes across multiple LLM architectures, revealing their significant impact on training stability and performance.
Introduction of AM, a novel optimizer with momentum reset and spike- aakee clipping that outperforms existing methods like Adam and Adafactor.
A memory-efficien version of SPAM that leverages sparse momentum reduce memory usage while maintaining superior performance compared to state-of- the-art art memory-efficient optimiz-
2

---

Published as a conference paper at ICLR 2025
2 GRADIENT SPIKES
In this section, we formally define gradient spikes and then present the intriguing findings from our investigation into the training loss and gradient dynamics during LLM training.
Gradient spikes refer to a phenomenon that occurs during training where the magnitude of certain gradients significantly exceeds their historical values. To more precisely identify and analyze in- stances of gradient spikes, we introduce the Gradient Spike Score as a measurement of the deviation of a gradient's 's magnitude from its typical behavior over time. By quantifying this relative change, we can monitor the dynamics of gradients during training.
9T-1, gr be the sequence of gradient obtained during the training process from time step 0 to T. The Spike Score of the gradient the tth step, denoted as gs)) , is defined as the ratio of the magnitude of the gradient at that step to the average magnitude of the gradients across all steps:
Definition 2.1 (Gradient Spike Score). Let {go. 91-
19
GSS(g) =
T+I
2j
A gradient is considered a spiked gradient if its GSS(gi) exceeds predetermined threshold 8, i.e. (g:) > 0 indicating a significant increase from typical fluctuations, often amounting to increases of two three orders of magnitude.
2.1 PRESENCE OF GRADIENT SPIKES DURING LLM TRAINING
Building upon the above concepts, we further explore the presence of gradient spikes during LLM training. Specifically, monitor the gradients of the entire model over the initial 1, 000 training steps and identify gradient spikes using the condition GSS(g:) ) > 50. Our investigation encompasses two widely adopted LLM architectures, LLaMA (Touvron al. 2023 and Pythia (Biderman al., 2023), with model sizes varying from 60M to 1B parameters. Experiments were conducted on two datasets: the well-known C4 dataset (Raffel al. 2020) and cleaner high-quality dataset, SlimPa- jama (Soboleva et al.) 2023). Please refer to Appendix D for more details. Our key observations can be summarized as follows:
@ Loss bumps accompanying gradient spikes occur irregularly during LLM training. Al- though we do not observe severe loss spikes that lead catastrophic divergence (Takase et al.] 2023 Chowdheey et al. 2023), we do observe subtle loss bumps that happen quite frequently. For instance, Figure topp illustrates the training loss of LLaMA-60M, 350M, and 1B models, where several loss bumps can be during training, marked with red circles. We further investigate the model's 's gradients at these moments and observe that gradient spikes coincide with the loss bumps, as demonstrated in Figure bottoo.. While gradients remain small for most of the training, they suddenly become extremely large when loss spikes occur.
Gradient spikes widely presented in different layers, across different architectures, model sizes, and datasets. Overall, we observed many gradient spikes across all layer types, as detailed in Figure (4) and Appendix [ & B with LayerNorm layers, in particular, experiencing an exceptionally high frequency of spikes.
that models of varying sizes, from 60M to 1B, all exhibit gradient spikes. To verify whether architecture is the root cause of these spikes, we conducted experiments with Pythia-70M, which also suffers from numerous gradient anomalies, as shown in F Additionally, we found that gradient spikes occur even when using cleaner, high-quality datasets such as SlimPajama, although the frequency of spikes is reduced with this cleaner dataset.
Advanced spike mitigation approaches cannot completely eliminate gradient spikes. We also evaluate whether previously proposed techniques for addressing spikes can eliminate gradi- ent spikes. Specifically, we assess multiple approaches, including Scaled Initialization (Nguyen & Salazar 2019; Shocybi et al.) 2019). Embed LN (Dettmers et al.) 2021). Scaled Embed (Takasc al..|1200 2003), and Embed Detach (Zeng et al. 2022). The results in Figure 4 show that while some approaches perform better than others, they cannot completely eliminate gradient spikes. More specifically, find that Scaled Embed and Embed LN significantly reduce the number of gradient
We adopt the I. LaMa models used in Lialin al. 2023b
Zhao
2024

---

Published as a conference paper at ICLR 2025
spikes, while the other methods offer little to improvement, consistent with the findings reported in Takase et al. (0233)
Our observation of loss bumps likely relates to the edge of stability (EoS) phenomenon (Cohen ctt al. 202I), where the sharpness of the network hovers near the stability threshold for the remain- der of training while the loss continues decrease, albeit non-monotonically. However, the EoS phenomenon has not been extensively studied at the scale of LLMs. Moreover, our study reveals that these loss bumps have harmful effects LLM training, which were observed in previous studies.
(1) LLaMA-60M
(2) LLaMA-350M
(3) LLaMA-1B
Training Loss
2-
2
1.0
Q
0.5
2-
16-2
Spike
1.0
Spike
Spike
Gradient
0.5
Spike
Spike
0.0
Spike
Spike
0.5
2
1.0
200 600
800
1000
200
400
600 800
1000
200 400
1000
Update steps
Update steps
Update steps
Figure 2: Training loss lumps and their corresponding gradient spikes. Gradient trajectories are collected with I.LMMa-0M,, 350M, 1B models on C4 datasets. Gradient spikes detected using GSS(g.) > 50.
(1) LLaMA-C4
(2) LLaMM--SiiPaajama
(3) Pythia-C4
(4) LLaMA-C4
15K
Gradlent
Spikes LOc
-2
5K
250
500
750
1000
250
520
750
1000
250
520
750
1000
30.20 20 25 25 45-4045 30.25
Update Steps
Update Steps
Update Steps
Layers
Figure 3: Spike gradients present across different architectures and datasets. (1) (3): Plots of 100 randomly selected spike gradients (using GSS(g.) > 50) of LLaMa-60M and Pythia- 70M on C4 and SlimPajama datasets. (4): Number of spiked gradients every layers during the first KK steps in L.aMa-60M on C4.
2 11-3
20-3
3p-1
Gradient
0
-1
-2
500
1000
500
1000
Update Steps
Update Steps
Update Steps
Update Steps
Embed LN
Embed Detach
Scaled Embed
Scaled Initialization
Figure 4: Advanced spike mitigation approaches can not completely eliminate gradient spikes.
Gradient trajectories are collected with LLaMa-60M on C4. The spike gradient is detected via GSS(g,)) > 50.
4

---

Published as a conference paper at ICLR 2025
Second moment
Perplexity
34.0
First moment
1.0
33.0
0.3
33.6
Spice
Threshold
Update Steps
Splle Nerr
Update steps
W/o Spike
W/ Spike
Spike Gradient Remove
Perplexity of the final model after zeroing spiked gradients using various 9, GSS(g.) > 0. Experiments are conducted using LLaMa-60M on C4. Middle and Right: Impact of spiked Gradients on the first and second Moments. Simulated gradients (g: NNu, o²) are used to visualize the prolonged effects of gradient spikes on the first and second moments, with a large spike noise introduced the 30th step.
Figure 5: Left:
2.2 EFFECTS OF GRADIENT SPIKES ON LLM TRAINING
After identifying the presence of gradient spikes during training, a crucial question arises: are these gradient spikes detrimental perhaps counterintuitively. beneficial to the training of LLMs? To address this, we conducted a series of experiments as follows. Our findings confirm that gradient spikes are indeed harmful LLM training, exerting prolonged negative effects both the first and second moments, as discussed below.
Gradient spikes negatively impact LLM training. One direct way assess the impact of gradient spikes is by nullifying the spiked gradients during training and observing the final training perfor- mance. We first detect spiked gradients using various thresholds 0 and then those gradients to zero. Figure5p Left reports the results of ALaMA-60M. on C4. Surprisingly. zeroing out these spiked gradients leads improved model performance, evidenced by reduction in perplexity. This obser- vation clearly indicates that gradient spikes hinder effective training, and their removal is beneficial to overall model performance.
Gradient spikes have prolonged detrimental effects on the first and second moments. Due to the exponential averaging of the momentum mechanism, the influence over time. To demonstrate this, we conduct simulation experiment using Adam. In this experiment, we model the gradients as random variables drawn from a Gaussian distribution with mean μ = 0.1 and variance x² 0.1, i.c. 9i N(1, We sample gradients and track their corresponding moments over 200 steps, introducing gradient spike at step 30 with a large magnitude of 10. As shown in Figure Middle and Right, the spike's 's amplification persists, influencing both moments across subsequent steps. For example, it takes approximately 50 steps for the first moment to recover from the spike, while the second moment takes significantly longer, with the effect persisting beyond 200 steps. Two key factors plausibly contribute to this difference: (1) the second moment typically employs a larger exponential decay rate than the first (0.999 vs. 0.9); and (2) the second moment depends on the squared gradients, making it more sensitive to large spikes.
gradient spike decays slowly
2.3 PRELINMINARY ANALYSIS WITH THEORY IMPLICATIONS
We hereby provide very preliminary analysis help probe why gradient spikes have significant impact on the regret bound of Adam-like algorithms. We strictly follow the setting and notations used Alacaoglu et al. (2020). Specifically, referring Theorem the paper, the regret bound consists of two main terms:
R(T) < DTVT d
a aVI+ log
T
1/2 UT
Σ
2
20(1
B1) 1=1
V 32)
where 2f Gradient spikes directly affect these terms by increasing the magnitudes of the gradients Gt. In their Lemma 3, it is shown that the norm ||m 1/- depends on the accumulated
5

---

Published as a conference paper at ICLR 2025
gradients:
(1 31)
d
|m l 1/2
V(I- B2)
i-1
When gradient spikes occur, the values of Ij. become significantly larger for some j and i, which in turn increases the bound on |m= This enlargement propagates through the analysis, par- ticularly affecting the accumulation term Σ ar arllmal1?- in their Lemma 4, which is bounded by:
T ag Im
(1 -1-31)ay log
d
T
Σ
1/- I
Σ
2
v(1 B2) 7)
Here, gradient spikes increase Σ G2,5 significantly, especially in the coordinates where the spikes occur, leading to larger bound.
Finally, in the main regret bound (Equation (9) in the paper), these enlarged terms result in looser (larger) overall regret bound due to the presence of gradient spikes. The increased 6/// and Ein I. directly contribute to the regret bound becoming less tight. This theoretical implication highlights that while adaptive algorithms like AMSGRAD adjust learning rates based on gradient history, they may perform worse in terms regret when large gradient spikes are present due to the increased cumulative squared gradients and decreased effective learning rate.
important to note that our goal is not to claim theoretical innovations, but rather to quantita- tively assess how gradient spikes degrade Adam-like optimization, and that is only explored in very limited context. We would like to clarify the limitations of this analysis: (1) The analysis as- sumes convexity, which may not apply in non-convex settings (but is often mitigated by assuming condition or so). (2) The assumption ||gell S G, where G denotes the maxi- mum allowable gradient bound, may be in conflict with the presence of gradient spikes if G' is not sufficiently large to capture them. (3) There is significant dependence on G. and if G is set too high accommodate spikes, the constants in the regret bound grow disproportionately, potentially mak- ing the bound meaningless. Nonetheless, we find that our analysis aligns well with our experimental results, and we leave more rigorous theoretical exploration for future work.
SPIKE-AWARE ADAM WITH WITH MOMENTUM RESET (SPAM)
In this section, we introduce Swire-Aware Adam with Momentum Reset AM). Unlike previous solutions that introduce architectural innovations to mitigate the decremental effects of gradient spikes (Nguyen & Salazar 2019 Zeng et al., 2022 Dettmers et al. 2021 Takase et al., (003),, we attempt to address this issue from optimization perspective. Concretely, Reset and Spike-Aware Clipping into Adam to deal with gradient spikes. In addition, we introduce memory-efficient version of SPAM, which incorporates Sparse Momentum, significantly reducing the memory footprint during LLM training. Pseudocode of SP AM is in Algorithm
integrate Momentum
Momentum Reset. To mitigate the detrimental effects of gradient spikes on training stability, we introduce Momentum Reset. Momentum Reset involves periodically resetting the accumulated first and second moments used by adaptive optimizers such Adam. These optimizers rely on exponen- tial moving averages of past gradients to inform parameter updates. However, when a gradient spike occurs, it can significantly inflate these moments, causing the impact of the spike to persist over many subsequent iterations. By resetting the momentum terms at regular intervals of AT training iterations, we can prevent the lingering influence of anomalously large gradients on the optimizer's state. This practice ensures that parameter updates are based recent, more normal gradients rather than being skewed by gradient spikes. To mitigate potential instability caused by momentum reset, perform N steps (N = 150 by default) of cosine warmup following each reset operation.
Spike-Award Clipping. To further mitigate gradient spikes during intervals, we introduce Spike- Aware Clipping. While our initial experiments indicate that setting spiked gradients to zero can enhance performance, this approach completely removes the learning signal for those parameters, including valuable directional information critical to the optimization process. To address this, SPAM

---

Published as a conference paper at ICLR 2025
identifies gradients that exceed predefined threshold 0 and scales them to manageable value, preserving their directional information while controlling their magnitude.
Detecting gradient spikes using GSS defined in Definition 2.1 would require knowing and storing all gradients in advance-a method that is impractical for LLM training due to memory constraints. We adopt more memory-efficient on-the-fly approach by leveraging the components already cal- culated by Adam. Formally, we detect gradient spikes by identifying gradients 9i that meet the
following condition: g gi V 0] where V, the second moment of Adam and 0 is the
threshold used for the approximate GSS
Note that we only use GSS defined Definition 2.1
for the gradient spike analysis in Section 2 For real training, we employ the above approximation version. Since V; is essentially the moving average of g². this method efficiently identifies spikes without incurring additional overhead or the need to store the entire gradient history. Once detected, these spikes are clipped by scaling them to a manageable value. Specifically, for each spike gradient, we apply the operation: 9: = sign(g:) (g:) v0V. This technique is particularly useful when combined with Momentum Reset. By incorporating these strategies, SPAM effectively mitigates the negative impact of gradient spikes, improving training stability and performance.
Note that unlike the Update Clipping used in Adafactor (Shazeer & Stern 2018), which is applied to the whole weight update matrix when its Root Mean Square is larger than 1, our spike-aware clipping is directly applied to the spiked gradients gi whose magnitudes are significantly larger than its U/ e.g., > 50x.
Sparse Momentum. Momentum reset paves the way for the development of sparse momentum, a technique designed to reduce memory usage and computation during the training of LLMs. In traditional momentum-bast based optimizers, such as Adam, momentum is updated and stored for all parameters, which can be memoryyintensive for large-scale models. Sparse momentum offers more memory-efficien alternative by updating and maintaining only a dynamically selected subset of moments at each iteration. The percentange of selected subset is denoted by %d.
Key questions surrounding sparse momentum include how to effectively select parameter subsets, how to determine the sampling frequency, and whether to retain momentum for weights that are sampled consecutively Our empirical analysis shows that random sampling is the most effective strategy for selecting subsets of parameters. For the other questions, we find that they align well with the momentum reset strategy. Specifically, setting the sampling frequency to match the mo- mentum reset frequency, and resetting the momentum of all weights, even when they sampled consecutively, yield the most robust results.
Table 1: Comparison with various optimizers pre- -reiirannnn various sizes of LLaMA models C4. Perplexity is reported.
Table 2: Perplexity of Applying Advanced Techniques on LLaMA-60M.
Model Size
60M
130M
350M
1B
Optimizer
Perplexity
Adam-mini
34.10
24.85
19.05
16.07
Adam
34.09
Adam
34.09
24.91
18.77
16.13
Adam+Embed LN
33.61
Adam+Gradicet-Cllin
Vlll--aa Value
33.65
24.72
18.52
15.77
Adam+Embed Detach
34.48
33.33
24.88
18.51
15.22
Adafactor
32.57
23.98
17.74
15.19
Adam+Scaled Embed
33.87
SP AM
30.46
23.36
17.42
14.66
Adam+Scaled Initalization
34.29
Training Tokens
1.1B
2.2B
6.4B
11.6B
SPAM
30.46
EXPERIMENTS
To demonstrate the efficacy our proposed method, we conduct experiments on both pre-training and supervised fine-tuning using various sizes of the LLaMA model C4 dataset.
Baselines. We adopt several widely-used optimizers as our baselines. Since SPAM is built upon Adam, Adam serves our most direct baseline. We also incorporate two common gradient clip- ping approaches with Adam: (1) Value Clip, which clips all gradients when their absolute value exceeds threshold; and (2) Norm Clip, which scales the entire gradient if the L2 norm of the gradient vector exceeds a certain threshold. Additionally, we compare against another widely- optimizer, Adafactor (Shazeer & Stern 2018). In terms of spike mitigation techniques, eval-
7

---

Published as a conference paper at ICLR 2025
uate SPAM against previous approaches, including Scaled Initialization (gguyen & Salazar 2019
Shoeybi et al. 2019), Embed LN (Dettmers et al. 2021), Scaled Embed Takase et al. 2023 ) and
Embed Detach (Zeng et al.
2022). For mnemory-efhciem optimization methods, include Adam-
Mini Zhang et al. 2024a).
Galore Zhao et al. 2024). LoRA (Hu et al. 2021), and Rel.oRA (Lialin
et al. 2023a).
Architecture and hyperparameters.. Following (Lialin et al. 2023a Zhao et al. 2024), we conduct
our experiments using the LLaMA-based architecture with various sizes from 60M to IB parameters, incorporating RMSNNrm (Shazeer] 2020) and SwiGLU activations (Zhang & Sennrich, 2019). For each model size, we use the same set of hyperparameters across methods, varying only the learning rate, where we sweep over a set of learning rates from le - 4 to 1e 3, incrementing by 2e for each optimizer. All experiments are conducted using the BF16 format. We set clip threshold as 1 and 1e 3 for Norm Clip and Value Clip, respectively, following the setting in Takase et al. (2023). We hyper-parameters for Adafactor following the original paper (Shazeer Stern 2018) where E = 10 10-30 62 = 10-3 and = 1.0. For SPAM, we set reset intervals AT 500, Ir 5000. Detailed descriptions of our task setups and
warmup step N = 150 and GSS threshold
hyperparameters provided in the Appendix D
4.1 PERFORMANCE OF LLM PRE-TRAINING
Pre-training. We report the training curves of various LLaMA models on the C4 dataset as well as the final perplexity in Figure | and Table [ respectively. Overall, we observe that SPAM con- sistently achieves superior performance. As a memory-efficient approach, Adam-min performs on par with Adam, consistent with the results reported in Zhang et al. (2024a). Commonly used gradient clipping techniques such as Value Clip and Norm Clip improve performance over Adam, with the latter achieving slightly better results. Adafactor further outperforms the aforementioned approaches, demonstrating its effectiivneess SPAM consistently outperforms all baselines across various LLaMA model sizes, highlighting the benefits of integrating momentum reset and spike- aware clipping techniques. All spike mitigation approaches fall short of SP AM as shown in Table Additionally, Appendix Ejshows that SPAM can perform on par with or better than Adam in vision tasks and Appendix G shows that SPAM outperforms Adam in time series forescasting tasks.
Table 3: Comparison with memory-eflicien algorithms pre-training various sizes of LLaMA models on C4 dataset. Validation perplexity reported, along with memory estimate the total of parameters, optimizer states based on BF16 format.Ths results of GaLore, Full-Rank, LoRA and ReLoRA are obtained from Zhao
ct al. (024).
60M
130M
350M
1B
Adam
34.06 (0.36G)
(0.76G)
18.80 (2.06G)
15.56 (7.80G)
ReLoRA
37. 07 (0.36G)
29.37 (0.80G)
29.08 (1.76G)
18.33 (6.17G)
LoRA
34.99 (0.26G)
33.92 (0.54G)
25.58 (1.08G)
19.21 (6.17G)
GaLore
34.88 (0.24G)
25.36 (0. (2.)2))
18.95 (1.22G)
15.64 (4.38G)
SPAM
32.39 (0.24G)
23.98 (0.52G)
18.28 (1.22G)
15.60 (4.38G)
Training Tokens
1.1B
2.2B
6.4B
11.6B
Memory-efficient Pre-training. We evaluate SPAM by specifying d% such that its memory usage, including both parameters and optimizer states, matches that of Galore. For Galore, LoRA, and 128, 256, 256, 512 for the 60M, 130M, 350M, and 1B
ReLoRA baselines, we set the ranks models, respectively, following the setup in Galore (Zhao et al. 2024).. The results in Table 3 show that SPAM consistently outperforms all the baselines by a good margin, demonstrating its effectiveness as memory-efficte optimizer.
512
4.2 PERFORMANCE OF LLM FINE-TUNING
In this section, we evaluate the effectiveness of SPAM for supervised fine-tuning. Following Li eet al. (2024a), we fine-tune LLaMA2-7B
Commonsensel7001 (Hu et al.] 2023) and test on downstream tasks. We do not apply layer-wise weight updates for GaLore and SP SPAM. The rank is set to 8 for all low-rank baselines. Correspondingly. the density of SPAM is set to 0.25% to maintain a comparable memory cost. The results are reported in Table We observe that SPAM substantially outperforms other memory-efficient methods, exceeding full fine-tuning by notable margin.
8

---

Published as a conference paper at ICLR 2025
Table 4: Fine-tuning performance of LLaMa2-7B on various downstream tasks. The "Mem." de- notes the running GPU memory. The mean and standard deviation of 10 repeated experiments are reported.
Method
Mem.
BoolQ
PIQA
SIQA
HellaSwag
WinGGradde
ARC-e
ARC-c
OBQA
Avg.
79.110.1
51.3:0005
58.510.02
74.8102
48.210001
36.210.2
63.4:0.1
LORA
29G
7584104
790:00.1
56.3=0.1
50.9:004
79.6:0.2
7766::11
46.9:00
34.4:0.
63.7:002
Gal.rr
366
82.8+0.7
78.40022
558-04
56.3+0.5
79.0+0.1
75.++0.4
46.2+0.5
34.2+0.1
63.6+04
SPAM 0.25%
360
85.010.2
78.9102
55.7=0.2
57.810.1
78.9102
76.5100.
47.31.02
35.110.3
64.4:022
100%)
87.140.2
79.5+0.1
58.3-0.1
58.14004
83.3:0.2
48.6+0.11
40.1+0.2
66.7=0.1
5 ABLATION STUDY
Random
27
27
Perplexity
40
Max, Weight
26
26
28
25-
25
26
35
24
24
0.25 0.5
0.750.9
2.5K 1K 500 250 100 50
None 10K 1K
100
10
0 10
50
100 150 200
(1) Sparsity (1-d)
(2) Interval AT
(3) Theshold
(4) Warmup Steps N
Figure 6: Ablations for sparse subset selection strategy, momentum reset inteval, GSS threshold and warmup steps. "None" denote that the spike-aware clipping is not applied.
Selection strategy for sparse momentum. Many strategies have been proposed to select subsets of parameters for sparse training, such random selection (Liu al. (2022a), max weight magnitude (Mocanu et al. (2018), and max gradient magnitude (Evci et al. (020). Among these strategies, the most effective approach for sparse momentum training remains unclear. To investigate this, we con- duct experiments with LLaMA-60M on the C4 dataset. The results are reported in (igure(((() In- terestingly, we find that randomly selecting subsets parameters performs significantly better than the other two strategies for our sparse momentum. One plausible explanation for this discrepancy that random selection allows for rapid exploration across model parameters, whereas gradient- weight- based strategies might be confined to the same subset of parameters during training.
Momentum reset interval AT. To investigate the impact of interval AT, we conduct experiments based (LaMA-130M and C4 with varying AT fromm 50 to 2500. The warmup steps is set to 150 and the thresthold 0 is set to 5000. The results are reported in Figure (2). We observe a performance improvement as the interval AT decreases from 2500 to 500. However, when AT is further shortened, performance begins to degrade. This suggests that while momentum resets can enhance performance, excessively frequent resets may be detrimental overall results.
GSS threshold 0. Threshold 0 decides which gradient are detected as spikes. To illustrate the T with varying from 20000 to 10. The warmup steps is set to 150 and the interval AT is set to 500. We observe that performance improves
impact of on SPAM, we present the results of LLaMA-130M
is reduced from extremely large values smaller values, such as 1000, indicating that spike gradient clipping and momentum reset techniques have mutually reinforcing effect. However, excessively small may interfere with the true gradient, ultimately leading degradation in performance.
Warmup steps N. We assess the impact of the warmup procedure following each momentum reset by presenting the performance of LLaMA-130M with different warmup steps, ranging from 0 to 200, in Figure (4). The results indicate a significant performance drop when no warmup is applied (N 0), compared to when warmup is used. In addition, performance reach to optimal when the warmup duration is set to approximately 150 steps.
6 RELATED WORK
Instability of Training Large Language Models. LLMs well-known for their training insta- bility (Molybog et al. 2023). often experiencing irregular loss spikes that can lead to catastrophic divergence (Chowdhery al. 2023). To address this issue, researchers have developed various sta-

---

Published as a conference paper at ICLR 2025
bilization techniques. While outline several key approaches, we acknowledge that this overview may not cover all significant contributions in the field.
One prominent approach involves architectural modifications, Xiong al. (2020) demonstrated that using Post-LN Transformers leads to larger gradients near the output layer, resulting in training instability, especially with large learning rates. In contrast, Pre-LN helps maintain well-behaved gradients during initialization, promoting more stable training. Embed LN, introduced by Dettmers) al. (202I), adds an additional LayerNorm after the embedding layer to improve stability, though it may cause performance degradation, as noted by [Scao et al. (2022). Embed Detach, proposed by Ding et al. (202T) and further extended by Zeng al. (2022) for LLMs, addresses loss spikes by shrinking embedding gradients. DeepNorm, developed by Wang et al. (2024), enhances stability in deep Transformers by scaling up the residual connection before applying LayerNorm. Addition- ally, aRReparam (Zhai et al.) 2023) re-parameterizes all linear layers using spectral normalization to
prevent attention entropy collapse.
Another set of approaches focuses improving initialization to mitigate training instability. Scaled Embed, proposed by Takase et al. (2023), scales up embeddings to stabilize LayerNorm gradients. Scaled Initialization (Nguyen & Salazar 2019) introduces parameter initialization strategy using smaller normal distribution // (0. 2/ 2/5d/ v2N) to stabilize training dynamics. Additionally, Fixup Zhang et al. 2019; Huang et al. 2020) claims that proper initialization can entirely eliminate the need for LayerNorm.
Very recently, Sun et al. (2025) introduced the concept of "the Curse of Depth", emphasizing that nearly half of the deep layers in modern LLMs underperform relative to expectations. They attribute this issue to the widespread adoption of Pre-LN (Li et al. 2024b) and propose scaling down the LayerNorm output
effective solution.
Momentum Reset. Momentum reset is not a new approach. It has been used in Gu et al. (2013) Nesterov 2013) to solve the rippling behavior of Nesterov's Accelerated Gradient (NAG) Nesterov 1983) in the high-momentum regime, particularly in the context of convex optimization problems. O'donoghue & Candes (2015) further proposed adaptive reset where the momentum will be reset when an increase in the function value is observed. Unlike these earlier work, we leverage mo- mentum reset to mitigate the detrimental effects of gradient spikes that arise during the training of billion-parameter language models, which present a large-scale, non-convex optimization challenge.
Memory-E -ffficient Optimizers. There have been several efforts to reduce Adam's memory foot- print. SM3 (Anil et al. (019)) a lightweight variant of AdaGrad (Duchi et al.] (011). selects the learning rate for the i-th parameter by taking the minimum value from set of candidates, each associated with the maximum squared gradient under predetermined cover. Adafactor (Shazeer & Stern 2018) and its variant CAME (Luo et al., 2023) utilize non- negative low-rank factorization over Adam's 's momentt estimate, U. Adam mini (Zhang et al. al.] (2024a) partitions the parameters into blocks and assigns a single learning rate v to each block to reduce memory. Similar approaches were proposed in (Zheng & Kwok] [019: Ginsburg et al. 2019). Low-precision optimizers stud- ied in (Dettmers et al.) (021). Recently, GaLore (Zhao et al.) [2224] Zhang et al. 2024b) enables the full-parameter training of LLMs through low-rank gradient updates.
CONCLUSION
In this paper, we presented comprehensive study of gradient and loss spikes in LLM training. demonstrating their detrimental impact on training stability and performance across a variety of architectures and datasets. To address this issue, we propose Spike-Aware Adam with Momentum Reset (SP AM), a novel optimizer designed to counteract gradient spikes through momentum reset and spike-aware gradient clipping. The effectiveness of SPAM backed up with extensive experiments across various L.M model sizes, where SPAM consistently outperformed Adam and other state-of- the-art optimizers by a good margin. When operating under memory constraints, SPAM motivates the feasibility of sparse momentum training, outperforms state-of-the-at tht memory-effiiien optimizers such as GaLore and Adam-Mini.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work used the Dutch national e-infrastructure with the support of the SURF Cooperative using the funding of the projects EINF-12538, EINF-10925 and NWO-2023.0277 This work is partially supported by the SLAIDER project. We would like to express our deepest gratitude the anony- mous reviewers whose insightful comments and suggestions significantly improved the quality of this paper. Shiwei Liu is supported by the Royal Society with the Newton International Fellowship.
REFERENCES
Ahmet Alacaoglu, Yura Malitsky, Panayotis Mertikopoulos, and Volkan Cevher. new regret anal- ysis for adam-type algorithms. In International conference on machine learning, pp. 202-210. PMLR, 2020.
Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimiza- tion. Advances in Neural Information Processing Systems, 32, 2019.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2397-2430. PMLR, 2023.
Tom B Brown. Language models are few-shot learners. arXiv preprint 200i::2005.11166 2020.
Pratik Chaudbari and Stefano Soatto. Stochastic gradient descent performs variational inference, converges limit cycles for deep networks. 2018 Information Theory and Applications Work- shop (ITA), pp. 1-10. IEEE, 2018.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240): 1-113, 2023.
Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs the edge of stability. In International Conference Learning Representations, 2021.
Tim Dettmers. Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. arXiv preprint wxXiv2210.00880) 2021.
Ming Ding, Zhuoyi Yang, Wenyi Hong, Wendi Zheng, Chang Zhou, Da Yin, Junyang Lin, Xu Zou, Zhou Shao, Hongxia Yang, al. Cogview: Mastering text-to-image generation via transformers. Advances in neural information processing systems, 34: 19822-18855 2021.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Utku Evei, Trevor Gale, Jacob Menick, Pablo Samuel Castro, and Erich Elsen. Rigging the lottery: Making all tickets winners. In International conference on machine learning, pp. 2943-2952.
PMLR, 2020.
Boris Ginsburg, Patrice Castonguay, Oleksii Hrinchuk, Oleksii Kuchaiev, Vitaly Lavrukhin, Ryan Leary. Jason Li, Huyen Nguyen, Yang Zhang, and Jonathan M Cohen. Stochastic gradi- ent methods with layer-wise adaptive moments for training of deep networks. arXiv preprint arXiv: 2019.
Ming Gu, Lek-Heng Lim, and Cinna Julie Wu. Parnes: a rapidly convergent algorithm for accu- rate recovery of sparse and approximately sparse signals. Numerical Algorithms, 64(2)33213377, 2013.
Edward J Hu, Yelong Shen. Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.00 09680, 2021.
11

---

Published as a conference paper at ICLR 2025
Zhiqiang Hu, Lei Wang, Yihuai Lan, Wanyu Xu, Ee-Peng Lim, Lidong Bing, Xing Xu, Soujanya Poria, and Roy Ka- Wei Lee. Llm-adapters: An adapter family for parameter-efficient fine-tuning of large language models. arXiv preprint 2rX1v:2304.01933 2023.
Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving transformer opti- mization through better initialization. In International Conference on Machine Learning, pp- 4475 4483. PMLR, 2020.
Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6990, 2014.
Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilié, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, et al. Bloom: A 176b- parameter open-access multilingual language model. arXiv preprint 1..::2//100100 2023.
Pengxiang Li, Lu Yin, Xiaowei Gao, and Shiwei Liu. Owlore: Outier- layerwise sam- pled low-rank projection for memory-efficient IIm fine-tuning. arXiv preprini arXiv:2405./883 2024a.
Pengxiang Li, Lu Yin, and Shiwei Liu. Mix-ln: Unleashing the power of deeper layers by combining pre-In and post-In. arXiv preprint arXiv:2412.13775, 2024b.
Vladislav Lialin, Sherin Muckatira, Namrata Shivagunde, and Anna Rumshisky. Relora: High- rank training through low-rank updates. The Twelfth International Conference on Learning Representations, 2023a.
Vladislav Lialin, Namrata Shivagunde, Sherin Muckatira, and Anna Rumshisky. Stack more lay- ers differently: High-rank training through low-rank updates. arXiv preprint arXiv:2307.05695, 2023b.
Shiwei Liu, Tianlong Chen, Xiaohan Chen, Li Shen, Decebal Constantin Mocanu, Zhangyang Wang, and Mykola Pechenizkiy. The unreasonable effectiveness of random pruning: Return of the most naive baseline for sparse training. arXiv preprint arXiv:2202. 02640,,,.2,,.... 2022a.
Zhuang Liu, Hanzi Mao, Chao- Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF conference on computer vision and paltern recognition, pp. 1 11976-1 1986, 2022b.
I Loshchilov, Decoupled weight decay regularization. arXiv preprint arXiv: 2017.
Yang Luo, Xiaozhe Ren, Zangwei Zheng, Zhuo Jiang, Xin Jiang, and Yang You. Came: Confidence- guided adaptive memory efficient optimization. arXiv preprint urXiv:2307.02044, 2023.
Stephan Mandt, Matthew Hoffman, and David Blei. A variational analysis of stochastic gradient algorithms. In International conference on machine learning, pp. 354-363. PMLR, 2016.
Decebal Constantin Mocanu, Elena Mocanu, Peter Stone, Phuong H Nguyen, Madeleine Gibeseu, and Antonio Liotta. Scalable training of artificial neural networks with adaptive sparse connec- tivity inspired by network science. Nature communications, 91):23383, 2018
Igor Molybog, Peter Albert, Moya Chen, Zachary DeVito, David Esiobu, Naman Goyal, Punit Singh Koura, Sharan Narang, Andrew Poulton, Ruan Silva, et al. A theory on adam instability in large- scale machine learning, arXiv preprint arXiv:2304.0988 2023.
Yu Nesterov. Gradient methods for minimizing composite functions. Mathematical programming, 140(1): 2013.
Yurii Nesterov. A method for solving the convex programming problem with convergence rate o (1/k2). In Dokl akad nauk Sssr, volume 269, pp. 543, 1983.
Toan Q Nguyen and Julian Salazar. Transformers without tears: Improving the normalization of self-attention.. arXiv preprint arXiv: 1xXvv.99/0.05895, 2019.
Yuqi Nie, Nam H Nguyen, Phanwadee Sinthong, and Jayant Kalagnanam. A time series is worth 64 words: Long-term forecasting with transformers. arXiv preprint (tt::22/1.1473300 2023.
12

---

Published as a conference paper at ICLR 2025
Brendan O O'donoghue and Emmanuel Candes. Adaptive restart for accelerated gradient schemes. Foundations of computational mathematics, 15:715-734 72215---,,00 2015.
David Patterson, Joseph Gonzalez, Quoc Le, Chen Liang. Lluis- Muiis-Miuuel Munguia, Daniel Rothchild, David So, Maud Texier, and Jeff Dean. Carbon emissions and large neural network training. arXiv preprint arXiv:2 10350, 2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, (1(140)11667, 2020.
Matthias C Rillig, Marlene Agerstrand, Mohan Bi, Kenneth A Gould, and Uli Sauerland. Risks and benefits of large language models for the environment. Environmental Science & Technology, 57 (9):3464 3466, 2023.
Teven Le Scao, Thomas Wang. Daniel Hesslow, Lucile Saulnier, Stas Bekman, M Saiful Bari, Stella Biderman, Hady Elsahar, Niklas Muennighoff, Jason Phang, al. What language model train if you have one million gpu hours? arXiv preprint arXiv:2210.15424 2022.
Noam Shazeer. Glu variants improve transformer. arXiv preprini urXiv:2002.05202. 2020.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596-4604 PMLR, 2018.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model par- allelism. arXiv preprint arXiv:/909.0805: 2019.
Umut Simsekli, Levent Sagun, and Mert Gurbuzbalaban. A tail-index analysis of stochastic gradient noise in deep neural networks. In International Conference on Machine Learning, pp. 5827-5837. PMLR, 2019.
Daria Soboleva, Faisal Al-Khateeb, Robert Myers. Jacob R Steeves, Joel Hes- A 627B token cleaned and dedu-
tnsss, and Nolan Dey. SlimPajama:
plicated version of RedPajama.
https: htt/://wwwwwbrbas//// net/blog/
-aaaaaamm7--tooeeeeeee hww.aaaeneeeaanee---------e-a-e- 2023. URL httt::///agggggaa...eooaaattttte
s1 slimpa jamaaa-aaaaaa
www.rreedaaaao jama 627B
Wenfang Sun, Xinyuan Song, Pengxiang Li, Lu Yin, Yefeng Zheng, and Shiwei Liu. The curse of depth in large language models. arXiv preprint arXiv:2502.05795, 2025.
Sho Takase, Shun Kiyono, Sosuke Kobayashi, and Jun Suzuki. Spike no more: Stabilizing the pre-training of large language models. arXiv preprint arXiv:2312.1660 2023.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pp. 10347-10357. PMLR, 2021.
Hugo Touvron, Thibaut Lavril, Gautier Izacard. Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Ar- mand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. arXiv preprint 230. 2023.
Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang. Dongdong Zhang, and Furu Wei. Deep- net: Scaling transformers to 1, 1,000 layers. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pp. 10524-10533 PMLR, 2020.
Aohan Zeng, Xiao Liu, Zhengxiao Du, Zihan Wang, Hanyu Lai, Ming Ding, Zhuoyi Yang. Yifan Xu, Wendi Zheng, Xiao Xia, et al. Glm-130b. An open bilingual are-trained model. arXiv preprint arXiv:2210.00 14.4444, 2022.
13

---

Published as a conference paper at ICLR 2025
Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Joshua M Susskind. Stabilizing transformer training by preventing attention entropy collapse. In International Conference Machine Learning, pp. 10770-40803, PMLR, 2023.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Infor- mation Processing Systems, 32, 2019.
Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. arXiv preprint uxXiv:991.0993: 2019.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christo- pher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022.
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Yinyu Ye, Zhi-Quan Luo. and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more. arXiv preprint arXiv:2406.1679, 2024a.
Zhenyu Zhang, Ajay Jaiswal, Lu Yin, Shiwei Liu, Jiawei Zhao, Yuandong Tian, and Zhangyang Wang. Q-galore: Quantized galore with int4 projection and layer-adaptive low-rank gradients. arXiv preprint N.XKv::24770020000 2024b.
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient IIm training by gradient low-rank projection. arXiv preprint arXiv:2403.33507, 2024.
Shuai Zheng and James T Kwok. Blockwise adaptivity: Faster training and better generalization in deep learning. arXiv preprint arXiv: 1905 1909,9999, 2019.
14

---

Published as a conference paper at ICLR 2025
A STATISTICS ANALYSIS OF GRADIENT SPIKES ACROSS VARIOUS TYPES OF LAYERS
It is important examine whether gradient spikes exhibit preference for certain layers. To do so, we report the number of gradient spikes across various types of layers and the ratio of gradient spikes the number of parameters five types of layers: Embedding Layer, Attention Layer, FFN Layer, LayerNorm Layer, and LM Head Layer. The experiments were conducted with LLaMA-60M on the C4 dataset, with gradient spikes detected over 1000 training steps. The detailed statistics are provided in Table We observe the following: The Embedding Layer exhibits the highest number of gradient spikes, also it has the largest parameter count. @ The LayerNorm Layer, however, experiences exceptionally high frequency of spikes, even with the smallest number of parameters.
Table 5: Number and Ratio of Gradient Spikes in each layer style of LLaMA. # #Siikes are collected from 1000 training steps. Experiments are conducted with LLaMA-600 on C4.
Module Name
#Total Spikes
#Total Params
#Total Spikes
Total Params
Embed
11954001
16384000
0.729
Attention
86302
8388608
0.010
FFN
105415
16908288
0.006
LayerNorm
949302
8704
109.06
L.M_
13893
16384000
0.000848
B LOCATIONS OF LOsS BUMPS AND GRADIENT SPIKES
To further investigate the correlation between loss bumps and gradient spikes, we present the lo- cations of gradient spikes associated with the loss bumps Table 6 The results reveal two key findings: 1 Gradient spikes are presented in different layers associated with the loss bump; Ø Gra- dient spikes typically occur before loss bumps, indicating that these gradient spikes may trigger loss bumps.
Table 6: Location of Spike Gradient at Each Layer for Different Tasks. The spike gradient is detected via GgS(g.) 50. The experiments are based on LLaMA-60M and PYthia-70M.
Model
Training Step
Loss Bump Occurs
0th
5th
Training Step Spilke Gradient Occurs
10th
15th
20th
25th
311th
35th
:04t
Layer 454h
50th
56th
196
202
197
196
196
197
197
197
197
197
LLMMA-6M (C4)
2005
1917
205
198
198
198
278
198
198
201
199
202
198
205
207
206
2006
2016
2006
206
MaMM-6GKK
325
394
207
206
207
208
207
210
207
200
394
393
207
328
358
571
573
337
337
578
577
571
Pythiia700: (C4)
578
677
576
15

---

Published as a conference paper at ICLR 2025
C
PSEUDOCODE
Algorithm 1: SPAM
Input: A layer weight matrix E Rm learning rate decay rates 81 0.9, 32 0.999, initial parameters we, randomly initialize mask M with d density for each layer, the first moment m. the GSS, momentum rerest interval AT', warmup scale total steps N, small constant = x 10 T total training steps.
second moment e, threshold
Output: optimized parameters wy.
while t < 7 do
Get gr R==* f- ( (w+)
D.Gnerrae Gradients
warmap scale =1 Cosi Cos..eeAnnalllggMMMlll....
if Mod (t, AT) = 0 then
M random. rand(f.shape) < d
Random initialize the binary mask reset the first reset the second moment to zero
zeros (f
zero
v - zeros li)e(f
Spike M = 9: [M] +2 > 0
b Detect spiked gradients
if sum Spike_M) > 0 then
L 9. Spiike M] = sign(gn
(M] [annggMM|SSpkeeee
v v[Spike.M)
> Spike Gradients CLIP
m, = 8m + gg
V. 32v=-1
m; I-
v, Vi
I-33
Wr W1-1 warmup.scale Vesto
[=+1
Return: optimized parameters wtr
D ARCHITECTURE AND HYPERPARAMETERS
We introduce details of the LLaMA architecture and hyperparameters used for pre-training, follow- ing [Lialin al. 2023a); Zhao al. (2024). Table 7 shows the most hyperparameters of LLaMA models across model sizes. We use a max sequence length of 256 for all models, with a batch size of 512, with batch size of 131K tokens. For all experiments, we adopt learning rate warmup of 1000 training steps, and use cosine annealing for the learning rate schedule, decaying to 10% of the initial learning rate.
Table Configurations of LLaMA models used in this paper. Data amount are specified in #token.
Params
Hidden
Intermediate
Heads
Layers
Steps
Data amount
60M
512
1376
8
10K
1.3B
130M
768
2048
12
12
20K
2.6B
350M
1021
2736
16
21
60K
7.8B
B
2048
5461
24
32
89K
11.6B
For all methods across each model size (from 60M to 1B), we tune the learning rates from 1e 4 to 1e-3 with an increasing step of 2 x 10-4 for pre-training tasks, and the best learning rate selected based on the validation perplexity. We find that the hyperparamcters, Interval AT' and warmup step N, are insensitive to model size and remain stable with the same learning rate across different model sizes. The detailed hyperparameter of SP AM pre-training and fine tne-tuning are reported in Table and T
16

---

Published as a conference paper at ICLR 2025
Table Hyperparameters of SPAM for pre-training experiments in this paper.
Hyper-Parameters
LLMMA-6OM
LLaMA-130M
LLaMA-350M
LLaMA-1B
Standard Pretraining
Learning rate
1c 3
8e
4c 4
2c
Interval AT
500
500
Threshold
5000
5000
5000
5000
Warmup steps N
150
150
150
150
Memory-Lfficient Pretraining
Learning rate
4e 3 3
4e -
2e 3
5e
Interval AT
500
500
500
1000
Threshold
5000
5000
5000
5000
Warmup steps N
150
150
150
300
Table 9: Hyperparameters of SPAM for fine-tuning experiments this paper.
Hyper-Parameters
LLLMA2-7B
Standard Fine- Fine-tuning
Learning rate
5
Interval AT'
1000
Threshold 0
5000
Warmup steps N
300
Memory Efficient Fine-tuning
Memory-Efficienl
Learning rate
1c 4
Interval AT
250
Threshold 0
5000
Warmup steps
E VISION TASKS
We further evaluate SPAM on vision task. Specifically, we conducted experiments on ImageNet-1K using Con-NeXt-Tiny (Liu et al. 2022b) and ViT-Tiny (Touvron et al. 2021). We adopt the default training recipe from the official code of ConvNeXT and train all models for 120 epochs. We set
AT = 25K, N 20 and 5000 for SPAM. The results in Table I0 demonstrate that SPAM can achieve on par or better performance than vanilla AdamW.
Table 10: SPAM performs on par or better than Adam W on vision tasks.
Optimizer
Model
Metric
25% steps
50% steps
75% steps
100% steps
AdamW
ConNeXt-T
Test Acc (f)
68.15
74.00
78.83
80.89
SPAM
ConNeXt-T
Test Acc (1)
68.36
73.63
78.85
81.04
AdamW
ViT-Tiny
Test Acc (t)
48.09
56.93
65.06
69.71
SPAM
VIT-Tiny
Test Acc (f)
47.34
56.47
65.57
69.98
17

---

Published as a conference paper at ICLR 2025
F MORE ABLATION STUDY OF SUBSET SELECTIOIN STRATEGIES
Key questions surrounding sparse momentum include how to effectively select parameter subset and whether to retain momentum for weights that are sampled multiple times. To answer this questions, conduct comparative studies based on LLaMA-60M and C4 and the results shown in Figure Figure Left shows the performence of three subset selection strategies where we will reset all moments after each momentum reset and keep gradients for all unselected parameters. Figure 7 Middle shows the performence of three subset selection strategies where we will keep the overlapped moments after each momentum reset and keep gradients for all unselected parameters. Figure 7 Right shows the performence of three subset selection strategies where we will reset all the moments after each momentum reset and drop gradients for all unselected parameters in each updating step. We observe the following: Among the three subset selection strategies- Max weight magnitude- based, Max gradient magnitude-based, and Random selection- the Random selection consistently outperforms the other two approaches. @ Comparing Figure Left and Figure Right, that resetting all moments after each momentum reset yields better performance than preserving overlapping moments.
Zero Moments Keep Gradient
Keep Moments Keep Gradient
Zero Moments Zero Gradient
45
Perplexity
40
40
40
35
35
35
30
30
0.25
0.50 Density
0.75
1.00
0.25
0.50 Density (d)
0.75
1.00
0.25
0.50
0.75
1.00
Density
Weight-Max
Random
Gradient-Max
Figure 7: Ablations for subset selection strategies.
The experiments are conducted with LLaMA-
60M on C4.
G EXPERIMENTS ON TIME SERIES DATA
To showcase SPAM's ability to mitigate gradient spikes across broader range applications, we conducted additional experiments time-series prediction tasks. In these experiments, we intentionally introduced anomalous data with 10% probability to simulate gradient anomalies. Experiments conducted with 10 repeated runs on Weather time series date using PatchTST (Nie et al. 2023 model. The results are presented in Figure
The findings demonstrate that the severity of anomalous data increases, SPAM's performance advantage over Adam becomes more pronounced, highlighting its effectiveness in mitigating the adverse impact of gradient spikes.
H PROLONGED DETRIMENTAL EFFECTS OF GRADIENT SPIKES DURING REAL TRAINING
We also measure the values of gradient, first moment, and second moment during the training of LLaMA-60M on the C4 dataset. The results are now presented Figure 9
From the figure, we observe that during actual training, gradient spikes also have a significant and prolonged detrimental impact on moments, especially on the second moment, providing further evidence to support our claims.
wttpp::/// htt////www..mm bgc- jena mpg. de /wt/ttett////////
18

---

Published as a conference paper at ICLR 2025
No Anomaly Data
10% Anomelies (Severity:2)
10% Anomalies (Severity:5)
0.90
10%
Anomalios (Severity-10)
Test Loss
0.155
0.22
0.38
9.85
0.131
0.152
0.20
0.34
0.80
0.19
0.32
0.73
0.115
Training Epochs
Training Epochs
Training Epochs
Troining Epochs
Vanilla Adam
SPAM
Figure 8: Test Loss during Training Process on Weather Time-series Data. Anomalous data is generated by adding Gaussian noise to 10% of randomly selected input values. Specifically, the anomalies data are conducted with X inputs.
Max(X)) where X the
0.0004
1.0
0.0002
Second Moment
First Moment
Gradient
0.0000
0
0.6
-0,0002
0.4
0.0004
0.2
-0.0006
-0.0008
0.0
500
1000
0
500
1000
500
1000
Update Steps
Update Steps
Update Steps
Figure 9: Gradient spikes have prolonged detrimental effects the first and second moments. Experiments are conducted on C4 dataset with LLaMA-60M.
I SENSITIVITY ANALYSIS OF HYPERPARAMETER A ON LLM ARCHITECTURES
We conducted experiments to evaluate the sensitivity of the gradient spike clipping threshold, 0, across three widely used LLM architectures: LLaMA, Pythia, and OPT. These experiments were performed on pre-training tasks using the C4 dataset. The final perplexity is reported in Table [T]
The results indicate that the gradient spike clipping threshold is not highly sensitive the choice of LLM architecture. SPAM consistently outperforms Adam across a wide range of f. Furthermore, the optimal range for é lies between 1000 and 5000.
Table 11: Sensitivity Analysis of Hyperparameter on L.M architectures. Perplexity is reported.
Architectures
500
0 = 1000
2500
5000
10000
Adam
LLaMA-60M
30.77
30.59
30.57
30.46
30.82
34.09
PYM
34.4
34.1
34.1
34.2
35.1
38.34
OPT-125M
28.7
28.4
28.5
28.6
29.0
32.20
J GSS VS. DISTRIBUTION BASED CLIPPING
We conducted experiment using outlier detection mechanism based on the assumption that stochastic gradient distributions follow a Gaussian distribution, as suggested in (Simsekli et al.
2019 Chaudhari & Soatto 2018 Mandt al. 2016):
Ghatch 8
where Gbotch is the stochastic gradient, G represents the gradient over the entire dataset, and 62 is the variance. Since calculating G on-the-fly fly during training computationally infeasible, approximate it using the moving average of Ghatch- The variance 82 is estimated online
19

---

Published as a conference paper at ICLR 2025
82 En (hatek ((n) where N is the total training steps. Gradients then evaluated
element-wise, and any element ((n) batch satisfying: |G G( > 36 is identified as an outlier.
Such outlier elements are clipped to satisfy: |G G(n) = 35.
We conducted experiments using LLaMA-60M and LLaMA- 130M to evaluate the performance of this Gaussian-basee Clipping and compare it with our proposed GSS-based clipping. The results are reported in Table T2 As the table indicates, Gaussian-based clipping falls short of our GSS-based clipping. One possible explanation is that stochastic gradient distributions are very complex and Gaussian distribution can not reflect the true distribution.
Table 12: Comparison between SPAM with spike-aware clipping and Gaussian-based clipping.
Methods
LLMMA-60M
LLaMA-130M
SPAM w/GSS based clipping
30.46
23,36
SPAM w/ Gaussian based Clipping
30.83
25.93
K GSS BASED CLIPPING VS. NULLIFYING
We conducted experiments on LLaMA-60M and LLaMA-130M to compare the performance of Spike-Aware Clipping and Nullifying Gradient Spikes. As shown Table and Table [4 SPAM with Spike-Aware Clipping outperforms SPAM with Nullifying on both pre-training and fine-tuning tasks, demonstrating the effectiveness of Spike-Aware Clipping.
Table 13: Comparison between SPAM w/ spike-aware clipping and SPAM w/ nullifying gradient spikes.
Methods
LLaMA-1300 LLaMA-60M
SPAM w/ Spike Aware Clipping
30.46
23.36
SPAM w/ Nullifying
30.86
23.62
Table 14:
Comparison between SPAM w/ spike-aware clipping and SPAM w/ nullifying gradient spikes on fine-tume tinettnnig task. The experiments are based on a pre-trained OPT-1.3B model.
Methods
WinoGrande
COPA
SPAM w/ Spike Aware Clipping (d=100%)
59.4
79.0
SPAM w/ Spike Aware Clipping (d=0.25%)
58.3
75.0
SPAM w/ Nollifying(d=100%)
58.0
78.0
SPAM w/ Nullifying (d=0.25%)
57.4
75.0
L COMPUTATIONAL ANALYSIS
We measured the running time per iteration for both LLaMA-60M and LLaMA-130M The results, presented in Table 15 indicate that SPAM incurs a slightly higher computational overhead compared to Adam, Adam-min and Adafactor. This overhead is primarily due to the gradient spike detection operation and the gradient selection based on sparse masks. However, we believe that such a small overhead is negligible compared to the overall pre-training time which can be dozens or hundreds of hours.
20

---

Published as a conference paper at ICLR 2025
Table 15: Running Time per Iteration (second). The runtime is measured by the average of 100 iterations under one H100 GPU.
Method
Time per Iteration (LLaMA-60M)
Time per Iteration (LLaMA-130M)
Adam
0.3666 (s)
0.6397
Adam-mini
0,3614 (s)
0,6472 (s)
Adafactor
0.3778 (s)
0.6565 (s)
GaLore (rank=128)
0.3871 (s)
0,6702 (s)
SPAM((==00%%)
0,3814 (8)
0.6683 (s)
(PAM(d=25%)
0.3799 (s)
0.6658 (s)
21


Reviews:
[{"review_id": "W0SPytaXZe", "reviewer": "Reviewer_BUJ8", "summary": "The paper proposes a novel optimization method designed to mitigate training instabilities, specifically gradient and loss spikes in large language models (LLMs). The method, SPAM, introduces momentum reset and spike-aware gradient clipping to counteract the effects of significant gradient spikes that can disrupt the learning process. Extensive experiments suggest that SPAM outperforms traditional Adam and its memory-efficient variants, offering better stability and performance across different model scales and training setups.", "questions": "1. **Detrimental effects of gradient spikes**. It would be valuable to observe the middle and right plots of Figure 5 during actual training.\n2. **Moment reset.** Does the benefit of momentum reset lie in isolating the effects of gradient spikes, even at the cost of training intervals affected by these spikes?\n3. **Statistical significance**. Were the fine-tuning experiments conducted using multiple random seeds?\n4. **Baselines**. How does this method compare with other memory-efficient approaches such as MeZO [1], SparseMeZO [2], and Extremely Sparse MeZO [3]?\n5. **Ablation studies**.\n    1. There is no comparison between Spike-Aware Clipping and simply nullifying gradient spikes.\n    2. It would be interesting to see if the parameter for sparse momentum could be selected in a structured manner, as described in  [4].\n    3. What happens in the case $\\triangle T < N$?\n6. **Loss and Gradient plots of SPAM**. How do the loss and gradient plots of SPAM compare with those of other methods shown in Figures 2-4?\n7. **Computational analysis**. Can the authors report the computational overhead of SPAM compared to other methods?\n\n[1] Malladi, Sadhika, et al. \"Fine-tuning language models with just forward passes.\" *Advances in Neural Information Processing Systems* 36 (2023): 53038-53075.\n\n[2] Liu, Yong, et al. \"Sparse mezo: Less parameters for better performance in zeroth-order llm fine-tuning.\" *arXiv preprint arXiv:2402.15751* (2024).\n\n[3] Guo, Wentao, et al. \"Zeroth-Order Fine-Tuning of LLMs with Extreme Sparsity.\" *arXiv preprint arXiv:2406.02913* (2024).\n\n[4] He, Yang, and Lingao Xiao. \"Structured pruning for deep convolutional neural networks: A survey.\" IEEE transactions on pattern analysis and machine intelligence (2023).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The paper’s analysis highlighting the prevalence and impact of gradient spikes in LLM training is insightful and demonstrates an important issue.\n- The experiments show consistent improvements across different LLM sizes and benchmarks, suggesting the method’s robustness within these settings.\n- The introduction of sparse momentum is a useful addition for reducing the memory overhead of training large models.", "weaknesses": "- Although SPAM is compared to Adam and a few memory-efficient optimizers, it lacks comprehensive analysis against more recent memory-efficient methods. Furthermore, additional experiments, as outlined below, are necessary to strengthen the evaluation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "l23eReZMOg", "reviewer": "Reviewer_LnTM", "summary": "The paper introduces SPAM (Spike-Aware Adam with Momentum Reset), an innovative optimization approach aimed at stabilizing the training of LLMs. SPAM addresses the challenges of gradient and loss spikes that result in training instability, which can require costly interventions such as checkpoint recovery. The proposed method improves the ADAM optimizer via integrating two key mechanisms:  spike-aware gradient clipping and momentum reset. These novelties contribute to mitigating the accumulation of gradient spikes, therefore enhancing the training stability and efficiency. Extensive experiments show that SPAM outperforms ADAM and other memory-efficient optimizers like GaLore and Adam-Mini over various LLM sizes and tasks, supporting its potential to improve training under memory constraints.", "questions": "Could the spike-aware clipping method proposed be extended to handle other types of optimization tasks and scenarios, like adversarial training?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The integration of momentum reset and spike-aware gradient clipping into Adam is noval and addresses the persistent issue of gradient spikes in Large Language Model training.\n2. The experiments are thorough and extensive, with evaluations spanning multiple LLM architectures and scales. The results clearly manifest SPAM's superior performance over the standard and memory-efficient baselines.\n3. The approach is highly relevant, especially for large-scale training where stability and efficiency are paramount.\n4. SPAM's sparse momentum feature is especially useful for resource-constrained training, making it an important contribution to memory-efficient optimization approaches.", "weaknesses": "1. The paper mentions the efficient implementation of momentum reset and spike detection, but a moredetailed practical guidance or pseudo code might improve reproducibility.\n2. While SPAM performs excellently across various Large Language Model sizes, additional experiments on tasks beyond LLM training, such as CV models or multi-task learning, should illustrate broader applicability.\n3. The choice of the gradient spike threshold might affect performance to a great extent. More discussion on how to tune this parameter across different model architectures would be of benefits.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rMz8G4yU7Y", "reviewer": "Reviewer_BU76", "summary": "This paper analyzes the phenomenon of loss spikes observed in large language models (LLMs) and reveals that loss spikes are not restricted to specific layers or architectures, but occur in a wide range of environments. The authors experimentally demonstrate that loss spikes affect the performance of AI systems and mathematically show that loss spikes are influenced by momentum-based optimizers, such as Adam. The proposed method, SPAM (Stochastic Gradient Projection with Adaptive Momentum), effectively addresses this issue by using a threshold-based approach to manage the average gradient. The paper compares the performance of SPAM in both pre-training and fine-tuning stages with various other methods, showing its effectiveness.", "questions": "1. The use of theta in GSS seems heuristic. How about selecting outliers based on the known distribution (such as Gaussian) of gradients instead? [Umut Simsekli, Levent Sagun, & Mert Gurbuzbalaban. (2019). A Tail-Index Analysis of Stochastic Gradient Noise in Deep Neural Networks.]\n2. If the method used in the experiments is SPAM with the sparse momentum approach, according to Algorithm 1, when m and v are set to zero, some weights are not updated, similar to dropout. It’s unclear whether the performance improvement is due to this or the actual spike gradient clipping.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper addresses the loss spike problem from the perspective of gradient clipping and demonstrates the algorithm's validity through an ablation study on the hyper-parameters used in the algorithm, along with various performance improvements. Additionally, the paper proposes a memory-efficient algorithm using sparse momentum, aiming to solve both the loss spike issue and the out-of-memory problem simultaneously.", "weaknesses": "1. Clipping gradients based on a threshold seems to lack novelty. It might be worthwhile to consider methods that prevent gradient spikes altogether.\n2. In sparse momentum, a random mask is applied, setting certain gradients to zero. It would be helpful to explain in detail how this actually reduces memory usage. From an algorithmic perspective, it appears as though the entire matrix, including the zero elements, is still being stored.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)