Query-Key Normalization for Transformers (2020)

Abstract: Low-resource language translation is a challenging but socially valuable NLP task. Building on recent work adapting the Transformer’s normalization to this setting, we propose QKNorm, a normalization technique that modifies the attention mechanism to make the softmax function less prone to arbitrary saturation without sacrificing expressivity. Specifically, we apply l2-normalization along the head dimension of each query and key matrix prior to multiplying them and then scale up by a learnable parameter instead of dividing by the square root of the embedding dimension. We show improvements averaging 0.928 BLEU over state-of-the-art bilingual benchmarks for 5 low-resource translation pairs from the TED Talks corpus and IWSLT’15.

## introduction the transformer @xcite remains the architecture of choice for machine translation. since its introduction, various architectural and functional modifications have been made to improve its performance on nmt datasets @xcite @xcite @xcite . translating low-resource languages presents special challenges. recent strategies for adapting transformers to this socially valuable task include exploiting transfer learning with many-to-many multilingual models @xcite , reducing model depth @xcite , and adding a regularization penalty for diverging from the predictions of a monolingual language model pretrained on the target language @xcite . this paper builds on recent work on layer normalization for 1 code to reproduce our experiments is available at https: //github.com/cyndxai/qknorm low-resource language pairs, introducing a normalization technique that tries to keep the input to softmax attention within an appropriate range. layer normalization. for transformers and other nlp models, layer normalization @xcite yields significantly better performance than batch normalization @xcite , in part because nlp models tend to exhibit greater variance in batch statistics during training, for example compared to computer vision @xcite . layer normalization boosts performance in deeper networks chiefly by controlling their gradients @xcite . it re-scales and re-centers activation distributions (though re-centering may be unnecessary, see @xcite . the type of normalization used and the placement of that normalization within the transformer are both crucial to transformer performance @xcite . ## dot products and the softmax function since the dot product is unbounded, differences between elements that may be insignificantly small on a relative basis can silence all other signals in the attention weights applied to v . we conjecture that this limits the complexity of the patterns that attention heads can learn. the impact is more obvious in less sophisticated transformer implementations (perhaps in part because subsequent advances have mitigated the same figure 1 : scaled dot product attention. self-attention heatmaps for 4 heads from one encoder layer displaying more "concentrated" attention, consistent with the conjecture that unnormalized dot products in qk t saturate the softmax and limit the attention patterns that can be learned. issue in different ways). figures 1 and 2 show a heatmap comparison of encoder weights trained using the code for the annotated transformer 2 , the first with scaled dot product attention and the second with qknorm. the models containing these encoders were trained for 10 epochs on iwslt 2016 @xmath0, with the baseline model scoring 19.4 bleu and the qknorm model scoring 24.33 bleu on the test set, computed with the sacrebleu python package @xcite . though this heatmap comparison is obviously not systematic, we think the visual at least provides a plausible intuition for the incremental gain this technique achieves, with scaled dot product attention exhibiting the kind of "winner-take-all" behavior we would expect from a softmax near saturation. in comparison to dot products, cosine similarities are bounded by [-1, 1] which creates the opposite problem as input to softmax -the differences 2 https://nlp.seas.harvard.edu/2018/04/ 03/attention.html between values are too small for softmax to let the model effectively ignore connections between words it should not attend to. instead of dividing by √ d as in scaled dot product attention we scale up using a learnable parameter that we initialize with a value that depends on the length of the sequences in the training data (and hence on the number of elements in qk t ): ## experiments and results we follow the implementation in the repository for @xcite , both in replicating their performance and as a starting point for our version (and also for computing bleu as reported in ## conclusion in this paper, we introduced a normalization technique that modifies the attention mechanism in transformers and demonstrated its utility for lowresource bilingual translation by building it into an existing transformer implementation with stateof-the-art performance on 5 low-resource language pairs. qknorm improves performance for each of the 5 pairs, with an average test bleu increase of 0.928. we pointed to possible explanations for its effectiveness but identifying exactly where it helps and why requires further research. first, we plan to combine our approach with the fairseq transformer implementation @xcite and apply it to the flores dataset @xcite , investigating the effect of qknorm on the optimal depth, number of attention heads, and warmup schedule for low-resource translation, in combination with recent advances like bpe-dropout @xcite . next, we plan to look at high-resource settings to see whether the benefits of query-key normalization dissipate with access to more training data. lastly, we intend to study how qknorm impacts what attention heads actually learn, adapting methods from bert attention studies such as @xcite .