Transformers without Tears: Improving the Normalization of Self-Attention

Introduction

The Transformer Vaswani et al. (2017) has become the dominant architecture for neural machine translation (NMT) due to its train-time parallelism and strong downstream performance. Various modifications have been proposed to improve the efficiency of its multi-head attention and feedforward sublayers Guo et al. (2019); Sukhbaatar et al. (2019). Our work focuses on layer normalization (LayerNorm) Ba et al. (2015), which we show has an outsized role in the convergence and performance of the Transformer in two ways:

The original Transformer uses post-norm residual units (PostNorm), where layer normalization occurs after the sublayer and residual addition. However, Chen et al. (2018) found that pre-norm residual units (PreNorm), where layer normalization occurs immediately before the sublayer, were instrumental to their model’s performance. Wang et al. (2019) compare the two, showing that PreNorm makes backpropagation more efficient over depth and training Transformers with deep, 30-layer encoders.

Our work demonstrates additional consequences in the base ($\leq$6-layer encoder) Transformer regime. We show that PreNorm enables warmup-free, validation-based training with large learning rates even for small batches, in contrast to past work on scaling NMT Ott et al. (2018). We also partly reclaim PostNorm’s stability via smaller initializations, although PreNorm is less sensitive to this magnitude and can improve performance. However, despite PreNorm’s recent adoption in many NMT frameworks, we find it degrades base Transformer performance on WMT '14 English-German.

Choice of normalization.

On five low-resource pairs from the TED Talks Qi et al. (2018) and IWSLT '15 Cettolo et al. (2015) corpora, we first train state-of-the-art Transformer models (+4.0 BLEU on average over the best published NMT bitext-only numbers). We then apply PreNorm, FixNorm, and ScaleNorm for an average total improvement of +1.1 BLEU, where each addition contributes at least +0.3 BLEU (Section 3), and attain a new 32.8 BLEU on IWSLT '15 English-Vietnamese. We validate our intuitions in Section 4 by showing sharper performance curves (i.e., improvements occur at earlier epochs) and more consistent gradient norms. We also examine the per-sublayer $g$’s learned by ScaleNorm, which suggest future study.

Background

The original Transformer applies LayerNorm after the sublayer and residual addition (PostNorm):

We conjecture this has caused past convergence failures Popel and Bojar (2018); Shazeer and Stern (2018), with LayerNorms in the residual path acting similarly to $\lambda_{i}\neq 1$; furthermore, warmup was needed to let LayerNorm safely adjust scale during early parts of training. Inspired by He et al. (2016b), we apply LayerNorm immediately before each sublayer (PreNorm):

This is cited as a stabilizer for Transformer training Chen et al. (2018); Wang et al. (2019) and is already implemented in popular toolkits Vaswani et al. (2018); Ott et al. (2019); Hieber et al. (2018), though not necessarily used by their default recipes. Wang et al. (2019) make a similar argument to motivate the success of PreNorm in training very deep Transformers. Note that one must append an additional normalization after both encoder and decoder so their outputs are appropriately scaled. We compare PostNorm and PreNorm throughout Section 3.

2 Weight initialization

Our experiments with this default initializer find that PostNorm sometimes fails to converge, especially in our low-resource setting, even with a large number of warmup steps. One explanation is that Xavier normal yields initial weights that are too large. In implementations of the Transformer, one scales the word embeddings by a large value (e.g., $\sqrt{d}\approx 22.6$ for $d=512$), giving vectors with an expected square norm of $d$. LayerNorm’s unit scale at initialization preserves this same effect. Since feedforward layers already have their weights initialized to a smaller standard deviation, i.e., $\sqrt{\frac{2}{d+4d}}$, we propose reducing the attention layers’ initializations from $\sqrt{\frac{2}{d+d}}$ to $\sqrt{\frac{2}{d+4d}}$ as well (SmallInit), as a corresponding mitigation. We evaluate the effect of this on PostNorm vs. PreNorm in Section 3.1.

This can be viewed as projecting $d$-dimensional vectors onto a $(d-1)$-dimensional hypersphere with learned radius $g$. This expresses the inductive bias that each sublayer’s activations has an ideal “global scale,” a notion we empirically validate in Section 4.2. ScaleNorm replaces the $2d$ scale and shift parameters of LayerNorm with a single learned scalar, improving computational and parameter efficiency while potentially regularizing the loss landscape.

This bias has an explicit interpretation at the final layer: large inner products sharpen the output distribution, causing frequent words to disproportionately dominate rare words. This led Nguyen and Chiang (2018) to introduce $\textsc{FixNorm}({\bm{w}})=g\frac{{\bm{w}}}{\lvert\lvert{\bm{w}}\rvert\rvert}$ with fixed $g$ at the last linear layer, to maximize the angular difference of output representations and aid rare word translation. By making $g$ learnable, we can apply ScaleNorm and FixNorm jointly, which means applying the following at the final linear layer:

Note that this combination at the last layer is equivalent to cosine normalization Luo et al. (2018) with a learned scale.

4 Learning rates

Despite using an adaptive optimizer, Adam Kingma and Ba (2015), Transformer training uses a learning rate (LR) schedule with a linear warmup and an inverse square root decay (InvSqrtDecay):

where $d$ is the hidden dimension of the self-attention layers, and $\lambda$, $n_{\text{warmup}}$ are hyperparameters that determine the highest learning rate achieved and the number of steps to reach it, respectively. These two hyperparameters have been the subject of much empirical study Popel and Bojar (2018); Ott et al. (2018). In light of our modifications however, we revisit various aspects of this schedule:

We conjectured that warmup is primarily needed when using PostNorm to gradually learn LayerNorm parameters without gradient explosion/vanishing (Section 2.1). Hence, we evaluate both PreNorm and PostNorm without warmup in Section 3.3.

Large learning rates.

To speed up training, one often explores using larger learning rates. In the context of Transformer, Ott et al. (2018) and Aharoni et al. (2019) take $\lambda\in\{2,3\}$ instead of the conventional $\lambda=1$. Ott et al. (2018) showed that one can scale up Adam’s learning rate to $10^{-3}$ with an extremely large batch (400k tokens). However, the improved convergence provided by our modifications could enable higher learning rates with much small batch sizes (4k tokens), as examined in Section 3.3.

Validation-based decay.

For similar reasons, one might wish to adopt a classic validation-based decay, i.e., training at a high learning rate for as long as tenable, decaying rapidly when development scores flatline. This has inspired usage of fixed decay schemes upon convergence with InvSqrtDecay Dong et al. (2018); Salazar et al. (2019). We revisit ValDecay under our modifications, where we still perform a linear warmup but then multiply by a scale $\alpha_{\text{decay}}<1$ when performance on a development set does not improve over $patience$ evaluations.

Experiments and results

We train Transformer models for a diverse set of five low-resource translation pairs from the TED Talks Qi et al. (2018) and the IWSLT '15 Cettolo et al. (2015) corpora. Details are summarized in Table 1. For more information motivating our choice of pairs and for exact training details, refer to Appendix A.

To see the impact of weight initialization, we run training on the en→vi dataset using warmup steps of 4k, 8k, 16k (Table 2). With default initialization, PostNorm fails to converge on this dataset even with a long warmup of 16k steps, only reaching 5.76 BLEU.

The second row shows that taking a smaller standard deviation on the attention weights (SmallInit) restores convergence to PostNorm. Though the $\sqrt{2/5}\approx 0.63$ adjustment used here seems marginal, operations like residual connections and the products between queries and keys can compound differences in scale. Though both models now achieve similar performance, we note that PreNorm works in all setups, suggesting greater stability during training. For all remaining experiments, we use PostNorm and PreNorm with SmallInit. We find this choice does not affect the performance of PreNorm.

Table 3 shows our results along with some published baselines. First, note that our Transformer baselines with PostNorm + LayerNorm (1) are very strong non-multilingual NMT models on these pairs. They outperform the best published numbers, which are all Transformer models in the past year, by an average margin of +4.0 BLEU. Then, we see that PreNorm (2) achieves comparable or slightly better results than PostNorm on all tasks. FixNorm (3) gives an additional gain, especially on ar→en ($p<0.01$).

Finally, we replace LayerNorm with ScaleNorm (4). ScaleNorm significantly improves on LayerNorm for two very low-resource pairs, gl→en and sk→en. On the other tasks, it performs comparably to LayerNorm. Upon aggregating all changes, our final model with ScaleNorm and FixNorm improves over our strong baseline with PostNorm on all tasks by an average of +1.1 BLEU ($p<0.01$), with each change contributing an average of at least +0.3 BLEU. In Section 4.2 and Appendix B, we further examine where the performance gains of ScaleNorm come from.

Moreover, ScaleNorm is also faster than LayerNorm. Recall that for each vector of size $d$, LayerNorm needs to compute mean, standard deviation, scaling, and shifting, which costs $O(7d)$ operations. For ScaleNorm, we only need $O(3d)$ operations to perform normalization and global scaling. This does not account for further gains due to reduction in parameters. In our implementation, training with ScaleNorm is around 5% faster than with LayerNorm, similar to the speedups on NMT observed by Zhang and Sennrich (2019)’s RMSNorm (which can be viewed as ScaleNorm with per-unit scales; see Section 4.2).

3 Learning rates

We compare the original learning rate schedule in equation 7 (InvSqrtDecay) with validation-based decay (ValDecay), possibly with no warmup (NoWarmup). We use $\lambda=1$, $n_{warmup}=8\text{k}$ for InvSqrtDecay and ValDecay. For NoWarmup, we instead use a learning rate of $3\cdot 10^{-4}$ for all datasets. For both ValDecay and NoWarmup, we take $\alpha_{decay}=0.8$ and $patience=3$. For experiments with high learning rate, we use either ValDecay or InvSqrtDecay with $\lambda=2$ (giving a peak learning rate of $\approx 10^{-3}$). All experiments use PreNorm + FixNorm + ScaleNorm.

In Table 4, we see that NoWarmup performs comparably to InvSqrtDecay and ValDecay except on gl→en. We believe that in general, one can do without warmup, though it remains useful in the lowest resource settings. In our 2$\times$LR experiments, we can still attain a maximum learning rate of $10^{-3}$ without disproportionately overfitting to small datasets like gl→en.

One might hypothesize that ValDecay converges more quickly to better minima than InvSqrtDecay by staying at high learning rates for longer. However, both schedulers achieve similar results with or without doubling the learning rate. This may be due to the tail-end behavior of ValDecay methods, which can involve multiplicative decays in rapid succession. Finally, our 2$\times$LR experiments, while not yielding better performance, show that PreNorm allows us to train the Transformer with a very high learning rate despite small batches (4k tokens).

Since PreNorm can train without warmup, we wonder if PostNorm can do the same. We run experiments on en→vi with NoWarmup, varying the number of encoder/decoder layers. As seen in Table 5, PostNorm often fails without warmup even with 5 or 6 layers. Even at 4 layers, one achieves a subpar result compared to PreNorm. This reaffirms Section 3.1 in showing that PreNorm is more stable than PostNorm under different settings.

4 High-resource setting

Since all preceding experiments were in low-resource settings, we examine if our claims hold in a high-resource setting. We train the Transformer base model on WMT '14 English-German using fairseq and report tokenized BLEU scores on newstest2014. Implementation of our methods in fairseq can be found in Appendix C.

In Table 6, ScaleNorm and FixNorm achieve equal or better results than LayerNorm. Since ScaleNorm is also faster, we recommend using both as drop-in replacements for LayerNorm in all settings. Surprisingly, in this task PostNorm works notably better than PreNorm; one observes similar behavior in Wang et al. (2019). We speculate this is related to identity residual networks acting like shallow ensembles Veit et al. (2016) and thus undermining the learning of the longest path; further study is required.

Analysis

Figure 1 shows that PreNorm not only learns faster than PostNorm, but also outperforms it throughout training. Adding FixNorm also gives faster learning at first, but only achieves close performance to that with PreNorm and no FixNorm. However, once paired with ScaleNorm, we attain a better BLEU score at the end. Because of the slow warmup period, ScaleNorm with warmup learns slower than ScaleNorm without warmup initially; however, they all converge at about the same rate.

To visualize how PreNorm helps backpropagation, we plot the global gradient norms from our runs in Figure 2. PostNorm produces noisy gradients with many sharp spikes, even towards the end of training. On the other hand, PreNorm has fewer noisy gradients with smaller sizes, even without warmup. LayerNorm has lower global norms than ScaleNorm + FixNorm but it has more gradient components corresponding to normalization.

2 Activation scaling and the role of g𝑔g

One motivation for ScaleNorm was that it expressed a good inductive bias for the global scaling of activations, independent of distributional stability (Section 2.3). In contrast, a contemporaneous work Zhang and Sennrich (2019) proposes root mean square layer normalization (RMSNorm), which still follows layer normalization’s motivation but reduces overhead by forgoing additive adjustments, using only a scaling $g_{i}$ per activation $a_{i}$. Despite their differing motives, tying the $g_{i}$ of RMSNorm and dividing by $\sqrt{d}$ retrieves ScaleNorm.

Hence we can frame our comparisons in terms of number of learnable parameters. We rerun our PreNorm experiments with RMSNorm. We also consider fixing $g=\sqrt{d}$ for ScaleNorm, where only FixNorm has learnable $g$. Table 7 shows that ScaleNorm always performs comparably or better than RMSNorm. Surprisingly, the fixed-$g$ model performs comparably to the one with learnable $g$. However, at higher learning rates (ValDecay with and without 2$\times$LR), fixed-$g$ models perform much worse on ar→en, en→he and en→vi. We conjecture that learning $g$ is required to accommodate layer gradients.

In Figure 3, we plot the learned $g$ values for pairs with 100k+ examples. For all but the decoder-encoder sublayers, we observe a positive correlation between depth and $g$, giving credence to ScaleNorm’s inductive bias of global scaling. This trend is clearest in the decoder, where $g$ linearly scales up to the output layer, perhaps in tandem with the discriminativeness of the hidden representations Liang et al. (2018). We also note a negative correlation between the number of training examples and the magnitude of $g$ for attention sublayers, which may reflect overfitting.

Finally, to affirm our intuition for interpreting $g$, we plot $g$ values with and without label smoothing (Figure 4). We see a difference in later layers of the decoder; there, removing label smoothing results in lower $g$ values except at the output layer, where $g$ increases sharply. This corresponds to the known overconfidence of translation models’ logits, on which label smoothing has a downscaling effect Müller et al. (2019).

Conclusion

Acknowledgements

The authors would like to thank David Chiang and Katrin Kirchhoff for their support of this research.

References

Appendix A Training details

The pairs are English (en) to Hebrew (he), Vietnamese (vi), and Galician (gl), Slovak (sk), Arabic (ar) to English (en). Because the data is already preprocessed, we only apply BPE Sennrich et al. (2016b) with fastBPEhttps://github.com/glample/fastBPE. Depending on the data size, we use different numbers of BPE operations.

We wanted to compare with the latest low-resource works of Neubig and Hu (2018); Aharoni et al. (2019) on the TED Talks corpus Qi et al. (2018). In particular, Aharoni et al. (2019) identified 4 very low-resource pairs ($<$70k); we took the two (gl→en, sk→en) that were not extremely low ($\leq$6k). They then identified 4 low-resource pairs with 100k-300k examples; we took the top two (ar→en, en→he). To introduce a second English-source pair and to showcase on a well-understood task, we used the en→vi pair from IWSLT '15 with an in-between number of examples (133k). In this way, we have examples of different resource levels, language families, writing directions, and English-source versus -target.

Model configuration.

We set the hidden dimension of the feedforward sublayer to 2048 and the rest to 512, matching Vaswani et al. (2017). We use the same dropout rate for output of sublayers, ReLU, and attention weights. Additionally, we also do word dropout Sennrich et al. (2016a) with probability 0.1. However, instead of zeroing the word embeddings, we randomly replace tokens with UNK. For all experiments, we use label smoothing of 0.1 Szegedy et al. (2016); Pereyra et al. (2017). The source and target’s input and output embeddings are shared Press and Wolf (2017), but we mask out words that are not in the target’s vocabulary at the final output layer before softmax, by setting their logits to $-\infty$.

Training.

We use a batch size of 4096 and optimize using Adam Kingma and Ba (2015) with the default parameters $\beta_{1}=0.9$, $\beta_{2}=0.999$, $\epsilon=10^{-8}$. Gradients are clipped when global norm exceeds 1.0 Pascanu et al. (2013). An epoch is a predefined number of iterations for each pair. We stop training when a maximum number of epochs has been met or the learning rate becomes too small ($10^{-6}$). We also do early stopping when the development BLEU has not improved for 20 evaluations. For gl→en, this number is 50. When doing validation-based decay, we use $\alpha_{decay}=0.8$ and $patience=3$. For complete data and model statistics, please refer to Table 1. The best checkpoint is selected based on the development BLEU score during training.

Evaluation.

We report tokenized BLEU Papineni et al. (2002) with multi-bleu.perl to be comparable with previous works. We also measure statistical significance using bootstrap resampling Koehn (2004). For WMT '14 English-German, note that one needs to put compounds in ATAT formathttps://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/utils/get_ende_bleu.sh before calculating BLEU score to be comparable with previous works.

Appendix B Further analysis

We ask if improvements from ScaleNorm on our low-resource tasks are due to improved regularization (a smaller generalization gap) or improved overall performance. We record smoothed train and test perplexities of our PreNorm models in Table 8. We see suggestive results but no conclusive trends. For ar→en, gl→en, and sk→en, train and test drop slightly, with test more so than train. For en→vi, train perplexity increases and test perplexity decreases an equivalent amount. For en→he, our smallest change between ScaleNorm and LayerNorm, train perplexity negligibly increased and test perplexity remains the same.

Appendix C Listings

fairseq.

We follow fairseq’s tutorialhttps://github.com/pytorch/fairseq/blob/master/examples/scaling_nmt/README.md and train a PostNorm Transformer base model using the following configuration:

For ScaleNorm, we replace all LayerNorms in fairseq/models/transformer.py and fairseq/modules/transformer_layer.py with ScaleNorm (implemented above). For FixNorm, we change the word embedding initialization to uniform with range $[-0.01,0.01]$ and normalize with torch.nn.functional.normalize.

We note that fairseq uses Xavier uniform initialization, which is big compared to our SmallInit (Section 3.1). We conjecture that fairseq training remains stable thanks to its large batch size, which gives more stable gradients.