Alignment Attention by Matching Key and Query Distributions

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.

arXiv commentNeurIPS 2021; Our code is publicly available at https://github.com/szhang42/alignment_attention

References in corpus (17)

• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• Sequence to Sequence Learning with Neural Networks
• Decoupled Weight Decay Regularization
• Show, Attend and Tell: Neural Image Caption Generation with Visual Attention
• Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs)
• XLNet: Generalized Autoregressive Pretraining for Language Understanding
• Training Very Deep Networks
• Enhanced LSTM for Natural Language Inference
• ALBERT: A Lite BERT for Self-supervised Learning of Language Representations
• MixMatch: A Holistic Approach to Semi-Supervised Learning
• The Cramer Distance as a Solution to Biased Wasserstein Gradients
• Evaluating Bayesian Deep Learning Methods for Semantic Segmentation
• Interpolating between Optimal Transport and MMD using Sinkhorn Divergences
• Latent Alignment and Variational Attention
• Contextual Dropout: An Efficient Sample-Dependent Dropout Module
• Sawtooth Factorial Topic Embeddings Guided Gamma Belief Network
• TopicNet: Semantic Graph-Guided Topic Discovery