B 2 T Connection: Serving Stability and Performance in Deep Transformers (2023)

Abstract: In the perspective of a layer normalization (LN) position, the architecture of Transformers can be categorized into two types: Post-LN and Pre-LN.Recent Transformers prefer to select Pre-LN because the training in Post-LN with deep Transformers, e.g., ten or more layers, often becomes unstable, resulting in useless models. However, in contrast, Post-LN has also consistently achieved better performance than Pre-LN in relatively shallow Transformers, e.g., six or fewer layers. This study first investigates the reason for these discrepant observations empirically and theoretically and discovers 1, the LN in Post-LN is the source of the vanishing gradient problem that mainly leads the unstable training whereas Pre-LN prevents it, and 2, Post-LN tends to preserve larger gradient norms in higher layers during the back-propagation that may lead an effective training. Exploiting the new findings, we propose a method that can equip both higher stability and effective training by a simple modification from Post-LN.We conduct experiments on a wide range of text generation tasks and demonstrate that our method outperforms Pre-LN, and stable training regardless of the shallow or deep layer settings.

## introduction to prevent the vanishing (or exploding) gradient problem in the training of a deep neural network (dnn), various techniques, such as batch normalization @xcite and residual connection @xcite , have been proposed and widely used in almost all recent dnns. transformer @xcite employs the layer normalization @xcite for this purpose. transformer is currently the most successful model architecture in dnns. it was firstly developed for applying sequence-to-sequence tasks, such as machine translation @xcite , summarization @xcite , and automatic speech recognition (asr) @xcite , and is currently used in speech, vision, and many other information processing research fields. as reported in the batch normalization literature @xcite , the position of the normalization layers primarily affects both the stability and resultant performance of a trained model. in transformers, some previous studies have investigated the impact of the layer normalization positions @xcite . there are currently two major layer normalization positions in transformers: pre-layer normalization (pre-ln) and post-layer normalization @xcite . pre-ln applies the layer normalization to an input for each sub-layer, and post-ln places the layer normalization after each residual connection. the original transformer @xcite employs post-ln. however, many recent studies have suggested using pre-ln @xcite @xcite because the training of deep transformers (e.g., those with ten or more layers) using post-ln is often unstable, resulting in useless models. figure 1 shows loss curves for an actual example; the training of 18 layered transformer encoder-decoders (18l-18l) on a widely used wmt english-to-german machine translation dataset. these figures clearly show that the post-ln transformer encoder-decoders fail to train the model. however, in contrast, @xcite reported that post-ln consistently achieved better performance than pre-ln on a machine translation task when they used 6 layered (relatively shallow, 6l-6l) transformers. this paper focuses specifically on such discrepancies between pre-ln and post-ln in configurations with various number of layers. we investigate the sources of the instability of training in deep configurations and the superior performance in shallow configurations for post-ln, compared with that for pre-ln, to understand the essentials of the differences between @xcite we discover that the layer normalization in post-ln is the main source of the vanishing gradient problem that leads to unstable training, whereas pre-ln prevents it, as shown in figure 1 . in particular, we clarify that the layer normalization is a significant factor of the vanishing gradient problem by comparing the input/output vector norms of gradient flows for each layer normalization during back-propagation. these analyses bring us a novel idea that can satisfy higher stability by skipping over layer normalizations and provide better performance than pre-ln regardless of their layer sizes. consequently, we propose a method that is based on post-ln transformers but has additional residual connections to enable stable training. we conduct experiments on a wide range of text generation tasks, namely machine translation, summarization, language modeling, and asr. the experimental results lead to the following three new major findings: ## post-ln and pre-ln transformers we briefly describe post-ln and pre-ln transformers in this section. the original transformer @xcite uses post-ln, in which layer normalizations are located after each residual connection. let x be an input of a sublayer, and f(•) be a sub-layer of a transformer, such as a feed-forward network or multi-head attention. post-ln is defined as follows: in contrast, pre-ln places the layer normalization before an input of each sub-layer: ## gradients of transformer layers as described in @xcite , the vanishing gradient problem often occurs in post-ln transformers. figure 3 shows the gradient norms of each layer for the (a) encoder-side and (b) decoderside at the beginning of training, when 18l-18l transformer encoder-decoders are trained on a widely used machine translation dataset (the wmt english-to-german dataset). focus on the decoderside of post-ln as illustrated in figure 3 (b) . this figure shows that shallower layers have smaller gradient norms. in other words, the vanishing gradient occurs in the decoder-side of post-ln because its gradient norms exponentially decay as they are back-propagated to shallower layers. this result is consistent with the previous study @xcite . we consider that this vanishing gradient causes the difficulty of stacking many layers with the post-ln setting, as shown in figure 1 . to investigate the vanishing gradient empirically in more detail, we measure the gradient norms of parts (1) -(5) of figure 2 ( @xcite . figure 4 shows the gradient norms of each part in the 18th layer foot_0 . this figure shows that the gradient norms decrease drastically from (4) to ( 3 ) and ( 2 ) to ( 1 ). these parts correspond to layer normalizations, as shown in figure 4 . this suggests that layer normalizations in post-ln transformers are probably the cause of the vanishing gradient problem. 13 10 0 10 1 10 1 10 -1 10 0 (a) encoder side (b) decoder side 10 0 10 1 10 1 10 -1 10 0 figure 3: gradient norms of 18 layered transformerbased encoder-decoder architectures. to investigate the difference between the gradient flows of post-ln and those of pre-ln theoretically, we calculate the derivatives of equations ( 1 ) and (2), as follows: layer norm attention ffn layer norm layer norm attention the difference between these equations implies that the residual connection in pre-ln prevents the vanishing gradient because it retains the gradients of upper layers even if the derivative of the layer normalization decreases gradients drastically. ## transformations by each layer as described, it is difficult to stack many layers in post-ln transformers because the vanishing gra-3080 dient problem occurs. although pre-ln is more stable in training, post-ln can achieve better performance if training succeeds (see section 6). in this section, we explore the reason for this difference in performance. focus pre-ln in figure 3 . in contrast to post-ln, in pre-ln, a deeper (higher) layer has a smaller gradient norm. thus, the parameters of higher layers are not required to change dramatically from their initial values. this implies that higher layers in pre-ln are not sufficiently effective. to investigate the effectiveness of higher layers, we focus on the transformations by each layer. figure 5 shows the average cosine similarities between the outputs of each pair of layers for 6l-6l transformer encoder-decoders trained on the wmt dataset when several sequences are input. this figure indicates that the lower-left similarities of pre-ln are higher than those of post-ln. this result means that the outputs of shallow layers are similar to the output of the final layer in pre-ln, but not in post-ln. consequently, higher layers in pre-ln are less effective than those in post-ln if training succeeds. we consider that the residual connection in pre-ln causes this phenomenon. as equation (2) shows, in pre-ln, an input x skips over the sublayer f(•) by the residual connection. thus, the input x is directly connected to the final layer output. this property makes the training stable, as described in section 3, but causes high similarities between the outputs of the various layers. therefore, we consider that pre-ln underperforms post-ln because the residual connection in pre-ln reduces the effectiveness of its higher layers. in contrast, in post-ln, larger gradient norms in higher layers (as shown in figure 3 ) make higher layers more ef-fective (as shown in figure 5 ) but it is necessary to prevent the vanishing gradient problem in shallow layers when we stack many layers. ## modification for stable training in post-ln: bottom-to-top connection this section introduces a modification that makes the training of post-ln more stable while preserving its high performance. this modification comprises an additional residual connection to mitigate the vanishing gradient in post-ln by enabling many layers to be stacked. as discussed in the previous sections, we need a term that retains gradients in the derivatives, as in equation ( 4 ), to prevent the vanishing gradient. to satisfy this requirement, we propose a residual connection that skips over all layer normalizations except the final one in each layer. our introduced connection ties an input of a layer to the result of the feed-forward network (ffn), as illustrated by the red arrows in figure 2 (c ). we call this connection bottom-to-top (b2t) connection, which is formalized in the following equation: the derivative of this equation is the following equation: because this derivative contains i, which is unrelated to the derivatives of internal layer normalizations, our b2t connection (i.e., x inp ) helps to propagate gradients. for a decoder-side, we can prove this property in the same manner. figure @xcite indicates that b2t connection mitigates the vanishing gradient of 18l-18l encoderdecoders. moreover, we locate b2t connection before the final layer normalization in each layer to avoid a direct connection to the final layer output based on the discussion in section 4. thus, b2t connection preserves the property of post-ln with respect to the transformations performed by each layer, as illustrated in figure 5 (c) foot_1 . ## experiments through experiments, we indicate following three findings. we describe the essential experimental configurations in this section. appendix a presents more details, such as the hyper-parameters and computational budgets. ## related work layer normalization @xcite is a useful technique for training neural networks but its mechanism has been unclear @xcite . the transformer, which is the standard architecture for various tasks, also contains layer normalizations. the original transformer architecture adopted the post-ln configuration @xcite . however, recent transformer implementations have adopted pre-ln configurations @xcite @xcite . to construct deep transformers that achieve better performance, recent studies have focused on the behavior of layer normalizations. @xcite indicated the difficulty of training deep transformers with post-ln due to the vanishing gradient problem, and demonstrated that pre-ln enables the stacking of many layers through machine translation experiments. in addition, they proposed a method to connect all layers to increase the effectiveness of deep transformers. @xcite and @xcite also proposed such connection methods to stack many layers. @xcite introduced additional connections into attention sub-layers to improve the performance. @xcite explored the relation between the warm-up strategy and layer normalizations in transformers. through theoretical and empirical analyses, they indicated that post-ln requires the warm-up strategy to stabilize the training. @xcite analyzed the training dynamics of post-ln and pre-ln transformers. they then proposed admin, which consists of additional weight parameters to control the variances of outputs from each sub-layer. in contrast, we indicated that we can stabilize the training of post-ln transformers by adding only a residual connection that skips over layer normalizations that cause the vanishing gradient. some studies have proposed initialization methods to make the training of deep neural networks stable @xcite . @xcite proposed the depth-scaled initialization to prevent the vanishing gradient problem in transformers. @xcite proposed the fixed-update initialization to remove normalizations in neural networks. inspired by these studies, @xcite proposed t-fixup, which enables both warm-up and layer normalizations to be removed from transformers. in addition to the initialization scheme, @xcite introduced weights into residual connections before layer normalizations, following @xcite . ## conclusion in this study, we addressed the stability of training post-ln transformers. through theoretical and empirical analyses, we indicated that layer normalizations cause the unstable training when many layers are stacked. in addition, we investigated the reason for the different performance of pre-ln and post-ln by transformations of each layer. we introduced b2t connection to prevent the vanishing gradient while preserving the advantage of post-ln. we conducted experiments on various tasks. the experimental results led to the following three findings; 1, post-ln achieved better performance than pre-ln if its training succeeded. 2, our modification enabled the training of deep transformers (e.g., those with ten or more layers). 3, our modification preserved the benefit of post-ln, and therefore outperformed pre-ln.