Alignment Attention by Matching Key and Query Distributions
This work addresses a bottleneck in attention-based models for researchers and practitioners in NLP and related fields, offering an incremental improvement through a simple regularization method.
The paper tackled the problem of improving self-attention mechanisms in neural networks by introducing alignment attention, which matches key and query distributions within each head, resulting in enhanced performance on language tasks, generalization, and robustness.
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.