Rethinking Skip Connection with Layer Normalization (2020)

Abstract: Skip connection is a widely-used technique to improve the performance and the convergence of deep neural networks, which is believed to relieve the difficulty in optimization due to non-linearity by propagating a linear component through the neural network layers. However, from another point of view, it can also be seen as a modulating mechanism between the input and the output, with the input scaled by a pre-defined value one. In this work, we investigate how the scale factors in the effectiveness of the skip connection and reveal that a trivial adjustment of the scale will lead to spurious gradient exploding or vanishing in line with the deepness of the models, which could by addressed by normalization, in particular, layer normalization, which induces consistent improvements over the plain skip connection. Inspired by the findings, we further propose to adaptively adjust the scale of the input by recursively applying skip connection with layer normalization, which promotes the performance substantially and generalizes well across diverse tasks including both machine translation and image classification datasets.

## introduction there is a surge of research interest in solving the optimization problem of deep neural networks, which presents a considerable number of novel difficulties that are in need of practical solutions. batch normalization @xcite , layer @xcite and skip connection @xcite are widely-used techniques to facilitate the optimization of deep neural networks, which prove to be effective in multiple contexts @xcite . among them, batch normalization (bn) and layer normalization (ln) make progress by normalizing and rescaling the intermediate hidden representations to deal with the so-called internal co-variate shift in mini-batch training and collective optimization of multiple correlated features. skip connection bypasses the gradient exploding or vanishing problem and tries to solve the model optimization problem from the perspective of information transfer. it enables the delivery and integration of information by adding an identity mapping from the input of the neural network to the output, which may ease the optimization and allow the error signal to pass through the non-linearities. recently, several efforts have been made to deal with the optimization problem by further exploiting the existing solutions, such as exploring the implementation of skip connection @xcite and combining the skip connection with other normalization methods @xcite . some of the proposals are illustrated in figure 1 . most of these variations could be formulated as or coarsely equivalent to a missing piece from the existing work is how would the residual block perform if g is realized as layer normalization and λ does not equal one. although transformer @xcite has shown the effectiveness of the combination of layer normalization and skip connection, it is intuitive that the modulating factor λ may not always be one, especially when f is not well trained and not sufficient in representation ability. in addition, the intuition is supported by @xcite who found that if using a gate mechanism to determine the balance between the two components, the skip component wins in almost all situations. however, in the study of @xcite foot_1 , which discussed the combination of batch normalization, learned λ, and convolutional neural networks, they found empirically that the best performance is achieved when λ is fixed as one, suggesting the shortcut and the transformed input should be of equal importance. therefore, in order to explore whether the optimal ratio of x to f(x, w) is one in the residual block, and to extend the understandings of the residual block formed as @xmath0, w)), we propose three different implementations of the residual block in combination with layer normalization. the first implementation, which we name expanded skip connection, simply expands the skip connection and sets λ larger than one, following the findings of @xcite . the second variation, named as expanded skip connection with layer normalization, includes the layer normalization after the expanded skip connection, since layer normalization is observed to be helpful in facilitating the optimization of skip connection as in @xcite . the final proposal, recursive skip connection with layer normalization, is a novel combination that does not fit in the general form of the residual block, which takes the advantages of skip connection and layer normalization in a recursive manner, so that the input signal is amplified yet with layer normalization stabilizing the optimization. the illustrations of those combinations are shown in figure 2 . our experiments on diverse tasks, including machine translation and image classification with strong baselines, i.e., transformer and resnet, conclude the following: ## related work in recent years, a large number of skip connection methods have been used for neural networks @xcite @xcite @xcite @xcite @xcite @xcite . for the skip connection, there are two main types of problems. the first is how should the skip connection be incorporated into the existing neural network structures so that the best improvements can be achieved. the second is how should the neural network with skip connections be optimized so that its representation capability could be fully mined. therefore, we categorized the related work into on the connection problem and on the optimization problem. on the connection problem highway network @xcite built a highway connection from the input to the output, similar to the skip connection. however, a transform gate was proposed to control the balance of the input x and the transformed input f(x, w), instead of using identity mapping to combine the input and the transformed input. @xcite designed five types of skip connections and discussed the possible skip connections in detail. based on their theory and experiments, they argued that it is likely for a model to perform the best when the skip connection is directly combined into the output without any additional scaling such that the skip and the transformed input have equal contribution to the output. the reason is that with scaling, the gradient of the skip suffers from the gradient exploding or vanishing problem, which hinders the deep neural network from efficient optimization. on the optimization problem since the early work has established that it is sufficient to directly combine skip connection into forward propagation of neural networks without any scaling @xcite , the succeeding studies on the optimization problem mostly followed the simple yet effective architecture of the skip connection. in computer vision task, preact-resnet @xcite demonstrated that it is helpful to move the batch normalization applied to f(x, w) forward and apply it to x instead. from another point of view, it is similar to applying the batch normalization to y of the previous block. in natural language processing tasks, transformer @xcite made use of skip connection and layer normalization extensively in its architecture, and from our experiments, the layer normalization seems essential to the overall performance of the model and can help the optimization of the non-linear transformation to some extent. in all, we propose to use the form @xmath1, w)) to summarize the combination of normalization and skip connection in existing work. the form is in accordance with the findings that the normalization g should be placed after the addition of the skip and it could be beneficial to elevate the portion of the input information with λ. moreover, this work distinguishes itself in that (a) we mainly focus on the combination of layer normalization and skip connection, which is less investigated but deems promising, (b) we rethink the common practice that the skip need not scale in consideration of the help of layer normalization, (c) we conduct experiments on both representative computer vision and natural language processing tasks, while existing work only provides result in a single modality, and most importantly (d) we escape from the general form of the residual block that generalizes all previous work and propose a novel recursive construction of the residual block with layer normalization that outperforms all the variants of the general form that are examined in this work . ## architecture in this section, we describe three skip connection constructions, which are illustrated in figure 2 . they are based on the integration of more input information and application of the layer normalization. expanded skip connection (xskip) the expanded skip connection is defined similar to the constant scaling method in @xcite as expanded skip connection with layer normalization (xskip+ln) motivated by transformer, which combines skip connection with layer normalization, we further examine the effect of layer normalization on expanded skip connection, which takes the form of layer normalization may help mitigate the gradient malformation in optimization caused by the modulating factor. we use the parameterized layer normalization as suggested in @xcite , because different input features naturally have different importance. different from batch normalization which acts in "sample space", layer normalization acts in "feature space". considering the special case that f(x, w) is ill-learned and outputs random values and the examples @xmath2, batch normalization simply ignores the input feature f which results in uneventful optimization according to mini-batch statistics, while layer normalization still recognizes the input feature f and its value can be properly adjusted according to the other input features. the same may also apply to the situation where λx is dominant compared to f(x, w) and the transformed input falls into the range comparable to random noise. in such situations, which may happen when the neural network is hard to optimize, layer normalization could still help the learning of the shortcut, while batch normalization might fail. recursive skip connection with layer normalization (rskip+ln) another way to stabilize the gradient is to keep λ = 1 each time but repeatedly add the shortcut with layer normalization, such that more input information is also modeled. it is defined recursively as ## empirical study in this section, we examine the effect of the discussed combinations on two representative baselines through empirical studies. first, we conduct exploratory experiments on datasets that are relatively small to gain insights of the combinations, such as the choice of λ and the effectiveness of the constructions. then, we discuss the observed results and conduct further analysis. finally, we validate our findings on larger and more difficult datasets. ## conclusion and future work in this work, we focus on exploring the combination of skip connection and layer normalization to build a novel skip connection architecture able to incorporate reasonable skip information at output level and solve the gradient explosion problem in practice. according to the empirical studies, we find that expanding skip information should be helpful to the optimization of the model, which is hindered by the gradient malformation in practice, and layer normalization does better than batch normalization in mitigating such problem. the proposed recursive skip connection with layer normalization guarantees the stability of the shorcut gradient by recursively dividing the expanded skip connection and adding it to the output iteratively with layer normalization. the refined skip connection architecture may enable more efficient optimization. validation on cifar-100, wmt-2014 en-de and en-fr proves the findings and the effectiveness of the proposed recursive skip connection construction with best improvements over 1.0 absolute point, demonstrating its generalization ability to a wide range of existing systems. unfortunately, a limitation of our investigation is that λ is selected as integer and only a small number of values are used, which may not be of the best granularity for expanding the skip connection. a more desirable solution would be to adaptively determine the value of λ, maybe in terms of the gradient distribution during the learning so that more efficient optimization could be realized. we will further seek to address this limitation in future work. because the recursive definition, the inner gain parameters of layer normalization could penetrate through the normalization and effectively act as the learnable ratio. it is also important to note that the ratio is also dependent on the current layer distribution, which could serve as a regularization component. for any λ larger than 2, the general form of the ratio is 1 + λ @xmath3. it should be reminded that setting λ = 2 is sufficient for automatically learning the shortcut-residual ratio and dynamically adjusting the ratio based on instant data distribution. yet with larger λ, the distribution of the higher order combination of x and f(x) will be considered. the recursive definition introduces gradual and smooth combination of the input and the residual up to the specified order λ which could deal with different distributions better. our experimental results also demonstrate that the proposal performs better than the scaled shortcut.