Alignment Attention by Matching Key and Query Distributions

Shujian Zhang, XINJIE FAN, Huangjie Zheng, Korawat Tanwisuth, Mingyuan Zhou

Keywords: Probabilistic methods, distribution matching, self-attention

Abstract: 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.

Code Of Conduct: I certify that all co-authors of this work have read and commit to adhering to the NeurIPS Statement on Ethics, Fairness, Inclusivity, and Code of Conduct.

TL;DR: Aligning the distributions of key and query to improve the learning of attention-based models

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2110.12567/code)