Stable-Transformer: Towards a Stable Transformer Training

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

Abstract:
The scale of parameters in Transformers has expanded dramatically—from hundreds of millions to several trillion. A key challenge when scaling the model to trillions is the training instability. Although many practical tricks, such as learning rate warmup, query-key normalization and better weight initialization, have been introduced to mitigate the training instability, a rigorous mathematical understanding of why such instabilities happen and why the above-mentioned tricks work well is still unclear. In this paper, we give a theoretical analysis of the initialization, normalization and attention mechanism in Transformers, and present a set of stabilized designs of the initialization, normalization and attention mechanism, which are thus termed as StableInit, StableNorm and StableAtten, individually. In experiments, we demonstrate that each of our stabilized designs, i.e., StableInit, StableNorm and StableAtten, exhibits better stability. Furthermore, by putting the stabilized designs together, we propose a stabilized Transformer, termed Stable-Transformer, and show in experiments on large model (1B parameters) and deep model (200 layers) that our proposed Stable-Transformer
achieves a more stable training process.

Decision: reject
Reviewer scores: 5, 6, 6

Reviews and discussion:
<thread>
creation date:
1730540008038

thread_type:
official_review

<thread_object>
creator: reviewer_6zea

<content>
summary:
This paper is motivated by the observation that a rigorous mathematical understanding of why training instabilities occur in transformers—and why current stabilization techniques work—is still lacking. To address this, the authors provide a theoretical analysis of transformer initialization, normalization, and attention mechanisms, proposing a set of stabilization methods for each. Experimental results show that these individual stabilizations enhance training stability. Furthermore, by combining these approaches, the authors introduce a new model, the Stable-Transformer, which demonstrates a more stable training process in experiments.

strengths:
1.  The paper is well-organized and clearly articulated. The problem statement, proposed solution, and experimental results are presented in a logical and accessible way. The figures and diagrams effectively illustrate the key concepts and findings. 
2.  The paper holds significant value as it thoroughly analyzes the instability issues associated with training transformers. It also introduces a stabilized model, the Stable-Transformer, and demonstrates through experiments that this model achieves a more stable training process.
3. The proposed method is firmly based on robust theoretical analysis. The mathematical formulation and proofs are clearly articulated and well-supported.

weaknesses:
Overall, I'm satisfied with this paper and its core contributions. Other minor concerns/suggestions are listed as follows:

Have the authors considered assessing Stable-Transformer on large multi-modal models such as CLIP or Flamingo? Evaluating these models could showcase the wider applicability of the proposed stabilization techniques beyond just language and vision transformers.

questions:
Considering the encouraging outcomes with GPT and ViT models, what do the authors identify as the key next steps for validating and enhancing Stable-Transformer? Are there specific challenges you foresee in implementing these techniques on larger models or different modalities?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730587283580

thread_type:
official_review

<thread_object>
creator: reviewer_yoj2

<content>
summary:
The paper introduces several techniques to stabilize transformer training: StableInit, StableNorm, and StableAtten, which replace Xavier initialization, LayerNorm, and Attention, respectively. For each component, the authors first identify limitations in previous methods and key factors contributing to stability, then propose the solutions. They provide theoretical analyses and experimental validation of each technique on GPT2-S and ViT backbones, demonstrating improved training stability over baselines. Additionally, the supplementary materials show that these stabilization techniques enable training with larger learning rates.

strengths:
1. The authors identify instability issues in each vanilla module when training large models and propose insights to address them.

2. Experiments on GPT and ViT verify that the proposed approaches enhance training stability, with each experiment and theoretical analysis supporting the stability claims.

3. The authors demonstrate that these methods enable the network to tolerate larger learning rates.

4. The paper is well-written, with a clear structure that makes it easy to follow.

weaknesses:
1. While each method contributes to improved stability, the gains from each proposed module are incremental. It would strengthen the argument to demonstrate that alternative methods, such as enforcing Lipschitz continuity in normalization [1], cannot replace each proposed module (e.g., showing that StableNorm outperforms other normalization techniques). Comparing the proposed methods with alternative stability techniques, rather than just naive baselines, would be more compelling. How does StableNorm compare to other recent normalization techniques like DeepNorm? And how does StableAtten compare to L2 Self-Attention [2]?

2. Although these methods make training more stable and tolerant of larger learning rates, the authors do not present other potential benefits of stable training. Emphasizing the importance of stability improvements would strengthen the argument. Does improved stability benefit generalization, lead to faster convergence, or improve performance on out-of-distribution data?

3. It would be valuable to test stability across an increasing number of layers and demonstrate whether improved stability benefits downstream tasks or enhances robustness to sample or label noise. Can the proposed techniques enable training with more layers, such as 200 or 1000 layers? Does improved stability benefit fine-tuning ViT for image classification on smaller datasets like CIFAR-10/100, or make it more robust to distribution shifts, such as on CIFAR-10-C and CIFAR-100-C?



[1] Gouk et al., "Regularisation of neural networks by enforcing Lipschitz continuity," Machine Learning, 2021.

[2] Kim et al., The Lipschitz Constant of Self-Attention, ICML 2021.

questions:
Please see the weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730644198704

thread_type:
official_review

<thread_object>
creator: reviewer_xfs7

<content>
summary:
This paper proposes theoretical and experimental approaches to improve the training stability of Transformer models. To address the instability issues arising from the expansion of parameters in large-scale Transformers, the authors introduce new initialization (StableInit), normalization (StableNorm), and attention mechanisms (StableAtten). These three techniques are combined to form a "Stable-Transformer," which the authors demonstrate experimentally to improve training stability and performance.

strengths:
Theoretical Contribution: The paper provides a theoretical analysis of the training instability in Transformers and proposes improvements to initialization, normalization, and attention mechanisms to address this instability. This is a significant contribution to the understanding and enhancement of Transformer model training.

Novel Stabilization Techniques: The authors propose StableInit, StableNorm, and StableAtten to stabilize the initialization, normalization, and attention mechanisms, respectively. By combining these, the authors offer a comprehensive solution to stabilize the entire Transformer architecture.

weaknesses:
Lack of Scalability Validation for Large Models: While the paper provides a theoretical foundation and experimental validation of the stabilization techniques, it lacks empirical testing for scalability to larger models, such as massive language models.
Implementation Complexity of StableNorm and StableAtten: The proposed techniques, StableNorm and StableAtten, require additional hyperparameter tuning, which may increase implementation complexity. This complexity can limit the practical applicability of these methods in real-world scenarios.
Limited Experimental Scope: The experiments are limited to GPT and ViT, with no validation across a broader range of Transformer architectures or real-world applications. This raises questions about the generalizability of the proposed stabilization techniques.

questions:
see weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1732160402933

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of our revisions

comment:
We would like to thank all reviewers for their valuable time and great 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 that has 1B parameters, in **Appendix J**. (xFs7)

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

- We added experiments on a 200-layer vision Transformer, i.e., StableViT-200 that has 1.44B  parameters, in **Appendix L**. (yoj2, xFs7)

- Based on the newly added experiments, we further clarify our practical contribution in the abstract and introduction section.

The revised content is highlighted in purple in the revised paper. We kindly invite the reviewer to check the newly added experiments.

We sincerely thank all reviewers again for their valuable suggestions and insightful comments, which have greatly helped strengthen our paper. For any further questions or concerns, we would be happy to discuss them further.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732678824992

thread_type:
public_comment

<thread_object>
creator: ~yuzhu_wang1

<content>
title:
Some  confusion about Fig.1

comment:
Dear authors:

I read the ICLR submission with great interest. I think it provides a good contribution on training Transformers. But I have some confusion about Fig.1, and I look forward to the authors providing explanations.

(1) Why does the accuracy of baseline ViT-L-150-epoch suddenly increase at about 60 epochs?

(2) Why does the accuracy of StableViT suddenly exceed the baseline at the end of training, for example, at 250 epochs? Training 150 epochs has similar phenomena.

(3) Note that the classic Masked AutoEncoder provides training recipes and reports 82.6% accuracy on ImageNet for ViT-L trained from scratch. The result is higher than the baseline (81.3%) and StableViT-L (82.4%). Importantly, the training process ViT-L-300-epoch (green curve of Fig. 1)  showed no instability.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a set of techniques for stable Transformer training: StableInit, StableNorm and StableAtten. StableInit initializes parameters differently from the Xavier initialization, motived by a Lipschitz constraint; StableNorm generalizes RMSNorm by introducing an additional hyper parameter alpha, and StableAtten builds on top of QKNorm and StableNorm which normalizes Q K and scales the attention logits differently. Experiments are ran on GPT2 alike LM training and ViTs. This paper receives borderline ratings and it is indeed a borderline paper after the AC's assessment.

First and foremost, stabelizing Transformer training is a very important task, but also one that has a high bar for showing meaningful progress, and there are plenty of techniques that have been presented in the literature. Therefore, we need to carefully evaluate if this paper's contribution convincingly justifies the claim of a "Stable Transformer". Now if we look at the techniques, they are quite simple -- you can think of them as rescaling either the initial scale of parameters (StableInit), or the outputs activations at certain locations (StableNorm and StableAtten). In particular, in doing so, it introduces another hyper parameter alpha, which is shown to be important to performance -- and I find the reliance on this hyper parameter unsatisfying. Of course, a method should not be criticized for its simplicity (because it can be a plus if the results are outstandingly good). Now if we look at evaluations, there are much left to be desired. Essentially, there is a lack of comprehensive comparisons to other stable training techniques, also the scope of the evaluation seems to be quite limited (there are two tasks, but again the claim of stable Transformer is much broader than what these two tasks reflect, eg what about encoder decoder tasks or other modalities). I find it pretty hard to draw a clear conclusion that the proposed techniques are necesarily better than what we already know.

Based on these considerations, I don't think this work is ready for publication. I suggest either redefining the scope and claims of this work -- either to focus on a more specific aspect of Transformer training than claiming a general stable Transformer, and/or including more comprehensive comparisons to the literature.