Alignment Attention by Matching Key and Query Distributions
Year: 2022

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.

heading: 1 Introduction
text: Attention-based mechanisms aggregate features with learnable weights to introduce useful inductive biases for sequence models [1, 2]. Since the introduction of the self-attention based Transformer [3], attention has become the foundation for many state-of-the-art models. Exploiting its computational efficiency and scalability, it has been used to train unprecedented large models on big datasets [4]. Large attention-based models have been demonstrating their ability to learn good representations in an unsupervised manner and benefit downstream analysis, with tremendous success in various natural language processing (NLP) [4–9], compute vision [10, 11], and multi-modal learning tasks [12, 13].
Attention networks, including multi-head attention, are being effectively utilized to capture the correlations between each pair of input tokens through individual or multiple attention functions. More specifically, in a self-attention layer withH heads, assuming that the output of the previous layer consists of n tokens, each of which is represented as a feature vector of dimension dmodel = d ·H , then each token feature vector will be transformed by a dmodel × d query projection matrix into a query feature vector, by a dmodel × d key projection matrix into a key feature vector, and by a dmodel × d value projection matrix into a value feature vector. The inner products of the ith query feature vector with all n key feature vectors are then fed through a softmax function to define the relative weights of the n keys to that query, which are used to aggregate the n value vectors into the vector representation of the ith word token in a head.
Although such networks are simple to optimize and intuitive to understand, how the key and query projection matrices should differ from each other has not been well-studied and understood. It is thus
The code is available at https://github.com/szhang42/alignment_attention
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
unclear whether they would result in well-controlled interactions between the keys and queries. In particular, ignoring the dependence between the n token feature vectors, we can view the output of the previous layer as an empirical distribution supported on n points in Rdmodel . In each head, this empirical distribution is transformed by the query and key projection matrices into a query empirical distribution and a key empirical distribution, respectively, each supported on n points at the same feature space in Rd. Since at each head two different projections matrices are used to project the input, the distributions of key and query will be different. Intuitively, if these two distributions are clearly misaligned with each other, then the query and key of a token, whose input feature resides in a region with lower probabilities, may have increased risk of being pushed further away from each other in the shared projection space.
This paper proposes alignment attention, which regularizes the query and key projection matrices at each self-attention layer, by matching the empirical distributions of the query and key feature vectors. We focus on within-head alignment between empirical distributions and present three different options for distribution matching. In our framework, alignment attention, built as an unsupervised approach to match the query and key distributions, is trained jointly to maximize a combination of data likelihood and distribution agreement. This efficient architecture design enables us to easily add alignment loss to convert existing self-attention networks, including pre-trained ones, into alignment attention. Meanwhile, it naturally shares parameters and computation with the self-attention networks, allowing an end-to-end training.
With a generic architecture, alignment attention can convert any existing soft attention models, including pre-trained ones, while maintaining the inherent advantages of conventional attention, such as efficiency and being simple to optimize. The proposed method boosts the performance while remaining efficient in memory and computation cost. Our experiments show that the proposed alignment attention method outperforms state-of-the-art self-attentions in a wide variety of settings, including natural language understanding tasks, graph attention network, and visual question answering, in terms of accuracy and uncertainty estimation. We further demonstrate that alignment attention achieves strong performance in domain generalization and adversarial robustness.

heading: 2 Alignment attention
text: We introduce a general recipe for alignment attention: (a) build the alignment to match the key and query distributions within each head, (b) develop efficient distribution matching methods, and (c) leverage existing attention structure and optimize the model in an end-to-end manner. The resulting architecture can be efficiently learned with existing self-attention networks.

heading: 2.1 Attention modules
text: Attention uses keys and queries to obtain soft attention weights W , which are then used to aggregate the values to obtain the output features. Consider n key-value pairs with a key matrix K ∈ Rn×dk , a value matrix V ∈ Rn×dv , andm queriesQ ∈ Rm×dk , where in general the dimensions of queries and keys are equal. The scaled product between key and query [3] is: Φ = fdot(Q,K) = QKT / √ dk ∈ Rm×n. Alternative choices include dot-product [3, 4] and additive attention [14–16]. Attention weights W is defined as the softmax output of Φ: W = softmax(Φ), where Wi,j =
exp(Φi,j)∑n j′=1 exp(Φi,j′ )
represents the importance of the jth key to the ith query learned by the neural networks.
The multi-head attention, first proposed in Transformer [3], projects the queries, keys, and values into H subspaces withH different learnable linear projections. These projections are performed in parallel and then concatenated into a single latent representation. At the lth self-attention layer, we can obtain attention weight W l,h = softmax(f(Ql,h,Kl,h)), where Ql,h = QlM l,hQ , K
l,h = KlM l,hK , and V l,h = V lM l,hV for h = 1, ...,H , with M denoting the parametric matrices to learn. The attention results from all heads are then concatenated into the layer output as Ol = [W l,1V l,1, ...,W l,HV l,H ].

heading: 2.2 Alignment attention
text: Self-attention allows the model to attend to the information from each representation subspace at each position [17]. To encourage different attention heads to indeed capture distinct features, most previous studies focus on the disagreement regularization to explicitly encourage the diversity
among attention heads [17]. By contrast, our focus is on exploiting how to make the key and query distributions better interact with each other in the latent space. We propose agreement matching to encourage the distributions of key and query over different tokens to be consistent within each head. Given a source input x and its output y, a neural attention model is trained to maximize the conditional probability of y given x over a training corpus. We introduce distribution alignment, an auxiliary regularization term in order to encourage the alignment between the learned key and query distributions. Considering a supervised learning problem with training data D := {xi,yi}Ni=1, the likelihood parameterized by θ is denoted by pθ(yi |xi). For notational convenience, below we drop the data index i. The whole model is differentiable to directly maximize the likelihood. Formally, the training objective with alignment attention is expressed as:
L(x,y) = log pθ(y | x)︸ ︷︷ ︸ likelihood +λ ∗ LAlign(Q,K)︸ ︷︷ ︸ alignment , (1)
where λ is the alignment weight [18–21]. The auxiliary regularization term LAlign guides the distribution matching between key (K) and query (Q).
The proposed alignment provides a sample-and-head-dependent matching between the key and query. Assuming Q and K have the same batch size and number of heads; Q is of dimension [B,H, n, dq], where B represents the batch size, H the number of heads, n the number of queries, and dq the hidden dimension within a head; and K is of dimension [B,H,m, dk]. Focused on the self-attention networks, we assume n = m = w and dq = dk = d in our alignment attention. We use Q,K to calculate the point-to-point difference from the query to key for each head at each sample, resulting in a tensor with dimension [B,H,w,w]. At each of the B samples of the minibatch and each of the H heads, the training objective is to minimize the expected difference between the empirical distributions of query and key, both of which are supported on a set of w query/key features in Rd (see Figure 2). This flexible alignment method could be conveniently deployed into a single-head or multi-head attention mechanism.
With the alignment attention, the resulting key and query distributions should be close to each other (see t-SNE plots [22, 23]in Figure 1). We also visualize the attention weight from both soft and alignment attentions in Figure 1. It is clear that while many semantically and sentimentally important words and their combinations, such as “where, they, sitting,” “color, ribbon,” “what, time, picture,” “boxes, room,” “where, bicycles, chained,” “bus, gaudy,” and “animals, visible,” are overlooked by vanilla soft attention, they are appropriately highlighted by the proposed alignment attention.

heading: 2.3 Alignment methods
text: To align the key and query distributions, we need a method that can quantify the difference between two distributions given their empirical samples. Under this requirement, we consider three different distribution matching methods, including the discriminator-based adversarial training [24], which is directly related to the Jensen–Shannon (JS) divergence [25], the Wasserstein distance in its primal form, which can be defined by solving an optimal transport problem [26, 27], and bi-directional conditional transport [28], which is developed by exploiting both the chain rule and Bayes’ rule to quantify the difference between two probability distributions.

heading: 2.3.1 Adversarial training-based alignment
text: Adversarial training, the key part to enable GANs [24], has been successfully exploited to minimize the distributional discrepancy [29, 30]. Under a minimax two-player framework, the discriminator D is trained to distinguish the distributions of key and query, while both the query and key projection matrices, MQ and MK , are treated as the generator G and trained to confuse D. Here we consider a discriminator D that is shared across the H heads for alignment. Denote the empirical distributions of query and key as pQ and pK , respectively. For a sample of w tokens and at one of the H heads, we have pQ = ∑w i=1 1 w δqi and pK = ∑w j=1 1 w δkj , where qi and ki are projected from the ith token’s input feature via the query and key projection matrices at that head, respectively. Thus the alignment loss is the sum over H different head-dependent loss, each of which on a sample can be expressed as (we drop the head index for notational simplicity)
min G max D LAlign−GAN(Q,K) := Eq∼pQ [logD(q)] + Ek∼pK [log(1−D(k))]. (2)
In summary, the discriminator is trained to maximize the alignment loss, and the generator, i.e., the query and key projection matrices, is trained to minimize the alignment loss, so that the key and query distributions are learned adversarially to align with each other.
Discriminator-based modules. We leverage the highway architecture [31] to construct the discriminator. With a feature map X ∈ {Q,K}, the highway network first computes a transform gate τ , as τ = σ (Φτ (X)) , where Φτ is a fully-connected layer and σ is the sigmoid activation function. Then, the output of the highway network is, IX = τ ReLU(Φh(X)) + (1− τ) X, where Φh is another linear layer followed by ReLU activation and denotes the element-wise multiplication. We further apply a two-layer MLP to obtain the classifier probability, D(X) = σ(Φ2 (FNL (Φ1 (IX))))), where Φ1 and Φ2 are fully-connected layers connected by FNL, a leaky ReLU activation function. To optimize the discriminator-based modules, instead of alternately updating the adversaries, like in GAN [24], we use the gradient-reversal layer [32] to jointly optimize all the components.

heading: 2.3.2 Optimal transport-based alignment
text: An alternative to the adversarial training-based alignment is to consider optimal transport (OT) [26], which has been widely used for distribution matching. In OT, the transport plan Π(pQ, pK) is the set of all possible joint distributions of query and key, whose marginals over the query and key are pK and pQ, respectively. We define the OT-based alignment cost as
LAlign−OT(Q,K) = min π∈Π(pQ,pK) Eq,k∼π [c(q, k)] , (3)
where c(·, ·) is the transport cost between two points. For discrete pQ and pK , fining the optimal transport plan often involves solving a computationally expensive linear programming problem. It is also known in the literature that OT is sensitive to the outliers in pQ and pK due to the two marginal constraints [33]. To reduce the computation burden, we consider an entropy-regularized OT, which allows the OT problem to be solved iteratively using the Sinkhorn iterations [34].
Optimal transport-based modules. We use the Sinkhorn algorithm [34, 35] to estimate the transport plan, defined as π = argmin
π
∑w i=1 ∑w j=1(c(qi, kj) ·π(qi, kj)− π(qi, kj) log π(qi, kj)), where
= 0.01. The OT-alignment cost now becomes LAlign−OT = ∑w i=1 ∑w j=1(c(qi, kj) · π (qi, kj)).

heading: 2.3.3 Bidirectional conditional transport-based alignment
text: Instead of requiring π ∈ Π(pQ, pK) as in OT, we follow Zheng and Zhou [28] to consider probabilistic bidirectional conditional transport (CT), which exploits both the chain rule and Bayes’ theorem to constrain the joint π with both pQ = ∑w i=1 1 w δqi and pK = ∑w j=1 1 w δkj in two different ways. First, we consider a query-to-key CT that define π as π(qi, kj) = pQ(qi)πK(kj | qi), where πK(kj | qi) is a conditional distribution defined as πK(kj | qi) = pK(kj) exp(τφ(kj)
T τφ(qi))∑w j′=1 pK(kj′ ) exp(τφ(kj′ ) T τφ(qi)) and τφ(·) is
a neural network based transformation parameterized by φ. Second, we consider a key-to-query CT that defines π(qi, kj) = pK(kj)πQ(qi | kj), where πQ(qi | kj) = pQ(qi) exp(τφ(kj)
T τφ(qi))∑w i′=1 pQ(qi′ ) exp(τφ(kj) T τφ(qi′ )) .
Third, we define a point-to-point cost as cη(q, k) = 1 − τη(k) T τη(q)
||τη(k)||2||τη(q)||2 , where τη(·) is a neural network based “critic” function whose parameter η will be adversarially learned. Combing them together leads to the bidirectional CT-based alignment loss as
LAlign−CT(Q,K) = 12Eq∼pQ, k∼πK(· | q)[cη(q, k)] + 1 2Ek∼pK , q∼πQ(· | k)[cη(q, k)]. (4)
Compared to OT, bidirectional CT is able to efficiently model the alignment with less computation. The structure of alignment attention with transport-based methods is presented in Fig. 2.
CT-based modules. The critic τη(·) is structured similarly as the discriminator in Section 2.3.1. It projects the input data onto a vector space instead of outputting logits for binary classification. For τφ(·), we use a two-layer MLP network. We optimize the query and key projection matrices and φ to minimize LAlign−CT(Q,K) in (4), and optimize η to maximize it. The gradient-reversal layer [32] is also used to optimize the critic adversarially.

heading: 3 Related work
text: Alignment Learning. Liang et al. [36] first assign agreement terms for jointly training word alignment in phrase-based statistic machine translation. The general bidirectional sequence alignment models with model inevitability regularization are then proposed by Levinboim et al. [37]. Recently, the alignment is studied in the multi-head attention model based on the Transformer architecture, where alignment extraction is improved by augmenting an additional alignment head to the multi-head source-to-target attention component [38]. Our proposed alignment attention adopts the alignment idea and matches the distributions of query and key. This general and efficient framework gives us the flexibility to better model attention weights. The domain generalization ability and adversarial robustness of alignment attention are also studied.
Distribution Matching. Distribution matching is a fundamental problem in statistics and machine learning [39]. The widely used distances include the Kullback–Leibler (KL) divergence [40] and Jensen–Shannon (JS) divergence [25]. GANs [24] are proposed with the adversarial-based objectives. Due to the inherent advantages of allowing the two distributions to have non-overlapping supports [41, 42], the Wasserstein distance from the optimal transport problem is also used for defining the transport cost [26, 27]. In our alignment attention, we consider discriminator-based and transportbased methods as alignment methods. Based on the alignment methods, we apply alignment attention to the widely used attention models and leverage the existing efficient attention architecture to build entire alignment attention network.

heading: 4 Experiments
text: Our method can be incorporated into any self-attention based models. To exam its effectiveness and general applicability, we apply alignment attention to a diverse set of tasks, including language understanding, graph attention, and visual question answering. Furthermore, the model’s generalization
across domains and robustness towards adversarial attacks are studied on language tasks. We study a variety of state-of-the-art models for these tasks including ALBERT [5], BERT [4], and RoBERTa [6]. Below we present the main experimental settings and results, with more training and hyperparameter details provided in Appendix B.

heading: 4.1 Alignment attention in natural language understanding
text: Since the self-attention based Transformer was proposed, it has been widely used in pretrained models for NLP and related areas, achieving state-of-the-art results on various downstream tasks. However, training a state-of-the-art pretrained model now requires substantial computational resources which demands considerable energy, along with the associated financial and environmental costs. Vaswani et al. [3] report that a Transformer-base model was trained on 8 Nvidia P100 GPUs for 12 hours and Strubell et al. [43] report a BERT-base model was trained on 64 V100 GPUs for 79 hours. Thus for alignment attention, we utilize it to finetune these pretrained models on large corpora, which is not only computationally and financially friendly but also accessible to more researchers.

heading: 4.1.1 In-domain language understanding evaluation
text: We first evaluate the alignment attention for in-domain language tasks where the training and testing data are from the same domain. We conduct experiments on eight benchmark datasets from General Language Understanding Evaluation (GLUE) [44] and two Stanford Question Answering Datasets (SQuAD) [45, 46]. Our experiments are based on the state-of-the-art pretrained model, ALBERT [5], a memory-efficient version of BERT [4] with parameter sharing and embedding factorization. Based on Huggingface PyTorch Transformer [47], our implementation uses the base version of ALBERT following the same setting from Lan et al. [5].
Results. In Table 1, we present the soft attention and the alignment attention (AA) with GAN, optimal transport, and CT, resuming from the same checkpoints. The mean accuracies are reported with 5 independent runs (see full results with the error bars in Table 8 in the Appendix). Alignment attention outperforms the soft attention, which indicates that matching the distributions of key and query gives better performance than soft attention and the results are not sensitive to the choice of alignment methods. Due to the expensive computation for distance measure at each step of OT, we will focus on GAN and CT as alignment methods in the following experiments. Overall, alignment attention improves the soft attention in both GLUE and SQuAD datasets even by only using alignment attention at the finetuning stage.

heading: 4.1.2 Generalization across domains
text: In real applications, it is very common to deploy a neural network model into a new domain with data unseen during training. The model’s generalization has been extensively studied in the machine learning community. Significant past work has studied cross-domain robustness using sentiment analysis [48–50]. The recent work from Desai and Durrett [51] has explicitly elected tasks where outof-domain performance is substantially lower and challenging domain shifts are exhibited. Following the setting in Desai and Durrett [51], we test the generalization ability of our alignment attention. For our in-domain and out-of-domain datasets, we split the development set in half to obtain a held-out, non-blind test set. We conduct experiments on three tasks: (1) Natural Language Inference. The Stanford Natural Language Inference (SNLI) corpus is a large-scale entailment dataset [52] as the in-domain data. Multi-Genre Natural Language Inference (MNLI) [53] can be used as unseen out-of-domain test dataset. (2) Paraphrase Detection. Quora Question Pairs (QQP) is used as in-domain data which includes sentence pairs from Quora that are semantically equivalent [54]. TwitterPPDB (TPPDB) [55] is considered as out-of-domain data. (3) Commonsense Reasoning. Situations With Adversarial Generations (SWAG) is a grounded commonsense reasoning task [56]. The out-of-domain data is HellaSWAG (HSWAG), which is a more challenging benchmark [56]. We report accuracy and expected calibration error (ECE) for both in-domain(ID) and out-of-domain(OD).
ECE is calculated as a weighted average of the difference between each bin’s accuracy and confidence: ECE:= ∑ i Gi N |acc(Gi)− conf(Gi)|, where Gi, acc(Gi), and conf(Gi) are the count, accuracy, and confidence of samples in the i’th group, respectively. The number of groups is 10 as in [51].
Results. In Table 2, we include open-source implementations of Decomposable Attention (DA) [57] and Enhanced Sequential Inference Model (ESIM) [58] as baselines. For pretrained models, we use BERT-base-uncased [4] and RoBERTa-base [6] from HuggingFace Transformers [47]. We incorporate the alignment attention in both pretrained models. Alignment attention consistently outperforms the corresponding soft attention not only for in-domain, confirming our results in Section 4.1.1, but also for out-of-domain. The relevantly larger gains on the out-of-domain setting indicate that alignment attention has the better generalization ability across domains. The improved results of ECE demonstrate the better-calibrated model for uncertainty estimation with alignment attention.

heading: 4.1.3 Robustness towards adversarial attacks
text: Machine learning models recently have been found vulnerable to adversarial examples that are legitimate input altered by small and often imperceptible perturbations [59]. Therefore, it becomes increasingly important to exam the model’s robustness against adversarial attacks. Our alignment attention imposes a regularization to ensure well-aligned key and query distributions so that it is expected to become more robust to the generated perturbation that would fool the model. We follow the same settings from Section 4.1.1 to test the robustness of finetuned ALBERT-base models with soft attention or alignment attention. We utilize the TextAttack [60] framework and apply three state-of-the-art untargeted black-box adversarial attacks (1) Textfooler [61]: counter-fitted word embedding swap; (2) Textbugger [62]: character-level insertion, deletion, swap, and substitution; (3) BAE [63]: generating BERT masked token prediction. 1000 adversarial attacks are conducted for each model with a maximum sentence length of 512. We report the percentages of failed adversarial attacks in Table 3. Higher percentages indicate more robust models.
Results. The alignment attention shows consistent improvements over the soft attention across all three attacks and achieves significant gains on the average failure rates. The alignment attention
demonstrates its robustness and gets along with our intuition that the alignment attention can learn better and more robust key and query distribution due to the use of distributional matching.

heading: 4.2 Alignment attention in graph neural networks
text: To test the general applicability of our alignment attention, we also experiment our method with graph attention networks (GAT) [64], where the graph structure is injected into the attention masks in which nodes are able to attend over their neighborhoods’ features in the graph. Leveraging masked self-attentional layers, GAT processes the node features for the node classification.
Experimental Setup. Following the setting in GAT [64], we conduct experiments on three standard citation network benchmark datasets— Cora, Citeseer and Pubmed [65]— in a transductive setting, indicating all nodes from training and test are on the same graph [66]. The details of three datasets and experimental settings are deferred to Appendix B.
Results. In Table 4, we report the mean classification accuracies on test nodes over 5 random runs, and the standard deviations of alignment attention. We experiment with both AA-GAN and AA-CT. Table 4 shows that alignment attention consistently improves upon the corresponding baseline models across all three datasets, which further confirms the efficient structure of this alignment attention. The CT-based alignment attention performs better than GAN-based alignment attention.

heading: 4.3 Attention in visual question answering
text: Visual question answering (VQA) [67] is a multi-modal learning task where the model predicts the answer given a question conditioning on the image. Self-attention architectures in MCAN [68] has been recently proposed to learn the fine-grained semantic meaning of both the image and question. We adapt the proposed alignment attention to MCAN and compare with soft attention. We conduct experiments on the VQA-v2 dataset [67] and follow the hyperparameters and other settings from Yu et al. [68]. In addition, to investigate the model’s robustness to noise, we construct the noisy dataset by incorporating the Gaussian noise (mean 0, variance 1) to image features [69, 70]. Four-layer encoder-decoder based MCAN is used as the baseline model with the soft self-attention. For each experiment, we report the accuracy on both the original data and the noisy data. As in Fan et al. [70] and Zhang et al. [71], we also report the Patch Accuracy vs Patch Uncertainty (PAvPU) [70, 72] as a measure of the uncertainty estimation where the p-value threshold is set to be 0.05 and the number of attention weight samples is 20. Please refer more detailed experimental settings in Appendix B.
Results. The results are summarized in Table 5. We report the accuracy and uncertainty of different attentions on both original and noisy data. In terms of accuracy, alignment attention shows consistent improvements over the soft attention on both original and noisy data. These results verify our conjecture that the alignment attention is more robust to the noise which aligns with our results on adversarial robustness in Section 4.1.3. For uncertainty, we observe that on both original and noisy data, alignment attention has better uncertainty estimations, meaning that alignment attention in general is more certain on its correct predictions and more uncertain on its mistakes.
Results Analysis. Visualizations. We plot the attention weights of both alignment attention and soft attention on one question for VQA. In Figure 3, attention weight represents the average importance of each query and key pair. For example, in (b), when the row is ‘Did’ and the column is ‘hit’, the color represents the average attention weight from the query ‘Did’ to the key ‘hit’. We observe that the alignment attention gives relatively sharper attention weights compared to the soft attention and therefore gives good prediction accuracy and uncertainty estimation.
Parameter Size and Running time. In Table 6, we provide the parameter sizes and step time for different attention types combined with MCAN where the attention module constructs the main model. It shows that alignment attention (AA) keeps the parameter size at the very similar level as soft attention while moderately increasing the step time compared to soft attention.
Ablation Study. We conduct ablation study with AA+CT to exam the role of the alignment-weight hyperparameter λ in Equation 1 by turning it from 0.01 to 1. We find that the experimental results are not sensitive to the choice of the value of the λ. Any number from 0.01 to 1 would give similar results. In all experiments considered in the paper, which cover various noise levels and model sizes, we have simply fixed it as 0.01. Please see detailed results in Table 9 in Appendix.

heading: 5 Conclusion
text: Our proposed alignment attention aims to match the key and query distributions. We leverage different alignment methods with a generic and efficient architecture design which requires surprisingly few modifications to standard soft attention and enables us to easily convert existing soft attention models, including pretrained ones, to alignment attention. Our experiments on a variety of language understanding tasks show that alignment attention achieves strong performance in accuracy, uncertainty estimation, domain generalization, and adversarial robustness even by only adding the alignment loss during the finetuning stage. In the real-life scenarios, the attention models have been deployed in many machine learning systems, such as self-driving [73] and healthcare [74]. However, the data from the real practice is biased and long-tailed. Therefore, we see opportunities of our method that can mitigate the risks with uncertainty estimation. Further, on graph node classification and visual question answering, alignment attention demonstrates its general applicability and effectiveness of
each component of the proposed structure, showing the great potential to be added as a plug-and-play component to many existing attention models.

heading: Acknowledgements
text: The authors acknowledge the support of Grant IIS-1812699 from the U.S. National Science Foundation, the APX 2019 project sponsored by the Office of the Vice President for Research at The University of Texas at Austin, the support of a gift fund from ByteDance Inc., and the Texas Advanced Computing Center (TACC) for providing HPC resources that have contributed to the research results reported within this paper.

Review summaries:
This paper aims to improve the performances of current widely-used attention models by explicitly aligning the query and key distributions. To this end, the authors propose an alignment attention model, which can be easily incorporated into existing self-attention based models, including the pre-training models, e.g., BERT, to boost their performances while remaining efficient in memory and computation cost. To prove the effectiveness of the proposed approach, the authors conduct extensive experimental studies on natural language understanding tasks, graph attention network, and visual question answering task. The domain generalization ability and adversarial robustness of alignment attention are also studied.  > After Rebuttal:  Thanks the authors for the response. I'd like to keep my initial score.

This work aims at solving the alignment problem between the query and key distribution within the self-attention module. It is based on the hypothesis that the key and query following identical distribution can help to improve the feature learning for the attention module.  Correspondingly, an alignment regularization is proposed to explicitly encourage self-attention to match the key and query within each head. Three kinds of alignment methods, including Adversarial training-based, Optimal transport-based, and, Bidirectional conditional transport-based, all of which can outperforms the baselines. Comprehensive experiments are performed on a variety of problem settings, which proves that the introduced alignment module can bring consistent improvements. Besides, the proposed method can work as a plug-and-play module to be added to many existing attention modules. 

The paper reviews the attention mechanism commonly used in deep neural networks. It proposes the alignment attention mechanism that match the query and key distributions in self-attention within each head. The paper regards the alignment attention as an unsupervised regularization in current attention. Three alignment methods have been introduced to explore the effectiveness of distribution matching. Extensive experiments are done to demonstrate the effectiveness and general applicability of alignment attention in domains of language understanding tasks, and even graph node classification and visual question answering. The alignment attention shows great potential of being a plug-and-play component to current attention models.