B ranch N orm: Robustly Scaling Extremely Deep Transformers (2024)

Abstract: Recently, DeepNorm scales Transformers into extremely deep (i.e., 1000 layers) and reveals the promising potential of deep scaling. To stabilize the training of deep models, DeepNorm attempts to constrain the model update to a constant value. Although applying such a constraint can benefit the early stage of model training, it may lead to undertrained models during the whole training procedure. In this paper, we propose BranchNorm, which dynamically rescales the non-residual branch of Transformer in accordance with the training period. BranchNorm not only theoretically stabilizes the training with smooth gradient norms at the early stage, but also encourages better convergence in the subsequent training stage. Experimental results on multiple translation tasks demonstrate that BranchNorm achieves a better trade-off between training stability and converge performance.

## introduction in recent years, transformers @xcite have been developed rapidly and achieved state-ofthe-art (sota) performance on a wide range of tasks. meanwhile, the model capacity gets substantially expanded by widening the model dimension @xcite @xcite . given that deep neural models learn feature representations with multiple layers of abstraction @xcite , it is more attractive to increase the capacity of the model by scaling depths rather than widths. unfortunately, due to the training instability of transformers, the depths of these sota models are still relatively shallow @xcite . to stabilize the training of transformers, there have been various efforts on better architectures @xcite corresponding author. 2023), or the implementation of proper initialization @xcite @xcite . among them, the most representative approach is deepnorm @xcite , which first scales transformers to 1000 layers and significantly outperforms existing shallow counterparts. specifically, deepnorm aims to constrain the model update to a constant level by upweighting the residual connections in transformer and reducing the variance of parameter initialization. as a result, the stability of transformers is improved in the early training stage. however, such constraints on the magnitude of parameter updates may ultimately yield undertrained models during the subsequent training stage. to verify the above conjecture, we first conduct experiments on shallow models to ensure that all models can be trained stably to convergence. as shown in figure 1 , it is observed that deepnorm brings a certain degree of performance decline on vanilla transformers, and this issue tends to get worse when models get deeper. to address the above issue, we propose a simple yet effective approach to robustly scale extremely deep transformers, named branchnorm. specifically, the non-residual branch foot_0 of the transformer is dynamically rescaled according to the training period. in the early stage of model training, branchnorm theoretically stabilizes the training with smooth gradient norms. while in the subsequent training stage, branchnorm progressively degenerates into a vanilla post-layernorm (i.e., post-ln) to promote better convergence. experiments on a wide range of tasks show that branchnorm brings consistent improvement over deepnorm and effectively alleviates the above undertrained issue. moreover, branchnorm performs more robustly on some key hyperparameters (e.g., warmup, and learing rate) than deepnorm, which makes it likely to be a portable alternative for scaling extremely deep transformers. the contributions of this paper can be summarized as follows foot_1 : ## background in this section, we first provide a brief overview of the difference between post-ln and pre-ln in, and subsequently introduce the existing deepnorm. shown in figure 2 . subsequently, @xcite further analyze that pre-ln may have an excessive reliance on its residual connections, which inhibits the model from unleashing its full potential. motivated by the above observation, we base our approach on post-ln in the remainder of our experiments. formally, within the l-th sub-layer of the post-ln transformer, given the input x l , the subsequent output x l+1 is computed as follows: deepnorm. deepnorm follows the post-layernorm (i.e., post-ln) architecture and rescales the residual branch with a scalar factor α > 1. formally, the l-th sub-layer in deepnorm conduct calculations as follows: note that β merely affects the model initialization, whereas α is used and remains fixed throughout the whole training procedure. moreover, with the model gets deeper, deepnorm will assign larger value to α, which leads to the model outputs x l+1 depend too much on the residual branch αx l . in turn, the model parameters θ l on the non-residual branch may get insufficient training. ## approaches in this section, we first analyze the instability of post-ln from the perspective of gradient norm, then demonstrate how deepnorm can alleviate the unbalanced gradients to a certain extent, and finally introduce our proposed method branchnorm. ## experiment settings we conduct extensive experiments on bilingual translation, text summarization and language modeling tasks to verify our approach. in this section, we will describe our experimental settings and present the main results. ## analysis in this section, we take the machine translation task as an example and conduct analytical experiments with a 200-layer (i.e., 100l-100l) transformer on wmt14 en-fr dataset. we first verify the robustness of branchnorm to hyperparameters, and then analyze the decoding speed of deep models with shallow decoders, and finally analyze the issue of parameter redundancy. ## conclusion in this paper, we first explore the undertraining problem of deepnorm from a gradient perspective and propose a more flexible approach, namely branchnorm, which theoretically stabilizes the training with smooth gradient norms at the early stage. once the early phase of training instability is passed, branchnorm can then degenerate to a standard post-ln, thus encouraging better convergence performance. experiment results on various tasks show that branchnorm achieves a better trade-off between training stability and converge performance. for instance, branchnorm improves deepnorm by up to 1.5 bleu, 1.1 rouge, and 2.1 perplexity on machine translation, text summarization, and language modeling tasks, respectively.