SoftSignSGD(S3): An Enhanced Optimize for Practical DNN Training and Loss Spikes Minimization Beyond Adam

Q2. The bound in Theorem 4 depends on dimension $d$. Compare the dependence on $d$ explicitly with other recent works on adaptive optimizers like Adam.

R2. We summarize the convergence rates of S3 and Adam (from its original paper and subsequent recent works) as follows:

1. The convergence rate of S3 (Our Work): From Theorem 4, $\frac{1}{T}\sum_{t=1}^T \mathbb{E}[\Vert\nabla F({x}_t)\Vert_1] \leq \mathcal{O}\left(\frac{d}{T^{1/4}}\right).$

2. ** The convergence rate of Adam from (Kingma and Ba, 2015) [1]**: Under convexity assumptions, $\frac{1}{T}\sum_{t=1}^T \left(f({x}_t) - f({x}^*)\right) \leq \mathcal{O}\left(\frac{d}{T^{1/2}}\right).$

3. The convergence rate of Adam in from (D{'e}fossez et al., 2020) [2]: $\frac{1}{T} \sum _{t=1}^T \mathbb{E}[\Vert\nabla F({x}_t)\Vert_2] \leq \mathcal{O}\left(\frac{d\ln(T)}{T^{1/4}}\right).$

   Since $\Vert a \Vert_1 \leq \sqrt{d}\Vert a \Vert_2$, the rate can be expressed as: $\frac{1}{T}\sum _{t=1}^T \mathbb{E}[\Vert\nabla F({x}_t)\Vert_1] \leq \mathcal{O}\left(\frac{d\sqrt{d}\ln(T)}{T^{1/4}}\right).$

4. The convergence rate of Adam from (Chen et al., 2022) [3]: $\frac{1}{T}\sum _{t=1}^T \mathbb{E}[\Vert\nabla F({x}_t)\Vert_2] \leq \mathcal{O}\left(\frac{G\sqrt{d}\ln(T)}{T^{1/4}}\right),$ where $G \ge \Vert {g}_t \Vert_2 \propto \sqrt{d}$. Thus, $\frac{1}{T}\sum _{t=1}^T \mathbb{E}[\Vert\nabla F({x}_t)\Vert_1] \leq \mathcal{O}\left(\frac{d\sqrt{d}\ln(T)}{T^{1/4}}\right).$

5. Adam (from Li et al., 2023) [4] In probabilistic bounds (Li el at. 2023) , $\frac{1}{T}\sum _{t=1}^T \Vert\nabla F({x}_t)\Vert_2 \leq \mathcal{O}\left(\frac{G^2\ln{\frac{1}{\delta}}}{T^{1/4}}\right), \text{ with probability at least } 1-\delta,$ where $G \ge \Vert {g}_t \Vert_2 \propto \sqrt{d}$. Hence, $\frac{1}{T}\sum _{t=1}^T \mathbb{E}[\Vert\nabla F({x}_t)\Vert_1] \leq \mathcal{O}\left(\frac{d\sqrt{d}\ln{\frac{1}{\delta}}}{T^{1/4}}\right), \text{ with probability at least } 1-\delta.$

From this comparison, S3 may achieve a weaker dependency on $d$ compared to Adam. We hope this clarification highlights the theoretical advantages of S3's convergence.

Q3. I also suggest adding the standard deviation for all reported metrics.

R3. Thank you for the valuable suggestion. We acknowledge that we only ran one trial per result due to computational constraints. Running multiple trials and reporting mean and standard deviation values would indeed make the results statistically more robust. However, like many in academia with limited computational resources, we have access to only two machines, each equipped with eight V100 GPUs. Each trial of training ViT-B/16 and ResNet-50 on ImageNet required approximately 20 hours, while training GPT-2 (345M) on OpenWebText took about 70 hours per trial. We spent over three months completing all experiments, running one trial per result, except for GPT-2 (7B), which required approximately 20 days on a rented cluster. Consequently, while we would have liked to conduct more experiments, our limited computational resources make this impractical.

We emphasize that the datasets used, such as ImageNet and OpenWebText, are large-scale and diverse. This, coupled with results obtained from a single trial, provides reasonable confidence in their robustness. Moreover, the related Adan and Lion from industrial community also reported results based on single trials, indicating that this practice is acceptable for large-scale experiments. Notably, as for the most related Adan and Lion that come from the industrial community, they also only ran one trial for each results. In fact, the size of the datasets ImageNet and OpenWebText in our experiments is much huge, so the even one trial for each result will be still robust and convincing.

Q4. It is better to provide loss curves for all setups in Table 3.

R4. Thank you for the suggestion. We initially refrained from plotting all curves in the tables to avoid overlapping and cluttered legends. As suggested, we have included all curves in the revised manuscript. For convenience, the figure is also available at the following link https://1drv.ms/b/s!Au3MrR-o69M4iMgyf5FGMpdbJEd2cA?e=akqLug.

We sincerely thank you for your voluntary efforts, constructive comments, and positive score. We are especially grateful for the note, “The authors propose an optimizer, S3, with theoretical convergence guarantee, primarily aimed at solving the training instability issues seen with Adam.” Furthermore, we are pleased that the solid and extensive experiments have been acknowledged. Below, we address your remaining concerns point-by-point.

Q1. Clarifying the contributions of the modifications over Adam ("NAG," "higher $p$," and "same $\beta$") to loss spike prevention.

We greatly appreciate your keen observations and thoughtful analysis of our experimental results. Based on Theorems 1 and 2, the analyses in Section 2, and our experimental findings, we provide the following clarifications:

1. NAG's impact: NAG has no direct connection to the reduction of loss spikes. According to Theorems 1 and 2, the maximum update magnitude remains the same regardless of whether NAG is applied, meaning it has minimal influence on loss spikes.

   Regarding the setting "lr=3e-3, diff. $\beta$, w/NAG, $p=2$" in Figure 5, where no spikes are observed, the reasons are as follows:

   • Model Size: The number of parameters in ViT-B is 86M, which is relatively small. The probability of loss spikes is positively correlated with model size, so spikes may not emerge with this random seed due to their low probability.
   • Recording Granularity: During ViT-B training, we recorded training loss per epoch, whereas for GPT-2, loss was recorded per iteration. Consequently, some loss spikes might have occurred between epochs but were not captured.

   To validate this, we retrained ViT-B/16 with the same configuration ($lr=3e-3$, diff. $\beta$, w/NAG, $p=2$) using a different random seed. The resulting curves, available at this link, show a loss spike, corroborating our explanation.

2. Effect of "higher $p$": From Theorem 2, we know that: $\frac{\vert {n} _t^{(j)} \vert}{{b} _t^{(j)}(p_1)}({g} _t^{(j)}, {g} _{t-1}^{(j)}, \dots) \geq \frac{\vert {n} _t^{(j)} \vert}{{b} _t^{(j)}(p_2)}({g} _t^{(j)}, {g} _{t-1}^{(j)}, \dots) \quad \text{if } p_1 \leq p_2.$ This implies that higher $p$ reduces the maximum possible value of $\frac{\vert {n}_t^{(j)} \vert}{{b}_t^{(j)}(p)}$, thereby lowering the risk of loss spikes indirectly.

3. Effect of "same $\beta$": As demonstrated in Theorem 2 and Section 3, "same $\beta$" directly minimizes the risk of loss spikes by controlling the update magnitude effectively.

In summary, while NAG does not contribute to loss spike prevention, "higher $p$" indirectly reduces their probability by lowering the maximum upper bound, and "same $\beta$" directly minimizes the risk. These analyses will be further clarified in the revised manuscript to improve the reader's understanding.

I sincerely thank the authors for their detailed responses to my comments and for conducting additional experiments. I greatly appreciate the effort and time invested. However, I still have some concerns regarding Q1, particularly in light of the figure provided in response to Q4.

To effectively demonstrate that S3 can mitigate loss spikes, the experimental setup needs to be sufficiently robust. Unfortunately, I find the explanations provided by the authors not entirely convincing. Specifically:

• The authors stated, "The probability of loss spikes is positively correlated with model size, so spikes may not emerge with this random seed due to their low probability." However, I would like to point out that loss spikes occur in all Adam runs, which makes the experimental results appear to be cherry-picked.
• While I understand the limitations posed by resource constraints, if the randomness of the seed can have such a significant impact on experimental behavior, I strongly suggest conducting multiple trials to strengthen the conclusions.
• The authors stated, "Consequently, some loss spikes might have occurred between epochs but were not captured." I believe this issue can be addressed through certain loss-averaging techniques.
• To convincingly demonstrate that S3 can prevent loss spikes in Adam due to consistent \beta, it would be most effective to operate in a setting where all runs with non-identical \beta exhibit loss spikes, while all runs with identical \beta do not.

Finally, regarding the diffusion model experiments: in the case of large-scale diffusion model pretraining, would loss spikes or training instability be a problem for Adam?

Once again, I appreciate the authors' efforts and look forward to further clarifications.

Thank you for your thoughtful feedback and for providing us with the opportunity to further clarify our work. We sincerely appreciate your detailed comments and address your concerns as follows:

Regarding Loss Spikes

1. Experimental Setup and Seed Selection. In our experiments, including those conducted prior to your comments, we used a fixed random seed ("6666") for training ResNet-50 and ViT-B/16 on ImageNet. This is documented in the source code provided in the supplemental materials. The results where S3 (lr=3e-3, diff. $\beta$, w/NAG, $p=2$) did not exhibit a loss spike while Adam did were not selectively chosen. These results emerged naturally under this fixed seed.

   The experiments presented in Section 5.3 and Figure 5 aimed to conduct ablation studies to evaluate the contributions of individual components to the overall performance. They were not intended to explore the underlying causes of loss spikes, which were specifically addressed in Section 5.5 and Figure 7. We clearly knew that training ViT-B/16 with Adam does not consistently lead to frequent loss spikes due to the model's size, making it less ideal for studying this phenomenon of loss spikes. For this reason, we did not mention loss spikes in Section 5.3, nor did we record per-iteration training loss to capture the detailed occurrences of loss spikes.

2. Impact of Loss Spikes on Final Performance. Based on our observations, infrequent loss spikes during training have minimal impact on final performance. For example, as shown in Section 5.3, the test accuracy of AdamW, which encountered a loss spike, remains comparable to S3 (lr=3e-3, same $\beta$, w/o NAG, $p=2$) that did not exhibit spikes. Similarly, the results in Figure 3(c-d) and Table 1 show that the accuracy of AdamW after 300 epochs aligns with the reference values in TorchVision (https://pytorch.org/vision/main/models.html and https://github.com/pytorch/vision/tree/main/references/classification), which we cloned from GitHub as the basic source codes for the vision tasks, as described in Section B.1.

   Given this robustness,the conclusions drawn from the presented results may remain valid without requiring multiple trials, especially considering the resource constraints that limit additional experiments at this time.

3. Detailed Analysis of Loss Spikes. To verify the root causes of loss spikes, we fortunately had conducted detailed experiments in Section 5.5 using GPT-2 (345M), including the scenarios you suggested (i.e., testing Adam with both identical and non-identical $\beta_s$). The results demonstrate that:

   • When $\beta_1$ and $\beta_2$ differ, loss spikes frequently occur in the early stages of training.
   • When $\beta_1$ and $\beta_2$ are identical under otherwise identical conditions, the loss spikes are eliminated.
   • Constraining $\max \left(\frac{|m_t^{(j)}|}{\sqrt{v_t^{(j)}}}\right) \leq 1$ while keeping $\beta_1 \neq \beta_2$ significantly reduces loss spikes.

   These findings confirm that Adam with non-identical $\beta$ values can produce overlarge gradients, causing loss spikes, and setting $\beta_1 = \beta_2$ effectively mitigates this issue, as deeply analyzed in Sections 2 and 3.

Regarding Diffusion Model Experiments

While we have not conducted experiments on large-scale diffusion model pretraining with S3 or Adam using identical $\beta$, we believe that the proposed approach can mitigate training instabilities arising from loss spikes. Our algorithm is designed to be generalizable and is not limited to specific tasks or architectures.

We hope these clarifications address your concerns comprehensively. Thank you once again for your valuable feedback, which has been instrumental in refining our work.

Thank you for your response. Based on your reply, I will keep my score at 6 (marginal accept). However, if you can show experiments where S3, i.e. Adam with the same \beta, helps with convergence in a reasonable setting where loss spikes affect performance, I’m happy to raise my score to 8.

Additionally, about my question on diffusion: I was asking if loss spikes are actually a problem when training diffusion models in practice since they also use the Adam optimizer.

We are highly grateful for your prompt feedback. We would respectfully ask if you could specify which parts of the experiments in Figures 2 and 7 you found unconvincing. We believe that these experiments provide strong evidence supporting our detailed analysis in Section 2 on how Adam can potentially cause loss spikes.

Clarification about the Experimental Results

We would respectfully ask if you could specify which parts of the experiments in Figures 2 and 7 you found unconvincing. We believe that these experiments provide strong evidence supporting our detailed analysis in Section 2 on how Adam can potentially cause loss spikes.

Before our work, the root cause of loss spikes was not well understood, leaving researchers and practitioners unable to reliably reproduce or prevent loss spikes during LLM training. Although [2] attempted to propose a theory for the phenomenon, they were unable to identify the fundamental cause. Their conclusion explicitly acknowledges this, stating: "We conclude that at this point, there is no silver bullet to solve the problem, and the appropriate remedy depends on the specific setup of the large-scale training run."

In contrast, our work identifies the root cause of loss spikes and provides a simple method to address them. While the statement in [2], "Loss spikes are difficult to study because any reproduction of these spikes at a smaller scale is not necessarily caused by or remediated by the same factors as in larger scales," may hold for their analysis, it does not apply to our findings.

As noted by Reviewer Xi7Z, "Loss spikes in LLM training is a long-standing problem, and this possible explanation and solution should be immediately useful for the academic community." We believe that uncovering the root cause of loss spikes and providing a straightforward solution is indeed a significant contribution to the field.

Consistency with Observations in Related Work

Our analysis in Section 2 aligns well with the descriptions of loss spikes in [1][2][3][4][5]. For your convenience, we provide the following excerpt from our paper:

"Theorem 1 indicates that when \textsf{\small Adam} is employed, there exists a probability that the update of each element $\frac{\vert{m} _t^{(j)}\vert}{\sqrt{{v} _t^{(j)}}}$ can reach an excessively large value. For instance, with typical settings of $\beta_1=0.9$ and $\beta_2=0.999$, the update $\frac{\vert{m}_t^{(j)}\vert}{\sqrt{{v}_t^{(j)}}}$ could approach its theoretical maximum of $\frac{1-\beta_1}{\sqrt{1-\beta_2}\sqrt{1-\frac{\beta_1^2}{\beta_2}}}\simeq 7.27$, while the normal absolute value of the update is less than $1$. While the probability of any specific parameter's update reaching this maximal value is low, the probability that at least one parameter's update reaches this value is high due to the large number of parameters in large models (LLM).When a parameter's update is excessively large and the learning rate is also high, it is likely to deviate substantially from its intended trajectory. Such deviations may propagate to neighboring parameters through interconnections, triggering a chain reaction that culminates in loss spikes. This mechanism explains the frequent occurrence of loss spikes during LLM training, particularly in the initial stages when learning rates are higher. Specifically, the probability of loss spikes increases with the size of the LLM. In conclusion, vanilla Adam poses a significant risk of loss spikes during large-scale model training. Mitigating this problem requires strategies to constrain the maximum update magnitude for each parameter coordinate."

From this analysis, we can conclude:

• Larger models experience more frequent loss spikes, a trend that is consistently recorded in [1][2][3][4]. This may explain why reports of training instability were rare before the era of LLMs.
• Loss spikes are more likely to occur during the early stages of training, as also observed in [2][3][4][5].
• The common shift in Adam settings for LLM training from the default $\beta_2 = 0.999$ to $\beta_2 = 0.95$ that is more close to $\beta_1=0.9$ can be explained by our analysis, as it helps mitigate the occurrence of loss spikes.

Additionally, [2] notes that spikes requiring the batch-skipping trick may not appear in the figures of [2][4][5], which could affect the observations of such phenomena.

We hope these clarifications address your concerns. Thank you once again for your enthusiastical continued feedbacks.

We sincerely appreciate your further comments and this opportunity to clarify them. Below, we provide detailed responses to your concerns.

Loss Spikes Affecting Convergence

Fortunately, Figure 7 in the manuscript visually illustrates that the frequent loss spikes caused by vanilla AdamW lead to a slight slowdown in convergence, while AdamW with the same $\beta$s helps restore the convergence rate. For your convenience, we provide Figure 7 at this link: https://1drv.ms/b/s!Au3MrR-o69M4iMVueg52QALbpzxStQ?e=b8lCTO .

In this figure, we used AdamW under different conditions to train GPT-2 (345M) on OpenWebText. The experimental results show the following:

• With vanilla AdamW and a learning rate of $1 \times 10^{-3}$, more frequent spikes occur compared to the baseline learning rate of $3 \times 10^{-4}$. This causes a slight slowdown in convergence and results in a higher final loss.
• At a learning rate of $3 \times 10^{-3}$, even more frequent spikes emerge, eventually causing divergence.
• When adopting AdamW with $\beta_1 = \beta_2$ and a learning rate of $1 \times 10^{-3}$, as predicted by our analysis in Section 3, the loss spikes disappear. This configuration achieves faster convergence and a lower final loss compared to vanilla AdamW with the baseline learning rate of $3 \times 10^{-4}$.

In summary, when training DNNs encounters a loss spike, the loss often recovers in a short time. However, there is a small probability that loss will deviate from the intended trajectory due to the abrupt parameter changes, as detailed in Section 2. When frequent loss spikes occur, the trajectory may derail intended trajectory completely with high probability, negatively affecting both the convergence rate and the final loss.

This phenomenon is widely recognized by practitioners in the field of large language models (LLMs). To avoid the slowdown convergence and degraded performance from frequent loss spikes, practitioners would rather to resort to resource-wasting and time-wasting engineering strategies, i.e., skipping data batches before spikes occur or restarting training from nearby checkpoints [1][2]. Furthermore, some famous LLM groups have invested significant time and resources to study and mitigate this issue [3][4][5].

References:

[1] Aakanksha Chowdhery, et al., "Palm: Scaling language modeling with pathways," JMLR, 2023.

[2] Igor Molybog, et al., "A theory on Adam instability in large-scale machine learning," arXiv:2304.09871, 2023.

[3] Aohan Zeng, Xiao Liu, et al., "GLM-130B: An open bilingual pre-trained model," arXiv:2210.02414, 2022.

[4] Hugo Touvron, et al., "Llama: Open and efficient foundation language models," arXiv:2302.13971, 2023.

[5] Aiyuan Yang, et al., "Baichuan 2: Open large-scale language models," arXiv:2309.10305, 2023.

Loss Spikes in Training Diffusion Models

We have no direct experience in training large-scale diffusion models. However, our analysis of loss spikes caused by AdamW is generalizable. Therefore, if frequent loss spikes occur during the training of diffusion models, we think they could negatively affect convergence and final performance.

We hope these clarifications address your concerns comprehensively. Sincerely thank you once again for your valuable feedback, which has been instrumental in refining our work.

We sincerely thank you for your voluntary efforts and constructive comments. We appreciate your acknowledgment of the significance of minimizing the risk of loss spikes via soft-sign descent and the solid experimental evidence supporting our approach. Below, we address your remaining concerns point by point.

Q1. The assumption under which Theorem 1 holds (the signs of gradients are the same.) seems quite strong and may not hold in general.

R1. Thank you for pointing this out. We agree that the assumption is stringent and may not always hold strictly. However:

1. The update $\frac{\vert{m} _t^{(j)}\vert}{\sqrt{{v}_t^{(j)}}}$ with respect to $({g} _k^{(j)}) _{k=1}^t$ is a continuous function. Thus, when most of the signs of $({g} _k^{(j)}) _{k=1}^t$ are consistent and the secondary condition is nearly satisfied, $\frac{\vert{m} _t^{(j)}\vert}{\sqrt{{v} _t^{(j)}}}$ will be close to the theoretical maximum.
2. Historical gradients (${g} _k^{(j)}$) from much earlier iterations have minimal influence on the current value of $\frac{{m} _t^{(j)}}{\sqrt{{v} _t^{(j)}}}$ due to the moving exponential average.

As a result, it becomes easier for $\frac{\vert{m} _t^{(j)}\vert}{\sqrt{{v}_t^{(j)}}}$ to reach a relatively high value, probabilistically leading to loss spikes. This is supported by Figure 2, which shows that the peaks of $\max _{j \in [d]} \left( \frac{\vert {m} _t^{(j)} \vert}{\sqrt{{v} _t^{(j)}}} \right)$ during a loss spike are close to, but not strictly equivalent to, the theoretical maximum of 7.27.

In summary, while the assumption in Theorem 1 may not hold strictly, it provides a strong theoretical basis for explaining loss spikes when using Adam. We will incorporate this discussion into the revised manuscript. Your feedback has greatly improved the clarity of the paper.

Q2. Theorem 4 still has a dependency on the problem dimensionality $d$.

R2. You are correct that the convergence bound of S3 in Theorem 4 depends on the dimensionality $d$. However, its dependency is weaker than Adam's, as demonstrated in recent works. Below, we summarize the convergence rates for S3 and Adam for comparison:

1. The convergence rate of S3 (Our Work): From Theorem 4, $\frac{1}{T}\sum_{t=1}^T \mathbb{E}[\Vert\nabla F({x}_t)\Vert_1] \leq \mathcal{O}\left(\frac{d}{T^{1/4}}\right).$

2. ** The convergence rate of Adam from (Kingma and Ba, 2015) [1]**: Under convexity assumptions, $\frac{1}{T}\sum_{t=1}^T \left(f({x}_t) - f({x}^*)\right) \leq \mathcal{O}\left(\frac{d}{T^{1/2}}\right).$

3. The convergence rate of Adam in from (D{'e}fossez et al., 2020) [2]: $\frac{1}{T} \sum _{t=1}^T \mathbb{E}[\Vert\nabla F({x}_t)\Vert_2] \leq \mathcal{O}\left(\frac{d\ln(T)}{T^{1/4}}\right).$

   Since $\Vert a \Vert_1 \leq \sqrt{d}\Vert a \Vert_2$, the rate can be expressed as: $\frac{1}{T}\sum _{t=1}^T \mathbb{E}[\Vert\nabla F({x}_t)\Vert_1] \leq \mathcal{O}\left(\frac{d\sqrt{d}\ln(T)}{T^{1/4}}\right).$

4. The convergence rate of Adam from (Chen et al., 2022) [3]: $\frac{1}{T}\sum _{t=1}^T \mathbb{E}[\Vert\nabla F({x}_t)\Vert_2] \leq \mathcal{O}\left(\frac{G\sqrt{d}\ln(T)}{T^{1/4}}\right),$ where $G \ge \Vert {g}_t \Vert_2 \propto \sqrt{d}$. Thus, $\frac{1}{T}\sum _{t=1}^T \mathbb{E}[\Vert\nabla F({x}_t)\Vert_1] \leq \mathcal{O}\left(\frac{d\sqrt{d}\ln(T)}{T^{1/4}}\right).$

5. Adam (from Li et al., 2023) [4] In probabilistic bounds (Li el at. 2023) , $\frac{1}{T}\sum _{t=1}^T \Vert\nabla F({x}_t)\Vert_2 \leq \mathcal{O}\left(\frac{G^2\ln{\frac{1}{\delta}}}{T^{1/4}}\right), \text{ with probability at least } 1-\delta,$ where $G \ge \Vert {g}_t \Vert_2 \propto \sqrt{d}$. Hence, $\frac{1}{T}\sum _{t=1}^T \mathbb{E}[\Vert\nabla F({x}_t)\Vert_1] \leq \mathcal{O}\left(\frac{d\sqrt{d}\ln{\frac{1}{\delta}}}{T^{1/4}}\right), \text{ with probability at least } 1-\delta.$

From this comparison, S3 may achieve a weaker dependency on $d$ compared to Adam. We hope this clarification highlights the theoretical advantages of S3's convergence.

[1] Diederik Kingma and Jimmy Ba. "Adam: A method for stochastic optimization". ICLR, 2015.

[2] Alexandre D{'e}fossez, L{'e}on Bottou, Francis Bach, and Nicolas Usunier. "A simple convergence proof of adam and adagrad". arXiv:2003.02395, 2020.

[3] Congliang Chen, Li Shen, Fangyu Zou, and Wei Liu. "Towards practical adam: Non-convexity, convergence theory, and mini-batch acceleration. JMLR, 2022.

[4] Haochuan Li, Alexander Rakhlin, and Ali Jadbabaie. "Convergence of adam under relaxed assumptions". NeurIPS, 2023.

We sincerely appreciate your valuable feedback and constructive suggestions, which have significantly improved the quality and presentation of our work. We kindly ask whether our response has addressed your concerns, and we would be happy to provide further clarifications or address any additional questions to ensure a clear understanding of our contributions. If you find our response satisfactory, we would be grateful if you would consider raising your rating.

We look forward to your reply.

We sincerely appreciate the time and effort you have dedicated to reviewing our work. As the deadline for the reviewer-author discussion is approaching, we kindly ask if you could confirm whether our response has addressed your original concerns. If you are satisfied with the response, we would be grateful if you would consider updating your rating.

Thank you once again for your valuable feedback and for your time.

We sincerely appreciate your voluntary efforts and positive feedback.We are particularly grateful for your observation: "The issue of loss spikes in LLMs is a very interesting question and warrants further investigation." Additionally, we are encouraged by the recognition of our theoretical analyses for S3 and the acknowledgment of our empirical experiments. We address your remaining concerns point-by-point below.

Q1. The connection with the paper "On the Variance of the Adaptive Learning Rate and Beyond" [1].

R1. Your valuable insights help us establish a strong connection between our work and [1]. Although we had carefully studied [1] earlier, we did not recognize this connection. [1] empirically and theoretically demonstrated that the variance of the update $\frac{1}{\sqrt{{v}_t^{(j)}}}$ for each coordinate is significantly larger in the early stages, often causing $\frac{\vert{m}_t^{(j)}\vert}{\sqrt{{v}_t^{(j)}}}$ to become disproportionately large. Consequently, according to our analyses in Section 2, the train loss will encounter frequent loss spikes if the learning rate is also large, finally leading to train instability. Similarly, Theorem 1 in our work suggests that $\frac{\vert{m}_t^{(j)}\vert}{\sqrt{{v}_t^{(j)}}}$ is likely to become disproportionately large in the early stages, as conditions such as consistent sign patterns in $({g} _t^{(j)}, {g} _{t-1}^{(j)}, ...)$ and geometric decay of gradients are more easily satisfied at this stage. Furthermore, the analyses in both [1] and our work explain why the warmup technique enhances training stability during the early stages.

Thank you for this insightful observation, which has deepened our understanding of early-stage training stability without warmup. We have incorporated this finding into the revised manuscript.

Q2. Could the authors consider manually setting a large $\frac{{m} _t^{(j)}}{\sqrt{{v} _t^{(j)}}}$ in AdamW in specific coordinates to intentionally trigger a loss spike?

R2. Thank you for this insightful suggestion. Following your suggestion, we conducted an experiment training a 6-layer ViT using AdamW on CIFAR-10. Due to the relatively small number of parameters in the 6-layer ViT, the probability of naturally occurring training loss spikes is low, as discussed in Section 2. We randomly selected parameters in the third layer every 1000 iterations and manually set their updates $\frac{{m} _t^{(j)}}{\sqrt{{v} _t^{(j)}}}$ to $20$. The training loss curves are available at the following link: [provide link]. The loss curves exhibit some spikes, validating your hypothesis. Notably, loss spikes do not occur every 1000 iterations because abrupt changes in individual parameters do not always impact the overall loss. In summary, this suggestion has significantly deepened our understanding of the mechanisms behind loss spikes.

Furthermore, in addition to Figure 2 in Section 2, we conducted additional experiments in Section 5.5 to further investigate the causes of loss spikes from Adam, providing additional evidence and insights.

Q3.* Clarifying the specific cases in which they applied or chose not to apply warmup.

R3. Due to space constraints in the main text, we moved the detailed experimental settings, including warmup configurations, to Section B.1 of the appendix, which may have been overlooked. n Section B.1, we specified that the same warmup technique was applied across all experiments for the optimizers. For vision tasks, we linearly increased the learning rate to its peak over the first 30 warmup epochs, followed by a cosine decay to zero in subsequent epochs. For language tasks, we linearly increased the learning rate to its peak over the first 5000 warmup steps, followed by a cosine decay to $0.1\times$ of the peak in subsequent steps.

We sincerely thank you for your voluntary efforts and high recognition. More importantly, we are gratified that our explanation of loss spikes originating from vanilla Adam is recognized and that our proposed S3, which minimizes the risk of loss spikes, is considered an immediate contribution to the academic community. We address your concerns point-by-point below. following.

Q1. If $p$ is not equal to 1, how will the wall clock time per iteration change? (compared to Adam and S3 $p=1$)

R1. As suggested, we have re-trained GPT-2(345M) with S3 (p=1), S3 (p=2), and AdamW for 1000 iterations to record the wall clock time. The results are summarized below.

Table 1. The mean time per iteration for training GPT-2(345M) with S3 ($p=1$), S3 ($p=3$) and AdamW

As shown in Table 1, S3 ($p=3$) costs slightly higher time per iteration than AdamW, while S3 ($p=1$) is slightly faster than AdamW.

Q2. How is the learning rate for different optimizers determined? Are they properly tuned for baselines?

R2. Highly thank you for taking attention to this important detail. Due to space constraints, we provided the detailed settings of the learning rates for the baseline optimizers in Section B.1 of the appendix. A brief summary is provided here for your convenience.

During training for ResNet-50 and ViT-B/16, learning rates for all optimizers are linearly increased to their peak over the first 30 epochs, followed by a cosine decay to 0 in subsequent epochs. The peak learning rates are as follows:

• For SGD, we set $lr_{\max}=0.3$ , the default value in the official PyTorch vision reference codes.

• For AdamW, we utilize $lr_{\max}=0.003$ , the default value in the official PyTorch vision reference codes and also widely adopted in related studies [1][2].

• For Adan, we employe $lr_{\max}=0.015$, following official recommendations[3].

• For Lion, we adopt $lr_{\max}=0.001$, following official recommendations [2].

For GPT-2 (345M), we linearly increase learning rates for all the optimizers to the peak in the initial $5k$ steps and then decrease to $0.1\times$ of the peak with a cosine decay in the subsequent steps. Peak learning rates for all optimizers were determined through coarse search while training GPT-2 (345M) (see Figure 8 in the appendix). For GPT-2 (7B), the learning rates for AdamW followed settings from the LLaMA paper [4].

[1] J. Zhuang, et al. "AdaBelief optimizer: Adapting stepsizes by the belief in observed gradients". NeurIPS, 2020.

[2] X. Chen, et al. "Symbolic discovery of optimization algorithms". arXiv:2302.06675, 2023.

[3] X. Xie, et al. "Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models". TPAMI, 2024.

[4] H. Touvron, et al. "Llama: Open and efficient foundation language models". arXiv:2302.13971, 2023.

Q3. Will tuning betas for Adam (and getting a smoother loss curve) close the gap between Adam and S3?

R3. Thank you for the thoughtful suggestion. As we currently have access to only one machine with 8 V100 GPUs, it is not feasible to perform an exhaustive grid search over $\beta_1$ and $\beta_2$ for AdamW within the short response period. We conducted experiments using AdamW, fixing $\beta_1=0.9$ and changing $\beta_2$ to train ViT-B/16 on ImageNet. The results are summarized below:

Table 2. Impact of $\beta_1$ and $\beta_2$ on the Accuracy of AdamW training ViT-B/16 on ImageNet

As shown in Table 2, tuning $\beta$s may not enable Adam to achieve performance comparable to S3.

We sincerely thank you for your voluntary efforts and positive feedback. We are pleased to hear that you found the experimental results supportive. More importantly, we are encouraged by your recognition of the theoretical significance of S3's convergence analysis. We address your remaining concerns point-by-point below.

Q1. Is the S3 optimizer sensitive to some parameters, such as the learning rate?

R1. Thank you for highlighting this important detail. While the learning rate plays a significant role in the performance of all optimizers, Figure 8 in the appendix shows that S3 is not more sensitive to the learning rate than other optimizers. We also conducted sensitivity analyses for $\beta$ and $p$ of S3 in Section 5.4. The results demonstrate that S3 is not highly sensitive to $\beta$ or $p$; even the worst-performing parameter combination outperforms baseline AdamW.

Q2. Are there any drawbacks of the S3 optimizer?

R2. The value of $p$ in S3 can be any value greater than or equal to 1. However, setting $p$ as a floating-point number rather than an integer increases computational cost, compared to the baseline Adam.

Q3. It would be better if more experimental results are provided.

R3. Thank you for this valuable feedback. We understand that additional experiments would make the results even more compelling. However, like many in academia with limited computational resources, we have access to only two machines, each equipped with eight V100 GPUs. Each trial of training ViT-B/16 and ResNet-50 on ImageNet required approximately 20 hours, while training GPT-2 (345M) on OpenWebText took about 70 hours per trial. We spent over three months completing all experiments, running one trial per result, except for GPT-2 (7B), which required approximately 20 days on a rented cluster. Consequently, while we would have liked to conduct more experiments, our limited computational resources make this impractical. Nevertheless, our experiments cover widely-used vision and language tasks, including popular CNN and Transformer architectures. These results, acknowledged as solid by Reviewer Xi7Z and Reviewer RDxB, provide extensive validation.