Query-Key Normalization for Transformers

Findings of the Association for Computational Linguistics: EMNLP 2020, pages 4246–4253 November 16 - 20, 2020. c⃝2020 Association for Computational Linguistics 4246 Query-Key Normalization for Transformers Alex Henry, Prudhvi Raj Dachapally, Shubham Pawar, Yuxuan Chen Cyndx Technologies {alex.henry,prudhvi.dachapally,shubham.pawar,ethan.chen}@cyndx.com Abstract Low-resource language translation is a chal- lenging but socially valuable NLP task. Build- ing 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 satu- ration without sacrificing expressivity. Specif- ically, we apply ℓ2 normalization along the head dimension of each query and key ma- trix prior to multiplying them and then scale up by a learnable parameter instead of di- viding by the square root of the embedding dimension. We show improvements averag- ing 0.928 BLEU over state-of-the-art bilin- gual benchmarks for 5 low-resource transla- tion pairs from the TED Talks corpus and IWSLT’15.1 1 Introduction The Transformer (Vaswani et al., 2017) remains the architecture of choice for machine translation. Since its introduction, various architectural and functional modifications have been made to im- prove its performance on NMT datasets (Ahmed et al., 2017; Zhang et al., 2018; Wang et al., 2019; Dai et al., 2019; Zhao et al., 2019). 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 (Aharoni et al., 2019), reducing model depth (van Biljon et al., 2020), and adding a regularization penalty for diverging from the predictions of a monolingual language model pretrained on the tar- get language (Baziotis et al., 2020). This paper builds on recent work on layer normalization for 1Code to reproduce our experiments is available at https: //github.com/CyndxAI/QKNorm low-resource language pairs, introducing a normal- ization technique that tries to keep the input to softmax attention within an appropriate range. Layer normalization. For Transformers and other NLP models, layer normalization (Ba et al., 2016) yields significantly better performance than batch normalization (Ioffe and Szegedy, 2015), in part because NLP models tend to exhibit greater variance in batch statistics during training, for ex- ample compared to computer vision (Shen et al., 2020). Layer normalization boosts performance in deeper networks chiefly by controlling their gradi- ents (Xu et al., 2019). It re-scales and re-centers activation distributions (though re-centering may be unnecessary, see Zhang and Sennrich 2019). The type of normalization used and the placement of that normalization within the Transformer are both crucial to Transformer performance (Nguyen and Salazar, 2019). Softmax attention. Given a matrix X embed- ding a sequence of tokens, attention transforms each embedding into a mixture of itself and other elements of the sequence according to the impor- tance of their connections for the modeling task at hand. In the case of multihead self-attention, the vectors of X are projected linearly into Query, Key and Value matrices. The operation softmax(QKT ) (1) defines a distribution for each token over all the oth- ers in its sequence that sums …

Content truncated to 500 words