Source-linked AI summary
Query-Key Normalization for Transformers
Alex Henry, Prudhvi Raj Dachapally, Shubham Pawar, Yuxuan Chen
TL;DR
Low-resource translation needs effective methods for adapting Transformers, particularly where softmax attention can become arbitrarily saturated. The paper introduces QKNORM, which normalizes queries and keys along each attention head and applies learnable scaling, improving performance across five low-resource translation pairs. The authors conclude that its exact mechanism of benefit remains to be determined.
Problem
Low-resource language translation is challenging, while unbounded attention dot products can make small relative differences silence other signals.
Method
QKNORM applies ℓ2 normalization to queries and keys along the head dimension, then replaces their dot product with learnably scaled cosine similarity.
Results
QKNORM improved performance on each of five low-resource translation pairs, with an average test BLEU increase of 0.928.
Takeaways & Limitations
QKNORM modifies Transformer attention and demonstrated utility for low-resource bilingual translation on five language pairs.
Takeaways & Limitations
The exact locations and reasons for QKNORM’s benefits require further research.
Abstract
from arXiv · showhide
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 $\ell_2$ 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.
1 Introduction
Low-resource translation remains challenging, motivating normalization changes that help Transformer attention avoid overly concentrated softmax behavior. QKNORM replaces attention’s dot product with learnably scaled cosine similarity and improves low-resource bilingual translation.
- Motivation: Low-resource language translation presents special challenges despite the Transformer’s established role in machine translation.Recent adaptations include multilingual transfer learning, reduced model depth, and target-language language-model regularization.
- Contribution: QKNORM introduces normalization intended to keep softmax attention inputs within an appropriate range.The technique targets low-resource language pairs.
- Attention mechanism: Unbounded dot products can let tiny relative differences silence other attention signals, potentially limiting the complexity of learnable attention patterns.The softmax then determines distributions over sequence elements, whose weighted values form updated token representations.
- Contribution: QKNORM replaces the dot product inside softmax attention with cosine similarity scaled by a learnable parameter.The authors conjecture this binds QK^T to a narrower range, facilitating more diffuse attention patterns where useful.
2 Background
QKNORM extends normalization-based Transformer adaptations for low-resource translation by normalizing queries and keys after head splitting while retaining layer normalization. Across the same five benchmarks, it surpasses the prior combined normalization model by an average of 0.928 test BLEU.
- Prior work: The prior model combined FIXNORM, PRENORM, and SCALENORM and reported state-of-the-art bilingual performance on five low-resource translation pairs.The benchmarks came from the TED Talks and IWSLT’15 corpora.
- Prior normalization: FIXNORM sets word embeddings to unit length, while PRENORM moves layer normalization before each Transformer sublayer.PRENORM is described as enhancing stability by preserving the residual path as an identity map.
- Prior normalization: SCALENORM applies ℓ2 normalization along the embedding dimension to Q, K, and V before multihead attention splits inputs into heads.It multiplies the normalized representations by a learnable scalar initialized using 1/√d.
- Results: 0.928 test BLEU is the average improvement over the prior model across the same five translation pairs.The authors report that QKNORM was combined with FIXNORM, PRENORM, and LAYERNORM.
- QKNORM: QKNORM applies ℓ2 normalization only to Q and K along the head dimension after multihead attention splits its input.The head dimension corresponds to the embedding dimension within each separate head.
- QKNORM: QKNORM complements LAYERNORM rather than replacing it and, like SCALENORM, uses a learnable scaling parameter.The distinction is both the normalized components and the location of normalization relative to head splitting.
3 Dot Products and the Softmax Function
Unnormalized dot products can saturate softmax attention, producing concentrated weights and potentially limiting learnable attention patterns. QKNORM instead bounds query-key similarities, uses a learnable scale, and is associated with more diffuse attention and higher BLEU in the illustrated experiment.
- Softmax saturation: Unbounded dot products can make small relative score differences silence other attention signals, potentially limiting the complexity of patterns attention heads learn.The paper conjectures that this occurs when softmax approaches saturation.
- Softmax saturation: Figure 1 contrasts scaled dot product attention with concentrated self-attention heatmaps across four encoder heads.The visualization is presented as consistent with softmax saturation from unnormalized dot products.
- Query-key normalization: Figure 2 shows the same four heads with more diffuse attention patterns under QKNORM than under scaled dot product attention.The comparison is qualitative rather than systematic.
- Empirical illustration: 24.33 BLEU versus 19.4 BLEU was obtained by QKNORM versus the baseline on IWSLT 2016 German-to-English.Both models were trained for 10 epochs using the Annotated Transformer implementation.
- Query-key normalization: QKNORM replaces the dot product with bounded cosine similarities and scales them using a learnable parameter rather than dividing by √d.Cosine similarities lie in [−1, 1], while the learnable initialization depends on the 97.5th-percentile training sequence length L.
4 Experiments and Results
The experiments evaluate QKNORM on five low-resource translation pairs using established preprocessing and BLEU-scoring procedures. QKNORM improves test performance across the benchmark set, with an average increase of 0.928 BLEU.
- Experimental setup: The experiments use the same five low-resource translation pairs as Nguyen and Salazar (2019), spanning TED Talks and IWSLT’15.The pairs include four TED Talks directions and English-to-Vietnamese from IWSLT’15.
- Evaluation: BLEU is scored with multi-bleu.perl for en→vi and multi-bleu-detok.perl for the four TED Talks pairs because tokenization differs.The scoring choice reflects the available preprocessing needed for comparable evaluation.
- Evaluation: SacreBLEU rescoring leaves detokenized BLEU results essentially unchanged except for en→vi, where the reported baseline and model scores differ slightly.The best reproducible en→vi baseline scores were 32.48 with Moses and 32.41 with SacreBLEU; the paper reports 32.79 from Nguyen and Salazar (2019) in Table 2.
- Experimental setup: The study uses fastBPE and the same Moses bootstrap-resampling code as Nguyen and Salazar (2019).
- Model settings: The best results use 8,000 steps of linear warmup despite earlier findings that PRENORM can make warmup less important.
- Results: 0.928 BLEU is the average test-set improvement from incorporating QKNORM and using layer normalization instead of SCALENORM.The comparison covers all five translation pairs.
5 Conclusion
The paper introduces QKNORM as an attention normalization technique and demonstrates its utility for low-resource bilingual translation. It improves performance on all five language pairs, while the mechanism of its effectiveness remains open for further study.
- QKNORM modifies the Transformer attention mechanism and is demonstrated in an existing implementation on five low-resource language pairs.
- QKNORM improves performance on each of the five pairs, with an average test BLEU increase of 0.928.
- The paper identifies possible explanations for QKNORM’s effectiveness but leaves exactly where and why it helps for future research.
A Varying the Number of Heads
QKNORM remains stable on the en→vi test set as the number of attention heads varies, including configurations with small head dimensions.
- QKNORM performance remains stable when the number of attention heads reaches 32, corresponding to a head dimension of 16.
B Equation 3
Equation 3 sets the scale for QKNORM through an empirically selected rule of thumb that initializes a learnable parameter. The parameter value is chosen using en→vi experiments across percentile settings.
- Equation 3 is an empirically derived scaling rule that initializes a learnable parameter rather than fixing the scale permanently.
- The best value of L is selected by running en→vi translation experiments with different percentile values.
C Ablation Experiments
Table 6 ablates specific QKNORM components on en→vi, showing that the learnable scaling factor g is especially important. Without scaling, normalized query-key products remain within [−1, 1], limiting softmax differences.
- The ablation experiments evaluate how specific QKNORM components affect en→vi test performance.
- The biggest performance drop occurs when omitting g, the learnable scaling factor.
- Without scaling, ˆQ ˆK^T values remain within [−1, 1], leaving softmax to operate on bounded differences.