Stable-Transformer: Towards a Stable Transformer Training
Transformer;Stable Transformer
Yes
No ethics review needed.
see weakness
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.
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.
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.
Yes
No ethics review needed.
Please see the weaknesses.
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.
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.
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.
Yes
No ethics review needed.
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?
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.
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.
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.
We introduce Stable-Transformer that achieves a more stable training process.
@inproceedings{
anonymous2024stabletransformer,
title={Stable-Transformer: Towards a Stable Transformer Training},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=lkRjnNW0gb},
note={under review}
}
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 \textit{StableInit}, \textit{StableNorm} and \textit{StableAtten}, individually. In experiments, we demonstrate that each of our stabilized designs, \ie \textit{StableInit}, \textit{StableNorm} and \textit{StableAtten}, exhibits better stability. Furthermore, by putting the stabilized designs together, we propose a stabilized Transformer, termed \textit{Stable-Transformer}, and show in experiments that our proposed \textit{Stable-Transformer} 
achieves more stable training process.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Transformer Stable Transformer
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/edd02ff0a6fa0fee8039f05904f204f3d9aeef95.pdf
optimization
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Stable-Transformer: Towards a Stable Transformer Training
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission