Transformers without Tears: Improving the Normalization of Self-Attention

Transformers without Tears: Improving the Normalization of Self-Attention Toan Q. Nguyen*,1 Julian Salazar*,2 1University of Notre Dame 2Amazon AWS AI tnguye28@nd.edu julsal@amazon.com Abstract We evaluate three simple, normalization-centric changes to improve Transformer training. First, we show that pre-norm residual connections (PRENORM) and smaller initializations enable warmup-free, validation-based training with large learning rates. Second, we propose ℓ2 normalization with a single scale parameter (SCALENORM) for faster training and better performance. Finally, we reaffirm the effectiveness of normalizing word embeddings to a fixed length (FIXNORM). On five low-resource translation pairs from TED Talks-based corpora, these changes always converge, giving an average +1.1 BLEU over state-of-the-art bilingual baselines and a new 32.8 BLEU on IWSLT '15 English-Vietnamese. We ob- serve sharper performance curves, more consistent gradient norms, and a linear relationship between activation scaling and decoder depth. Surprisingly, in the high-resource setting (WMT '14 English-German), SCALENORM and FIXNORM remain competitive but PRENORM degrades performance. 1. Introduction The Transformer [1] 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 [2, 3]. Our work focuses on layer normalization (LAYERNORM) [4], which we show has an outsized role in the convergence and performance of the Transformer in two ways: Placement of normalization. The original Transformer uses post-norm residual units (POSTNORM), where layer normal- ization occurs after the sublayer and residual addition. How- ever, [5] found that pre-norm residual units (PRENORM), where layer normalization occurs immediately before the sublayer, were instrumental to their model’s performance. [6] compares the two, showing that PRENORM makes back- propagation more efficient over depth and training Trans- formers with deep, 30-layer encoders. Our work demonstrates additional consequences in the base (≤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 *Equal contribution. 1Work was done during an internship at Amazon AWS AI. to past work on scaling NMT [7]. We also partly reclaim POSTNORM’s stability via smaller initializations, although PRENORM is less sensitive to this magnitude and can im- prove 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. [8] shows that batch normaliza- tion’s effectiveness is not from reducing internal covariate shift, but from smoothing the loss landscape. They achieve similar or better performance with non-variance-based nor- malizations in image classification. Hence, we propose re- placing LAYERNORM with the simpler scaled ℓ2 normal- ization (SCALENORM), which normalizes activation vectors to a single learned length g. This is both inspired by and synergistic with jointly fixing the word embedding lengths (FIXNORM) [9]. These changes improve the training speed and low-resource performance of the Transformer without affecting high-resource performance. On five low-resource pairs from the TED Talks [10] and IWSLT '15 [11] corpora, we first train state-of-the- art Transformer models (+4.0 BLEU on average over the best published NMT bitext-only …

Content truncated to 500 words