Understanding the Failure of Batch Normalization for Transformers in NLP
Abstract
Batch Normalization (BN) is a core and prevalent technique in accelerating the training of deep neural networks and improving the generalization on Computer
Vision (CV) tasks.
However, it fails to defend its position in Natural Language
Processing (NLP), which is dominated by Layer Normalization (LN).
In this paper, we are trying to answer why BN usually performs worse than LN in NLP tasks with Transformer models.
We ﬁnd that the inconsistency between training and inference of BN is the leading cause that results in the failure of BN in NLP.
We deﬁne Training Inference Discrepancy (TID) to quantitatively measure this inconsistency and reveal that TID can indicate BN’s performance, supported by extensive experiments, including image classiﬁcation, neural machine translation, language modeling, sequence labeling, and text classiﬁcation tasks.
We ﬁnd that
BN can obtain much better test performance than LN when TID keeps small through training.
To suppress the explosion of TID, we propose Regularized BN (RBN) that adds a simple regularization term to narrow the gap between batch statistics and population statistics of BN.
RBN improves the performance of BN consistently and outperforms or is on par with LN on 17 out of 20 settings, involving ten datasets and two common variants of Transformer.
Introduction
Deep learning [19] has revolutionized Computer Vision (CV) [18] and Natural Language Processing (NLP) [39].
Normalization layers are key components to stabilize and accelerate the training in
Deep Neural Networks (DNNs).
In CV, Batch Normalization (BN) [15] is the default normalization technique and reveals superior performance over other normalization techniques in image recognition tasks by enforcing the input of a neuron to have zero mean and unit variance within a mini-batch data.
Furthermore, a growing number of theoretical works analyze the excellent properties of BN in beneﬁting optimization [15, 34, 4, 12, 7, 8].
While BN almost dominates in CV with empirical success and theoretical properties, Layer Normalization (LN) is the leading normalization technique in NLP, especially for Transformer models that achieve the state-of-the-art performance on extensive tasks, including machine translation [39], natural language understanding [9], text generation [32], few shot learning [5], to name a few.
As a direct substitute of LN, BN performs poorly in Transformer for neural machine translation [36].
It remains elusive to explain the failure of BN in NLP community.
In this work, we are trying to take a step forward.
Our contributions are summarized as follows:
• We ﬁnd that the inconsistency between training and inference leads to the failure of BN in NLP, supported by our extensive experiments, including image classiﬁcation, neural machine translation, language modeling, sequence labeling, and text classiﬁcation tasks.
• We deﬁne Training Inference Discrepancy (TID) to quantitatively measure this inconsistency and show that TID can serve as an indicator of BN’s performance.
In particular, BN reaches much better test performance than LN when TID keeps small through training, e.g. , in image recognition and language modeling tasks.
• We propose Regularized BN (RBN) that adds a regularization term in BN to penalize and reduce the TID when the TID of BN is large.
We reveal the optimization advantages of RBN over LN by exploring the layer-wise training dynamics of Transformer. • We empirically show that RBN can exceed or match the performance of LN, sometimes with a large margin, on 17 out of 20 settings, involving ten datasets and two common variants of Transformer.
Besides, RBN introduces no extra computation at inference compared to LN.
Related Work
Analyses of BN’s Success As BN becomes an indispensable component in deep neural networks deployed in CV tasks, a bunch of works explore the theoretical reasons behind its success.
From the view of optimization, the original BN paper [15] argues that BN can reduce internal covariate shift and thus stabilize the training, while Santurkar et al.
[34] debate that BN could smooth the loss landscape and thus enable training of neural network with larger learning rate [4].
Daneshmand et al.
[7, 8] prove that a stack of randomized linear layers and BN layers will endow the intermediate features of neural network with sufﬁcient numerical rank as depth increases, which is beneﬁcial for optimization and learning discriminative hierarchical features.
Huang et al.
[12] show that BN could improve the layer-wise conditioning of the neural network optimization by exploring the spectrum of
Hessian matrix with block diagonal approximation [26].
From the view of generalization, Ioffe and
Szegedy [15], Luo et al.
[23], Li et al.
[20], Wu and Johnson [41] argue that BN serves as regularizer which reduces over-ﬁtting when its stochasticity is small and may have detrimental effect when it is large [41].
Huang et al.
[11] further propose Stochastic Normalization Disturbance (SND) to measure such stochasticity and shows that large SND will hinder the training of neural networks.
Training Inference Inconsistency of BN Normalizing along the batch dimension usually introduces training inference inconsistency since mini-batch data is neither necessary nor desirable during inference.
BN uses population statistics, estimated by running average over mini-batch statistics, for inference.
The training inference inconsistency usually harms the performance of BN for smallbatch-size training since the estimation of population statistics could be inaccurate [40].
One way to reduce the inconsistency between training and inference is to exploit the estimated population statistics for normalization during training [14, 6, 45, 48, 47].
These works may outperform BN when the batch size is small, where inaccurate estimation may be the main issue [15, 16], but they usually work inferior to BN under moderate batch-size training [22].
Another way to reduce the inconsistency is estimating corrected normalization statistics during inference only, either for domain adaptation [21], corruption robustness [35, 29, 2], or small-batch-size training [37, 38].
We note that a recent work [13] investigates the estimation shift problem of BN.
Unlike this work that addresses the accumulated estimation shift due to the stack of BNs for CNNs in CV tasks, our work pays more attention to how the training inference inconsistency of BN correlates with its performances for
Transformers in NLP tasks.
Besides, the estimation shift of BN deﬁned in [13], which addresses the differences between the estimated population statistics and the expected statistics, differs from our
TID of BN that addresses the differences between the mini-batch statistics and populations statistics.
Exploring the Failure of BN in Transformer Similar to our work, Power Normalization (PN) [36] also investigates the reason behind the failure of BN in Transformers.
Our work signiﬁcantly differs from PN [36] in the following facets.
PN attributes the failure of BN to the unstable training of BN incurred by ﬂuctuated forward and backward batch statistics with outlier values, while we observe that the training of BN is as good as LN and the inconsistency between training and inference of BN matters more.
Based on our observation, we propose a regularization term to reduce the TID of BN.
Compared with PN, which incorporates a layer-scale layer (root mean square layer normalization [49] without afﬁne transformation [43]), our method introduces no extra computation at inference.
Besides, we use a more reasonable index to measure inconsistency which is invariant to the scale of data.
Furthermore, we show that our RBN can improve the layer-wise training dynamics of LN, which reveals the optimization advantages of RBN.
[Figure]
The training of Transformer BN is better than Transformer LN while the validation loss/BLEU of
Transformer BN underperforms that of Transformer LN after 8 epoch.
At the end of the training,
Transformer BN falls behind Transformer LN with large BLEU scores.
Lower loss and higher BLEU scores indicate better performance.
Based on the inconsistency of training and validation performance of BN, we hypothesize that the training inference discrepancy of BN causes its performance degradation.
Analyses of Training Inference Inconsistency in Transformer BN 3.1 Preliminary
Batch Normalization (BN) [15] is typically used to stabilize and accelerate DNN’s training.
Let x ∈ R d denote the d -dimensional input to a neural network layer. During training, BN standardizes each neuron/channel within m mini-batch data by[-
-]
[Equation]
where µ B,j = 1 m (cid:80) mi
=1 x ( i ) j and σ 2 B,j = 1 m (cid:80) mi =1 ( x ( i ) j − µ B,j ) 2 are the mini-batch mean and variance for each neuron, respectively.
Note that an extra small number (cid:15) is usually added to the variance in practice to prevent numerical instability.
During inference, the population mean µ and variance σ 2 of the layer input are required for BN to make a deterministic prediction [15] as:
[Equation]
These population statistics { µ, σ 2 } are usually calculated as the running average of mini-batch statistics over different training iteration t with an update factor α as follows:
[Equation]
The discrepancy of BN for normalization during training (using Eqn. 1) and inference (using Eqn. 2) can produce stochasticity, since the population statistics of BN are estimated from the mini-batch statistics that depend on the sampled mini-batch inputs.
This discrepancy is believed to beneﬁt the generalization [15, 11] if the stochasticity is well controlled.
However, this discrepancy usually harms the performance of small-batch-size training [40] since the estimation of population statistics can be inaccurate.
To address this problem, a bunch of batch-free normalizations are proposed that use consistent operations during training and inference, e.g. , Layer Normalization (LN) [1].
Basic Observations To analyze the failure of BN in NLP tasks, we ﬁrst plot the training loss and validation loss/BLEU [31] of BN and LN on IWSLT14 (De-En) dataset with the original Transformer model (see Figure 1).
We observe that the training of Transformer BN is faster than Transformer LN .
The training nll_loss of BN is even smaller than that of LN, especially at the beginning.
However, validation loss/BLEU of BN is worse than that of LN after around the seventh epoch.
This phenomenon can not be attributed to over-ﬁtting since BN introduces more stochasticity than LN in the training phase.
The inconsistency between training and inference of BN may play a role.
epoch
[Figure]
Mean Deviation of layers
IWSLT
CIFAR       epoch
[Figure]
IWSLT
CIFAR
[Figure]
[Figure]
[Figure]
Variance Deviation over BN layers
IWSLT
[Figure]
Since BN in ResNet18 also involves training inference inconsistency, we guess the degree of such inconsistency has a difference between ResNet18 and Transformer BN .
Therefore, we plot the deviation of batch statistics to population statistics of BN in ResNet18 and Transformer BN in
Figure 2 (top) to make a comparison. ResNet18 is trained on CIFAR-10 [17] and accuracy will droppercent if we replace BN with LN. We ﬁnd that at the end of the training, Transformer BN has a much bigger mean and variance deviation than ResNet18. Besides, the last several BN layers that are close to the output in Transformer BN have large variance deviation (Figure 2 (bottom)), which negatively impact the model output.
Furthermore, the performance degradation of Transformer BN coincides with the increase of variance deviation by comparing Figure 1 (right) and Figure 2 (bottom right).
Based on these observations, we hypothesize that the inconsistency between training and inference of
BN causes BN’s performance degradation in neural machine translation .
We ﬁrst mathematically deﬁne the training inference discrepancy of BN in the next subsection.
3.2 Training Inference Discrepancy
By observing Eqns.
1 and 2, the normalized output during training can be calculated as:
[Equation]
where σ B,j > 0 and σ j > 0 are the standard deviation for the j -th dimension.
We can see µ j − µ B,j σ j and σ j σ B,j can be viewed as random variables.
Their magnitude can characterize the diversity of mini-batch examples during training and indicate how hard the estimation of population statistics is.
We thus deﬁne the training inference discrepancy to quantitatively measure the inconsistency as follows.
Deﬁnition 1 (Training Inference Discrepancy (TID)) .
Let p B be the distribution of batch data.
Given a mini-batch data X sampled from p B , we deﬁne the TID of its mean and variance (with respect to model parameter θ ) as:
[Equation]
In terms of computing the TID in practice, we add a small positive constant in the denominator to avoid numerical instability.
We save the checkpoint at the end of each epoch and before training.
We ﬁrst estimate the population statistics by running forward propagation one epoch and then compute mean and variance TID by another epoch.
We omit θ when it can be inferred from context without confusion.
We compute the average mean and variance TID of all BN layers in ResNet18 trained on CIFAR10 and that of Transformer BN trained on IWSLT14 throughout training.
At the end of the training, the average mean/variance TID of BN in ResNet18 is approximately 0 .
9% , while that in Transformer is around 2 .
TID in
Transformer is much larger than that in ResNet18.
The trends are the same as basic observations in
Section 3.1.
We will use Equation (5) to compute TID in the subsequent analysis due to its better theoretical formulation (Equation (4)).
3.3 Comprehensive Validation
To further verify our hypothesis that large inconsistency between training and inference of BN causes BN’s degraded performance, we conduct experiments on Neural Machine Translation (NMT),
Language Modeling (LM), Named Entity Recognition (NER), and Text Classiﬁcation (TextCls) tasks.
We test both Post-Norm [39] and Pre-Norm [42] Transformers.
Experimental Setup We brieﬂy illustrate the experimental settings.
More detailed description can be found in supplementary materials.
For neural machine translation, we use IWSLT14 German-toEnglish (De-En) and WMT16 English-to-German (En-De) datasets, following the settings in Shen et al.
Our code is based on fairseq [30] 2 .
For language modeling, we conduct experiments on PTB [28] and WikiText-103 (WT103) [27].
We follow the experimental settings in Shen et al. [36], Ma et al.
For named entity recognition, we choose CoNLL2003 (English) [33] and Resume (Chinese) [50] datasets.
We mainly follow the experimental settings in Yan et al.
For text classiﬁcation, we select one small scale dataset (IMDB) [25] and three large scale datasets (Yelp,
DBPedia, Sogou News).
We use the code 3 and follow most conﬁgurations in Bhardwaj et al.
Performance Result We ﬁrst verify the inefﬁciency of BN compared to LN on four natural language tasks.
Results for Post-Norm and Pre-Norm Transformers are listed in Table 1.
BN performs much worse than LN on NMT, slightly worse on NER and TextCls tasks, but performs much better on LM.
Although BN performs worse in most cases, it has remarkable improvement over
LN on LM, raising the question: what contributes to the failure or success of BN?
Analyzing the Statistics of BN
We compute the TID of the last BN layer in Table 1 and leave the average TID of all BN layers in supplementary materials.
The last BN layer, which is close to the output, signiﬁcantly impacts the model prediction.
We observe that TID is highly correlated with the performance gap between BN and LN.
When TID is large, e.g., on WMT16, BN performs much worse than LN.
However, when the TID of BN is negligible, e.g., on PTB and WT103, BN performs better than LN with a large margin.
We select one dataset from each task with Pre-Norm Transformer and deﬁne the total TID as the sum of mean and variance TID.
At the end of the training, the total
TID of the last BN layer for WMT16/CoNLL/IMDB/WT103 is around 38%/16%/9%/5%, and the performance gap is -2.1 BLEU scores/-1.1 F
score/-0.
% accuracy/6.8 perplexity (PPL).
Larger TID tends to hurt BN’s performance.
To explore the quantitative relation between TID and performance gap, we substitute L = 3 ∼[-
-]
LN layers with BN layers from the bottom in the Post-Norm Transformer encoder on IWSLT14.
As
L increases, the variance TID of the last BN layer grows, and the BLEU scores of Transformer BN [Table]
93.3 Performance Gap -0.
0 Mean TID of BN last 3.4% 7.9% 1.6% 2.4% 9.6% 10.0% 2.9% 7.5% 3.9% 12.1% Var TID of BN last 4.6% 30.1% 1.7% 2.5% 6.5% 6.4% 6.2% 7.1% 3.3% 8.6%      10 TID (%) 1.
G a p ( B L E U ) Performance Gap vs TID      30 epochs V a r i a n ce
T I D ( % )
TID vs Gap Through Training 0.
va li d l o ss g a p drops off.
We plot the variance TID and BLEU gap between Transformer BN and Transformer LN in
Figure 3 (left).
We can see that the two quantities are highly correlated.
In Figure 3 (right), we plot the variance TID of the last BN layer and the validation loss gap between Transformer BN and Transformer LN on IWSLT14 through training.
The validation loss gap is calculated by subtracting loss of Transformer LN from Transformer BN .
At the beginning of training, BN performs better than LN.
When the TID begins to explode, BN’s performance starts to degrade.
Based on the results in Table 1 and observations in Figure 3, we argue that TID serves as an indicator of BN’s performance in Transformers.
Large TID hurts BN’s performance, while BN with small
TID performs better than LN due to its more efﬁcient optimization (see experimental validation in
Section 4.3).
Suppressing High TID by RBN
In this section, we are devoted to reducing the TID of BN when it is large.
If TID is suppressed, the performance of BN will be improved and may exceed LN due to the training efﬁciency of BN.
4.1 Regularized Batch Normalization
Assume there are H layers of BN in a neural network.
We denote the batch statistics and running statistics of each layer by µ iB , σ iB , and µ i , σ i , i = 1 . . .
, H .
Assume the Cross-Entropy (CE) loss with respect to the neural network parameters θ is denoted by L ( θ ) .
To avoid undesirable training [Table] inference discrepancy, we pose the optimization as a constrained problem:
[Equation]
where d µ and d σ measure the inconsistency of mean and variance.
This is equivalent to
[Equation]
To simplify the problem, we set λ i
= λ , ν i =
ν , for i
= 1 , . .
, H .
When handling batch data, we apply gradient-based optimization to the following loss ( L B ( θ ) is the batch CE loss):
[Equation]
In particular, we choose d µ ( µ B , µ ) = (cid:107) µ B − µ (cid:107) 22 and d σ ( σ B , σ ) = (cid:107) σ B − σ (cid:107) 22 .
The sensitivity analysis of hyperparameter is given in Section 4.3.
Since back propagating through the running statistics µ and σ would trace back to the ﬁrst batch of data which is impractical, we simply stop the gradient of µ and σ in back propagation.
4.2 Experimental Result for RBN
We choose λ, ν both from { 0 , 0 . 01 , 0 . 1 , 1 } by validation loss.
Results are shown in Table 2.
The optimal hyperparameters are listed in supplementary materials.
Neural Machine Translation
On IWSLT14 datasets, we see that RBN signiﬁcantly improves BN and can exceed LN with 0.1 BLEU scores with Pre-Norm Transformer and match LN with Post-Norm
Transformer.
On WMT16 dataset, although RBN still falls behind LN, it could improve 1.5/1.0 BLEU scores over BN in Post-Norm/Pre-Norm setting.
The reason is that even though RBN can suppress a large amount of TID, the remaining is still large since the original TID is huge.
We speculate that the high data diversity in WMT16 contributes to the explosive TID of BN, which is hard to remove.
We leave the veriﬁcation as future work.
Language Modeling On Post-Norm Transformer, BN could boost the testing PPL of LN from 53.2 to 45.9 on PTB and from 20.9 to 17.2 on WikiText-103.
Furthermore, substituting RBN for BN improves the testing PPL to 44.6 on PTB and 17.1 on WikiText-103.
On Pre-Norm Transformer, BN elevates the testing PPL of LN from 54.5 to 45.9 on PTB and from 24.6 to 17.8 on WikiText-103.
Moreover, replacing BN with RBN improves the testing PPL to 43.2 on PTB and 17.1 on WikiText103.
Overall, RBN exceeds LN with 8.6/3.8 testing PPL with Post-Norm Transformer and 11.3/7. testing PPL with Pre-Norm Transformer on PTB/WikiText-103.
Named Entity Recognition BN performs worse than LN on both Resume and CoNLL datasets, especially for Pre-Norm Transformer.
RBN improves BN in all settings, matches or exceeds
LN in three out of four settings.
By taking the better performance of Post-Norm and Pre-Norm, RBN matches the performance of LN on Resume and exceeds LN on CoNLL2003. [Table]
[Figure] training and leads to better performance. Text Classiﬁcation We ﬁnd that BN performs similar to/worse than LN on 4/4 settings.
RBN improves the performance of BN consistently and can match/exceed LN on 1/7 settings.
RBN improves BN with 0.3% accuracy on average, which shows the beneﬁt of our regularization.
We do not intend to achieve the state-of-the-art performance but to verify the efﬁcacy of RBN.
Comparison to BN’s Variants We compare our RBN with Power Normalization (PN) [36], Batch Renormalization (BRN) [14], and Moving Averaing Batch Normaliazation (MABN) [45] in Table 3. These methods incorporate population statistics of BN in training, which is beneﬁcial for alleviating training inference inconsistency of BN.
PN and MABN are implemented by their ofﬁcial codes 4 .
BRN is implemented according to their paper [14]. The conﬁgurations of PN, BRN, and MABN are given in supplementary materials. We highlight that PN incorporates layer scaling (LS) [49], which is important for stabilizing training, as shown in the supplementary materials and ofﬁcial code of PN.
We thus report the results for PN only and PN with layer scaling (PN+LS). We can see that RBN performs the best in most settings. PN and MABN is not stable without layer scaling, especially for Post-Norm Transformers.
4.3 Analysis
Training Inference Inconsistency We compute the TID of the last BN layer (BN last ) in Table 4 and plot the average TID of BN and RBN on WMT16, WT103, CoNLL2003, and IMDB datasets for Pre-Norm Transformers through training in Figure 4.
Figures of TID for other datasets and Post-Norm Transformer can be found in supplementary materials.
We can see that RBN reduces BN ’s mean and variance TID at the end of training.
On neural machine translation and named entity recognition tasks, the original TID is large.
RBN signiﬁcantly decreases the TID of [Table] [Figure] BN and improves BN’s performance by a clear margin.
For language modeling and text classiﬁcation tasks, RBN also reduces the moderate TID of BN and gets better PPL or accuracy.
[Figure]
Sensitivity to Hyperparameters We test different penalty coefﬁcients for RBN on neural machine translation with Pre-Norm Transformer.
The results are shown in Figure 5.
Penalizing the mean and variance discrepancy can both improve the performance of BN.
Combining them with moderate coefﬁcients achieves the best performance.
Training Dynamics To show the optimization advantages of RBN over LN, we explore the layer-wise training dynamics of LN and RBN in Pre-Norm Transformer on IWSLT14.
We refer the reader to Huang et al.
[12] for detailed analysis about the correlation between optimization of neural network and layer-wise training dynamics.
We empirically observe that replacing LN with RBN signiﬁcantly improves the layer-wise conditioning [12] of Transformer.
We denote the intermediate embedding in Transformer by ˜ X ∈ R B × T × d , each ˜ X i,j, :
∈ R
d is a word vector.
We reshape ˜ X to a sequence of word vectors to X =
[ x 1 , x 2 , . . .
, x BT ] ∈ R BT × d .
We assume BT > d which is satisﬁed in our experiments.
We deﬁne the general condition number with [Figure] respect to the percentage as C p ( X ) =
σ 1 σ (cid:100) pd (cid:101) , 0 < p ≤ 1 .
(cid:100) a (cid:101) is the smallest integer that is larger than or equal to a .
Lower C p ( X ) is usually associated with faster convergence of training.
We plot C 50% , and C 80% of input features of transformer encoder layer 2/4/6 in Figure 6.
We can see that RBN signiﬁcantly reduces the C 50% and C 80% of LN, usually with orders of magnitude.
We also plot the layer-wise C max ( X ) = λ max (( X T X ) 12 ) in Figure 7.
Smaller C max usually permits higher learning rates which leads to faster training and better generalization [10].
RBN has much smaller C max than LN.
Conclusion and Limitation
In this paper, we deﬁned Training Inference Discrepancy (TID) and showed that TID is a good indicator of BN’s performance for Transformers, supported by comprehensive experiments.
We observed BN performs much better than LN when TID is negligible and proposed Regularized BN (RBN) to alleviate TID when TID is large.
Our RBN has theoretical advantages in optimization and works empirically better by controlling the TID of BN when compared with LN.
We hope our work will facilitate a better understanding and application of BN in NLP.
Limitation.
Our analyses on TID are almost empirical studies without theoretical guarantee.
It is better to further model the geometric distribution of word embedding, evolving along with the training dynamics and information propagation, with theoretical derivation under mild assumptions.
Besides, our proposed RBN cannot entirely suppress huge TID in training large-scale datasets with high diversity, leading to degraded performance.
One possible direction is to combine RBN and LN for both better optimization properties and small TID, as explored in [13, 46] for CV tasks.
Acknowledgement Ji Wu was sponsored by National Key Research & Development Program of China (2021ZD0113402), Tsinghua University Spring Breeze Fund (2021Z99CFZ010), National Key Research & Development Program of China (Grant Number: 2021YFC2500803), and
Tsinghua-Toyota Joint Research Institute Inter-disciplinary Program.
Lei Huang was supported by
National Natural Science Foundation of China (Grant No. 62106012) and the Fundamental Research
Funds for the Central Universities.