# Alignment Attention by Matching Key and Query Distributions

arXiv: https://arxiv.org/abs/2110.12567 · HTML: https://mlsanity.com/paper/2110.12567/

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Builds on 62 of its 93 references that are in the map.
- First posted on arXiv 25 Oct 2021.

## Builds on (most cited first, 30 of 62)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [Sequence to Sequence Learning with Neural Networks](https://mlsanity.com/paper/1409.3215/) (2014)
- [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019)
- [ALBERT: A Lite BERT for Self-supervised Learning of Language Representations](https://mlsanity.com/paper/1909.11942/) (2019)
- [A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference](https://mlsanity.com/paper/1704.05426/) (2017)
- [Making the V in VQA Matter: Elevating the Role of Image Understanding in Visual Question Answering](https://mlsanity.com/paper/1612.00837/) (2016)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [ViLBERT: Pretraining Task-Agnostic Visiolinguistic Representations for Vision-and-Language Tasks](https://mlsanity.com/paper/1908.02265/) (2019)
- [Unsupervised Domain Adaptation by Backpropagation](https://mlsanity.com/paper/1409.7495/) (2014)
- [A large annotated corpus for learning natural language inference](https://mlsanity.com/paper/1508.05326/) (2015)
- [Know What You Don't Know: Unanswerable Questions for SQuAD](https://mlsanity.com/paper/1806.03822/) (2018)
- [Show, Attend and Tell: Neural Image Caption Generation with Visual Attention](https://mlsanity.com/paper/1502.03044/) (2015)
- [MixMatch: A Holistic Approach to Semi-Supervised Learning](https://mlsanity.com/paper/1905.02249/) (2019)
- [Adversarial Discriminative Domain Adaptation](https://mlsanity.com/paper/1702.05464/) (2017)
- [Can You Trust Your Model's Uncertainty? Evaluating Predictive Uncertainty Under Dataset Shift](https://mlsanity.com/paper/1906.02530/) (2019)
- [Energy and Policy Considerations for Deep Learning in NLP](https://mlsanity.com/paper/1906.02243/) (2019)
- [Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs)](https://mlsanity.com/paper/1511.07289/) (2015)
- [SemEval-2017 Task 1: Semantic Textual Similarity - Multilingual and Cross-lingual Focused Evaluation](https://mlsanity.com/paper/1708.00055/) (2017)
- [Revisiting Semi-Supervised Learning with Graph Embeddings](https://mlsanity.com/paper/1603.08861/) (2016)

## Built on by (most cited first, 14 of 14)

- [Preference-grounded Token-level Guidance for Language Model Fine-tuning](https://mlsanity.com/paper/2306.00398/) (2023)
- [MEDA: Dynamic KV Cache Allocation for Efficient Multimodal Long-Context Inference](https://mlsanity.com/paper/2502.17599/) (2025)
- [POUF: Prompt-oriented unsupervised fine-tuning for large pre-trained models](https://mlsanity.com/paper/2305.00350/) (2023)
- [Score Forgetting Distillation: A Swift, Data-Free Method for Machine Unlearning in Diffusion Models](https://mlsanity.com/paper/2409.11219/) (2024)
- [Regularizing a Model-based Policy Stationary Distribution to Stabilize Offline Reinforcement Learning](https://mlsanity.com/paper/2206.07166/) (2022)
- [A Unified Framework for Alternating Offline Model Training and Policy Learning](https://mlsanity.com/paper/2210.05922/) (2022)
- [ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans](https://mlsanity.com/paper/2502.07962/) (2025)
- [SemRoDe: Macro Adversarial Training to Learn Representations That are Robust to Word-Level Attacks](https://mlsanity.com/paper/2403.18423/) (2024)
- [A Prototype-Oriented Clustering for Domain Shift with Source Privacy](https://mlsanity.com/paper/2302.03807/) (2023)
- [Dual-Space Knowledge Distillation with Key-Query Matching for Large Language Models with Vocabulary Mismatch](https://mlsanity.com/paper/2603.22056/) (2026)
- [Switchable Decision: Dynamic Neural Generation Networks](https://mlsanity.com/paper/2405.04513/) (2024)
- [Q-Align: Alleviating Attention Leakage in Zero-Shot Appearance Transfer via Query-Query Alignment](https://mlsanity.com/paper/2508.21090/) (2025)
- [Indirect Attention: Turning Context Misalignment into a Feature](https://mlsanity.com/paper/2509.26015/) (2025)
- [QTrans: A Quantum Transformer for Sentiment Classification](https://mlsanity.com/paper/2609.12011/) (2026)

## Abstract (arXiv)

The neural attention mechanism has been incorporated into deep neural networks to achieve state-of-the-art performance in various domains. Most such models use multi-head self-attention which is appealing for the ability to attend to information from different perspectives. This paper introduces alignment attention that explicitly encourages self-attention to match the distributions of the key and query within each head. The resulting alignment attention networks can be optimized as an unsupervised regularization in the existing attention framework. It is simple to convert any models with self-attention, including pre-trained ones, to the proposed alignment attention. On a variety of language understanding tasks, we show the effectiveness of our method in accuracy, uncertainty estimation, generalization across domains, and robustness to adversarial attacks. We further demonstrate the general applicability of our approach on graph attention and visual question answering, showing the great potential of incorporating our alignment method into various attention-related tasks.