Taming Transformer Without Using Learning Rate Warmup

Authors: Xianbiao Qi, Yelin He, Jiaquan Ye, Chun-Guang Li, Bojia Zi, Xili Dai, Qin Zou, Rong Xiao
Primary area: optimization
Keywords: Transformer, Training Dynamics, Model Crash

Abstract:
Scaling Transformer to a large scale without using some technical tricks  such as learning rate warump and an obviously lower learning rate, is an extremely challenging task, and is increasingly gaining more attention. In this paper, we provide a theoretical analysis for the process of training Transformer and reveal a key problem behind model crash phenomenon in the training process, termed *spectral energy concentration* of ${W_q}^{\top} W_k$, which is the reason for a malignant entropy collapse, where ${W_q}$ and $W_k$ are the projection matrices for the query and the key in Transformer, respectively. 
To remedy this problem, motivated by *Weyl's Inequality*, we present a novel optimization strategy, \ie, making the weight updating in successive steps steady---if the ratio $\frac{\sigma_{1}(\nabla W_t)}{\sigma_{1}(W_{t-1})}$ is larger than a threshold, we will automatically bound the learning rate to a weighted multiple of $\frac{\sigma_{1}(W_{t-1})}{\sigma_{1}(\nabla W_t)}$, where $\nabla W_t$ is the updating quantity in step $t$. Such an optimization strategy can prevent spectral energy concentration to only a few directions, and thus can avoid malignant entropy collapse which will trigger the model crash. We conduct extensive experiments using ViT, Swin-Transformer and GPT, showing that our optimization strategy can effectively and stably train these (Transformer) models without using learning rate warmup.

Decision: accept (poster)
Reviewer scores: 8, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1729865583953

thread_type:
official_review

<thread_object>
creator: reviewer_1y1q

<content>
summary:
This paper visualizes the training dynamics of transformer models in both successful and failing cases. Based on these visualizations, the authors observe that the rapid growth of the maximal singular value of the product of the query and key matrices,  W_q^T W_k , is a critical phenomenon associated with model crashes. They provide insights suggesting that when  W_q^T W_k  becomes concentrated along the directions of the top singular vectors, the attention probability matrix becomes sparse and low-rank. They also show crashed models typically exhibit significant spectral energy concentration. Finally, the authors propose a variant of AdamW that addresses this issue by clipping the learning rate to prevent the spectral norm of all parameters from increasing too quickly.

strengths:
This paper provides insightful experiments demonstrating that the rapid growth of the spectral norm of the query-key product matrix (and other parameter matrices) is a crucial factor in model crashes. The study shows that in crashed models, the spectrum of the query-key product matrices often concentrates on the top eigenvalues. This provides novel insight on understanding transformer optimization.

weaknesses:
- The theory part does not enhance the paper as effectively as it could.
    - In Theorem 1, the statement is somewhat vague. Terms like “malignant collapse” and “spectral energy concentration” are not formally defined, making the theorem unclear. Additionally, the proof of Theorem 1, which appears in Appendix D, lacks clarity and coherence. For example, the theorem aims to establish a sufficient condition for malignant collapse, yet the proof assumes the low rank of  W  due to malignant collapse. Furthermore, the proof relies heavily on the assumption of Gaussian distribution.
    - I would personally suggest rewriting Theorem 1 as a highlighted insight or intuition, supported by experimental results, such as Figure 1 and 4.
- In Figure 5, could you elaborate on how the concentrated and low-rank structure of  $\frac{\partial \text{vec}(P)}{\partial \text{vec}(W_q^T W_k)}$  leads  W_q^T W_k  to become low rank? It still requires reshaping the matrix, multiplying the activation gradient, and summing over data from  $\frac{\partial \text{vec}(P)}{\partial \text{vec}(W_q^T W_k)}$  to the update of  W_q^T W_k, so this implication is not straightforward.
- Could you clarify how one would generally set the parameter $\tau$ for an unseen model, and how $\tau$ might vary as the model scales up?
- Additionally, does benign collapse occur in real training tasks?

questions:
See weakness above
</content>
</thread_object>

</thread>
<thread>
creation date:
1730222862018

thread_type:
official_review

<thread_object>
creator: reviewer_9a7p

<content>
summary:
The paper presents a method to eliminate learning rate warmup for training transformes. It begins by analyzing training dynamics, specifically observing the norm of certain parameters to understand behavior patterns when training is successful versus when it crashes. The authors note that, in the presence of warmup, the parameter norms tend to increase steadily, while they diverge in failed training runs. However, this analysis does not introduce new insights, as it’s known that parameter norms typically blow up when training crashes.
The authors argue that warmup can be bypassed by controlling the spectral energy concentration (SEC), which maintains the singular values of the product $W_q^TW_k$​ thus allowing successful training. To achieve this, they introduce a modified version of the AdamW optimizer using a power iteration-based step to estimate the largest singular values of the $W_q^TW_k$ matrix. Experiments with the proposed method show that it can perform similar to warmup.

strengths:
1 .The authors propose an alternate method for transformer training that does not need learning rate warmup.

weaknesses:
1. The choice of parameter norms for tracking training dynamics seems arbitrary. It’s already established in literature that parameter norms diverge when training fails, so this observation does not seem novel [1]. The paper could benefit from a deeper analysis or rationale for the specific parameters chosen, or alternatively, from an exploration of novel insights that could provide a more compelling argument.
2. While the authors claim that the attention maps are sparse and low-rank, the plots are difficult to interpret. Hence it is impossible to make conclusions based on the provided plots. It might help to change the scaling of the plots to help visualization or simply estimate the rank of the attention maps.
3. My major concern is that in the proposed algorithm that controls the step size to ensure the SEC does not collapse, only ends up modifying the step size $\alpha$. It seems to me that this is another way of performing learning rate warmup and the analysis further justifies the use of warmup as it may help prevent SEC from collapsing. In order to disentangle these ideas, the authors might need to discuss how the modified $\alpha_t$ looks over time (does it resemble warmup, which seems to be the case).
Experiments also suggest that the proposed method is in fact very similar to warmup and might be a different parameterization offering the same benefit as warmup.
4. Further, while replacing warmup the authors introduce an additional power iteration step which is computationally more expensive and adds to each gradient update step, while learning rate warmup does not require this additional step.


[1] Pennington, Jeffrey, Samuel Schoenholz, and Surya Ganguli. "Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice." Advances in neural information processing systems 30 (2017).

[2] Defazio, Aaron, et al. "When, why and how much? adaptive learning rate scheduling by refinement." arXiv preprint arXiv:2310.07831 (2023).

questions:
See above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730250811384

thread_type:
official_review

<thread_object>
creator: reviewer_yrrp

<content>
summary:
The paper analyzes a training instability associated with the self-attention layer of a Transformer by studying its training dynamics. Specifically, the paper identifies a ```Spectral Energy Concentration``` condition where-in the weights of the product of query transpose key matrix shows a rank collapse. The paper proposes to bound the base learning rate per step when the optimizer (AdamW specifically) proposes very large update vectors. The paper shows how the need to use learning rate warmup can be  relaxed for training Transformers. Empirical results are provided to support analysis and proposed algorithm.

strengths:
- The analysis of the  query (transpose), key weight matrix product  is interesting. Specifically, using the singular values of this product matrix to show collapse is interesting
- Application of Weyl's inequality to the update equation for the optimizer to stabilize Transformer training is a nice development
- The biggest contribution of this paper is that the proposed algorithm built on analysis frees up a practitioner to use  learning rate warmup.

weaknesses:
- The paper contains limited empirical results. Specifically, the experimental setup used to study the algorithm is small and no results are provided with larger models (0.5B, 1B, 3B). This limits the contribution as its not clear how the proposed algorithm will work with larger models (albeit can still be considered ``small'') where training instabilities are more apparent

- Related to above, the QK layer norm paper by Deghani et al. (2023) show that training instabilities occur at 8B for vision transformers (ViTs). So, I wonder whether the new update along with no learning rate warmup will work as advertised

- The paper proposes ``Spectral Energy Concentration (SEC) '' to quantify collapse. Wouldn't a smooth rank measure also provide the similar  information while being easier to interpret? Furthermore, this contribution is not novel as large attention logits being a problem has been observed in literature below by Deghani et al (referred to in the paper) and by [Worstman et al.](https://arxiv.org/abs/2309.14322).

- The argument in Section 3.2 is not very clear. Especially the reader could use some help with understanding why the rank of X.T being small would imply the rank of X.T \otimes X.T is still small. The equation shows a square relationship. Firming up the arguments here would be very helpful to the reader

questions:
- Would the proposed optimizer update work at reasonable scales like the ones considered by [Worstman et al.](https://arxiv.org/abs/2309.14322)?

- How different is the instability shown by analyzing weights of Q.T, K product different from logit norm growth shown previously in literature?

- (clarification request) The paper describes the product matrix in line 345 in general as a non-symmetric positive quasi-definite
square matrix. A reference to what these technical terms mean and/or definition in the appendix would help clarify what this means to the reader. The reviewer was unable to come up with a definition of quasi-definite for non-symmetric matrices but this maybe perhaps because I missed looking at an appropriate reference. 

The concerns raised in the review is that the work has limited novelty and also limited empirical support. Therefore, the initial decision for this paper is a reject.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730713240462

thread_type:
official_review

<thread_object>
creator: reviewer_qdmw

<content>
summary:
This paper deals with the challenges of scaling Transformers without relying on techniques like learning rate warmup. It provides a theoretical analysis that identifies the concentration of spectral energy in weight matrices as a key factor behind model crashes, leading to entropy collapse. To resolve this issue, the authors propose a novel optimization strategy that smooths weight updates, preventing rapid concentration of spectral energy. Extensive experiments with various Transformer models, including ViT, Swin Transformer, and GPT, demonstrate that this approach effectively trains models.

strengths:
This paper demonstrates strong originality, with the necessary mathematical definitions and proofs included as required. It is well-structured and organized clearly. The essential mathematical framework is presented both in the main text and supplementary material. Additionally, experiments have been conducted to verify the theoretical claims made in the paper.

weaknesses:
This paper has several areas for improvement in terms of writing, particularly regarding the use of mathematical notations in the abstract without proper definitions. A similar issue is evident in the Introduction as well.

The experiments are inadequate; the paper should incorporate additional baseline models for comparison, as it currently only discusses or modifies three models. It would be advantageous to present empirical results using well-known recent models, such as those with linear complexity in attention, like Linformer or Flatten Transformers.

 It would be preferable to provide proper references instead of linking to internet sources (such as Wikipedia) as seen on page 6 and elsewhere.

questions:
I have a few questions and suggestions as follows.

1. In the abstract, there are mathematical notations that should either be predefined or the abstract should be written more clearly without them.

2. In the Introduction, some notations are defined later in the text, such as "$\delta vec(P)$" mentioned in the second contribution point, which should be defined earlier.

3. The paper claims that spectral energy concentration is a key issue behind model crashes. If attention is computed without using softmax (for example, in Linear attention models like Flatten or Vicinity Vision Transformer), will the same crash occur? Is the problem specifically due to softmax, which might cause sparsity or model instability?

4. The paper states that "the attention map is a sparse yet low-rank matrix" (in the third contribution point). Does this imply that "$A = softmax(\frac{P}{sqrt{d_q}})$" is sparse? If so, please provide a reference or proof for this assertion.

5. In the experimental results, please clarify the choice of using only pure encoder architectures (like ViT) or pure decoder architectures (like GPT). Is it possible to utilize an architecture that combines both encoder and decoder components?

6. It would be preferable to provide proper references instead of linking to internet sources (such as Wikipedia) as seen on page 6 and elsewhere.

7. Please include some recent variants of Transformers in the empirical results for comparison.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732324762531

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of our revisions

comment:
We would like to thank all reviewers for their time and effort in reviewing our paper!

To address reviewers' comments and concerns, we have made the following changes:

- We added experiments on a larger vision Transformer, i.e., StableViT-giant which has 1B parameters, in **Appendix J**.

- We added experiments on a larger language Transformer, i.e., StableGPT-large which has 774M (0.774B) parameters, in **Appendix K**.

- We added experiments on a Flatten Transformer in **Appendix L**.

- We visualized the actual step size for different weights in different layers over time in **Appendix M**.

- We revised the paper, updated the figures, clarified some concepts when we first mentioned them according to the reviewers' suggestions.

The revised contents are highlighted in purple in the revised paper.  **We kindly invite all reviewers to take a look at our new experimental results!** 

We sincerely thank all reviewers again for their valuable suggestions, which have greatly helped strengthen our paper.

If you have any further questions, we would be happy to discuss them!
</content>
</thread_object>

</thread>
<thread>
creation date:
1732478446781

thread_type:
public_comment

<thread_object>
creator: ~atli_kosson1

<content>
title:
Related Work on Warmup

comment:
Dear Authors,

I read your ICLR submission with great interest. I think it is a good contribution that provides an interesting perspective on warmup. I would like to draw your attention to two of our recent works which I think may be quite relevant. In our [NeurIPS 2024 paper](https://arxiv.org/abs/2410.23922) (which is parallel work with yours), we explore the need for warmup in GPT training. One of our findings is that the need for warmup can be significantly reduced and sometimes fully eliminated by limiting the Frobenius norm of the update relative to that of the weights. This differs from your method which is based on the spectral norm instead of the Frobenius norm. Another difference is that we control the size of the update throughout training instead of clipping it, which can be done without any additional hyperparameters (for Adam / Lion like optimizers). The first paper expanded upon a similar observation we made in our earlier [ICML 2024 work](https://arxiv.org/abs/2305.17212) on how weight decay affects the training dynamics of neural networks. Would you perhaps consider discussing these as related work in a future revision?

Thank you,

Atli Kosson
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper aims to explain the mechanism behind warump in Transformer training.
It starts with a visualization of a two-layer Transformer in Section 3 by monitoring several quantities, and then develops a metric spectral energy concentration based on the sigular value of $W_q^{T} W_k$. Hence a smooth weight updating rule is proposed for warmup. The extensive experiments support the effectiveness.

The AC suggests the authors to conduct more experiments on relatively deep Transformers (e.g., three layers) to verify the experimental findings. Besides, the proposed spectral engergy concentration is quite close to effective rank used in machine learning theory, e.g., benign overfitting. This also deserves to explore. Frankly speaking, massive known results are delivered in the main text, e.g., Proposition 1, Theorem 2 (Weyl inequality), it would be best to put them into the appendix.