Alignment Attention by Matching Key and Query Distributions
Authors: Shujian Zhang
XINJIE FAN
Huangjie Zheng
Korawat Tanwisuth
Mingyuan Zhou
Venue: NeurIPS 2021 Poster (2021)
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.

Alignment Attention by Matching Key and Query Distributions
Shujian Zhang
Xinjie Fan
Huangjie Zheng Korawat Tanwisuth Mingyuan Zhou
The University of Texas at Austin
(szhang19, xfan, huangjie zheng, korawat. tanvisuth) edu
mingyuan... www.@mccombs...oooomoomo wwwpe/ccmmmm...ee.....eeeee.ee.. wtw.as.exdeeeee edu
Abstract
The neural attention mechanism has been incorporated into deep neural networks achieve state-of-the-art performance various domains. Most such models use multi-head self-attention which is appealing for the ability to attend to informa- tion 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 unsupervised regularization in the existing attention framework. It is simple convert any models with ceff-attention, including pre-trained ones, to the proposed alignment attention. On variety of language understanding tasks, 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-relate tasks.
Introduction
Attention- baaeti mechanisms aggregate features with learnable weights to introduce useful inductive biases for sequence models [ 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, has been used to train unprecedented large models on big datasets [ Large attention-bases 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 compute vision [O and multi-modal learning tasks [12]
Attention networks, including multi-head attention, are being effectively utilized to capture the correlations between each pair of input tokens through individual multiple attention functions. More specifically, in self-attention layer with H heads, assuming that the output of the previous layer consists of n tokens, each of which is represented then each token feature vector will be transformed by a dymodet x d query projection matrix into query feature vector, by dmode! d key projection matrix into key feature vector, and by dmodel x d value projection matrix into a value feature vector. The inner products the ith query feature vector with all key feature vectors are then fed through softmax function 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 head.
feature vector of dimension dyoodst d H.
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 wtt.ppes//gggchhb..c////saaag//a
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

unclear whether they would result in well-controlled interactions between the keys and queries. In particular, ignoring the dependence between the token feature vectors, we can view the output of the previous layer as an empirical distribution supported on points in Romodde In each head, this empirical distribution is transformed by the query and key projection matrices into a query empirical distribution and 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 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 whth 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-en tnd training.
With a generic architecture, aligmment 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 answer- ing, in terms of accuracy and uncertainty estimation. We further demonstrate that alignment attention achieves strong performance in domain generalization and adversarial robustness.
2 Alignment attention
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 an end-to-end manner. The resulting architecture can be efficiently learned with existing self-attention networks.
2.1 Attention modules
Attention uses keys and queries 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 E IR "xd,, a value matrix V axds and m queries Q E Rxd., where in general the dimensions of queries and
keys are equal. The scaled product between key and query
3 is: tb = saa(Q, K) = QK ddz E
Rmxa. Alternative choices include dot-product
3 and additive attention 14.16). Attention
weights W is defined as the softmax output of Φ W softmax ()), where Wiy = >y exp((...)
exp(It.))
represents the importance of the jth key to the ith query learned by the neural networks.
The multi-head attention, first proposed in Transformer 33 projects the queries, keys, and values into Il subspaces with H different learnable linear projections. These projections are performed in parallel and then concatenated into a single latent representation. At the /th self-attention layer, we can obtain attention weight Wr. (f K²(A)), where Qi. = Q'M²'' Kt,h K M/A. and V' ML:" 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 O [W41y4] ...., W yIH
VLA =
2.2 Alignment attention
Self-attention allows the model to attend to the information from each representation subspace at each position [7]. To encourage different attention heads to indeed capture distinct features, most previous studies focus the disagreement regularization explicitly encourage the diversity
2

---

Where are shttrsstting
Where are
What color mibii
Whet coloris
What Erno the picture
theie any bosess in the rooms
Aew thire wrry Bicenni (000l are these Blsvsiied
Where these bicycles chaired
Dooalhe he look gaudy
the Bus looks BINKEE
How mans
Algureet
Figure 1: T-SNE visualizations of (a) soft attention and (b) alignment attention with bidirectional conditional transport. In (a) and (b), visualize the both key distribution and the query distribution. The blue dots and orange diamonds represent the features from the key distribution and these from the query distribution, respectively. We show attention heatmaps of VQA data (c) and (d), includes overall attention of seven examples' matrix embedding by summing the attention weight vectors, with darker color denoting higher attention probability.
among attention heads 17. By contrast, our focus is 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 head. Given a source input x and its output y. a neural attention model is trained to maximize the conditional probability of y given over training corpus. We introduce distribution alignment, an auxiliary regularization term in order to encourage the alignment between the learned key and query distributions. Considering supervised learning problem with training data D == {ar, y: J. the likelihood parameterized by 0 is denoted by po yo a..). For notational convenience, below we drop the data index i. The whole model is differentiable directly maximize the likelihood. Formally, the training objective with alignment attention is expressed as:
different tokens to be consistent within each
[ y)
log PD x) +d & Calign (Q, K),
(1)
likelihood
aligmment
where is the alignment weight 18 The auxiliary regularization term CAliga guides the distribution matching between key (K) and query (Q).
The proposed alignment provides a
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, dg), where B represents the batch size, H the number of heads, n the number of queries, and d, the hidden dimension within head; and K is of dimension [B, H, m, dæ). Focused on the self-attention networks, we assume n = m = w and dg = ds = 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 Id (see Figure 2). This flexible alignment method could be conveniently deployed into 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 []. We also visualize the attention weight from both soft and alignment attentions in Figure I 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.
2.3 Alignment methods
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 255, the Wasserstein distance in its primal form, which can be defined by solving 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.
3

---

Pre-Ammmmeet
One layer of alignment altendion
Kep
Single Head
Quaay
Luna
Query
fax
softnax
Post-Allgnment
Key
Suy
XQ
GRL
Critc
CRL
Single Head
OuRy
Query
Algement
Lass
Say
kay
QUTY
On the left, we visualize the alignment attention. The ellipse represents query and the rectangle represents the key. On the right, a demonstration of the difference and similarity between vanilla soft attention and our alignment attention. Alignment attention (in green) shares the same architecture as soft attention before obtaining key, query, and value. Then alignment attention adds the alignment structure to perform distribution matching, where GRL represents the gradient reversal layer.
Figure 2:
2.3.1 Adversarial maining-based alignment
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, Me and Mr, are treated as the generator G and trained 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 Рk. respectively. For a sample of wo tokens and at one of the H heads, we have po = Ei= 5 and pk Σ Oky, where qi and k; are projected from the ith token's 's input feature via the query and key projection matrices at that head, respectively. Thus the aligmment loss is the sum over H different head- wead-deepnneett loss, each of which on sample can be expressed as (we drop the head index for notational simplicity)
min max Lalign GAN (Q, K) := Egyroo Dlg D(4) + Ex-px [log(1 D(k))) G D
(2)
In summary, the discriminator is trained to maximize the alignment loss, and the generator, i.c., 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-basedd modules. We leverage the highway architecture 31 to construct the dis- criminator. With feature map X E (Q. K), the highway network first computes a transform gate T, as 7 (0, (X)), where 2 a fully-connected layer and is the sigmoid activation function. Then, the output of the highway network is, Ix = T c ReLU (X)) )) + (1 -) O X, where 4,, is another linear layer followed by ReLU activation and @ denotes the element-wise multiplication. We further apply two-layer MLP to obtain the classifier probability. D(X) = a ( (FxL (4, (Ix))))), where P1 and 2 are fully- culnecco layers connected by FNL. leaky ReLU activation function. To optimize the discriminator-bast based modules, instead of alternately updating the ad adversaries, like in GAN
24
we use the wradient-reversal layer [32 to jointly optimize all the components.
2.3.2 Optimal transport-basee aligmment
An alternative to the adversarial training-bass alignment is to consider optimal transport (OT) 26). which has been widely used for distribution matching. In OT, the transport plan II(pq, Pk) is the set of all possible joint distributions of query and key, whose marginals over the query and key are Pи and pQ. respectively. We define the OT-based alignment cost as
CAligu-OT (Q,
min rEll(pQ.Pk)
Ey (e(g,
(3)
where -).,-) is the transport cost between two points. For discrete PQ and Pk. fining the optimal transport plan often involves solving computationally expensive linear programming problem. also known in the literature that OT is sensitive to the outliers in PQ and px due to the two marginal constraints [33] To reduce the computation burden, we consider an entropy-regularizet 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 trans-
port plan, defined n. argmin E E k2)) (9:, k2) - K2) log x (q,, ((,,,))), where
0.01. The OT-alignment cost becomes LAlign-OT Σ 2- Σ (c(qi, ky) me ki)).
2.3.3 Bidirectional conditional transport-based alignment
Instead of requiring [p pk) as in OT, we follow Zheng and Zhou 28 consider probabilistic bidirectional conditional transport (CT), which exploits both the chain rule and Bayes' theorem to constrain the joint и with both PQ = 2i=1 w
S and px- = -j==1 w 00, in two different ways. First, TT (ea R3))
(Q(())TT (k; q;),, where Tx (k3 (q) Pae (k;) (xx (x and Tx(-) is is conditional distribution defined n (ký (q:) Σy'=1 Px (Agr)
consider a query-to-key CT that define
neural network based transformation parameterized by Φ. Second, we consider key-to-query CT that defines (a; ky) = Pk ( ky); where m2: 1ky) = po(g)) exp(ralk, ((
2 polar) exp(Ta (a(a)))
Third, we define a point-to- -oint cost as (g, = -1 (k) (0(()
(R) [,,,)1/ ? where T.(()) is neural network based "critic" function whose parameter 77 will be adversarially learned. Combing them together leads to the bidirectional CT-based alignment loss as
-Align-CrQQ. = 2 (- (1))) (g, + En-px wmaa (1) (1) k)). CT
(4)
Compared to OT, bidirectional CT is able to efficiently model the aligmment with less computation. The structure of alignment attention with transport-badd based methods is presented in Fig.
CT-based modules. The critic (-) is structured similarly as the discriminator in Section 2.3. projects the input data onto vector space instead of outputting logits for binary classification. For T.) we use a two-layer MLP network. We optimize the query and key projection matrices and ø to minimize CAlign-TT (Q, K) in (4). and optimize 7 to maximize it. The gradient-reversal layer [32 also used to optimize the critic adversarially.
Related work
Alignment Learning. Liang et al. 336 first assign agreement terms for jointly training word align- ment 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 studied in the multi-head attention model based on the Transformer architecture, where alignment extraction is improved by augmenting 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 a fundamental problem in statisties and machine learning (39] The widely used distances include the Kullback-Leible (KL) divergence [40] and Jensen Shannn-Shnnnnnnnn (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 the Wasserstein distance from the optimal transport problem also used for defining the transport cost 226] 27 In our alignment attention, we consider discriminator-base and transport- based 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.
41 42
Experiments
Our method can be incorporated into any self-attention based models. To exam its effectiveness and general applicability, we apply alignment attention a diverse set of tasks, including language under- standing, graph attention, and visual question answering. Furthermore, the model's generalization
5

---

across domains and robustness towards adversarial attacks are studied on language tasks. We study variety of state-of-the the-art art models for these tasks including ALBERT [5), BERT 4 and RoBERTa [ō]. Below we present the main experimental settings and results, with more training and hyperparameter details provided in Appendix B
4.1 Aligmment attention in natural language understanding
Since the self- celf-attention based Transformer was proposed, it has been widely used in pretrained models for NLP and related areas, achieving state-of-the-arr results on various downstream tasks. However, training a -tate--f-thhe-r 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-bas model was trained on 8 Nvidia P100 GPUs for 12 hours and Strubell et al. + 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.
4.1.1 In-domain language understanding evaluation
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-ar pretrained model, ALBERT (5], memory efffory-efficient version of BERT [ 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].
Table 1: Performance of alignment attention on GLUE and SQuAD benchmarks.
MRPC
CoLA
RTE
MNLI
QNLI
QQP
SST
STS
SQUAD 1.1
SQUAD 2.0
ALBERT-BASE
86.5
75.8
92.4
80.86/88.77
78.80/82.07
LLBERT-BASE-AA-CO
87.5
55.7
77.3
85.8
91.3
91.4
92.6
91.1
81.19/88.9 8.98
79.25/82.57
ALBEET-BASE-AA-OO
87.9
54.6
77.0
85.7
91.2
91.3
92.8
91.2
81.13/88.89
79.18/82.44 2.8
ALBERT-BASE+AA-CT
88.6
55.9
77.2
85,9
91.3
91.5
93.1
91.5
81.32/89.2
79.33/82.71
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 independent runs (see full results with the error bars Table 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, 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.
4.1.2 Generalization across domains
In real applications, it is very common deploy a neural network model into 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 (0). The recent work from Desai and Durrett (51) has explicitly elected tasks where out- of-domain performance is substantially lower and challenging domain shifts exhibited. Following the setting in Desai and Durrett (5T), 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 obtain held-out, non-blind test set. We conduct experiments on three tasks: (1) Natural Language Inference. The Stanford Natural Language Inference (SNLI) corpus is 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 grounded commonsense reasoning task [56].
The oft-of data HellaSWAG (HSWAG), which 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 weighted average of the difference between each bin 's accuracy and confidence: C. a conf(G.)), where Gü acc(G;), and conf(G;) are the count, accuracy, and ECE:= = E. x
confidence of samples in the i'th group, respectively. The number of groups is 10 as [51).
Table 2: Results of domain generalization. We report the accuracy and ECE of various models both in-domain data and out-of-domain data for three tasks: natural language inference, paraphrase detection, and commonsense reasoning.
ACCURACY 1
OD
NATURAL LANGUAGE INFURENCE
SNLI
MNLI
SNLI
MNLI
DA
SF
84.63
57.12
1.02
8.79
ESIM
88.32
60.91
1.33
12.78
BART-BASE 51
90.04
73.52
2.54
7.03
BERT-BASD-AA-GAN
90.59
74.15
2.02
5.82
BERT-BASE-AA-CT
90.65
74.23
1.89
5.65
ROBERTA-DASE
91.23
78.79
1.93
3.62
ROBERTA BASE+AA-GAN
91.52
79.55
2.70
3.31
BABERTA-BASE+AA-CC
91.68
79.60
2.79
PARAPIRASE DETECTION
QQP
TWITTER
QQP
TWITTER
85.85
83.36
9.79
ESIM SS
87.75
84.00
3.65
8.38
BERT-BASE STI
90.27
87.63
2.71
8.51
BASA
90.80
88.34
7.48
BASE-AATS-
90.62
88.25
1.74
7.52
ROBERTA-BASE
91.11
86.72
2.33
9.55
ROBERTA BASE+AA-GAN
91.66
87.28
1.78
9.45
ROBERIA-BASE-AA-C
91.53
87.33
1.89
9,40
COMMONSENSE REASONING
SWAG
HS WAG
SWAG
HS WAG
DA 57
46.80
32.48
5.98
40.37
ESIM 58
$2.09
32.08
7.01
19.57
BERT-BASE
79.40
34.48
2.49
12.62
BERT-BASE-AA-GAN
79.56
35.90
1.95
12.11
RERT-RASP-AA-CT
79.60
36.25
1.86
11.78
ROBERTA-DASE 5T
82.45
41 11.68
1.76
11.93
ROBERTA-BASE
83.03
42.51
9.97
ROBERTA-IASE-AA-CT
83.14
42.88
1.43
9.77
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, use SERT-base-und
[4] and RoBERTa-bor base [6] from HuggingFace Transformers (47). We incorporate the alignment attention both pretrained models. Alignment attention consistently outperforms the corresponding soft attention not only for in-domain. confirming our results in Section 4.1.I. 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.
4.1.3 Robustness towards adversarial attacks
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 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.I 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-an 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 conducted for each model with a maximum sentence length of 512. We report the percentages of failed adversarial attacks 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

---

Table 3: Results of pretrained large-scale models' robustness against udversarial attacks, We report the percentages of failed attacks under three adversarial attacks respectively.
ATTACK
ATTENTION
MRPC
COLA
RTE
QQP
SST-2
AVG.
BASE
6.5
2.6
16.2
25.4
7.0
11.5
TEXTFOOLER
AA-GAN
8.4
7.1
14.2
31.2
14,5
15.1
AA-CT
8.7
6.9
15.3
32.1
13.9
15.4
BASE
10.6
16.8
19.9
30.1
40.1
23.5
TEXTBUGGER
AA-GAN
12.6
22.4
20.7
36.0
56.1
29.6
AA-CT
36.5
BASE
44.8
4.9
35.6
48.
13.9
29.6
BAE
AA-GAN
44.2
7.3
35.8
46.5
17.8
30.3
AA-CT
45.3
6.7
36.3
49.4
17.5
31.0
demonstrates its robustness and gets along with our intuition that the alignment attention can learn better and more robust key and query distribution due the use of distributional matching.
4.2 Alignment attention in graph neural networks
To test the general applicability 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 GAT [64, we conduct experiments on three standard citation network benchmark datasets Cora, Citeseer and Pubmed [65 in 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
Table Classification accuracy for graphs.
Attention
Cora
Citeseer
PubMed
GAT
83.00
72.50
77.26
AA-GAN
83.78:10.2
73.32:0.11
78.77:10.2
AA-CT
83.80+00.3
73.49+0.2
78.79+0.2
Results. In Table 4. we report the classification accuracies test nodes 5 random runs, and the standard deviations of alignment attention. We experiment with both AA-GAN and AA-CT. Table that alignment attention consistently improves upon the corresponding bascline 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.
4.3 Attention in visual question answering
Visual question answering (VQA) 67 is 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 MCAN and compare with soft attention. We conduct experiments on the VQA-v2 dataset [67] and follow the hyperparameters and other settings from Yu al. 668. In addition, investigate the model's robustness 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 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 [71], we also report the Patch Accuracy vs Patch Uncertainty (PAvPU) 770 72 as
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
Table 5: Accuracies and PAvPUs of diffcrent attentions on both the original VQA-v2 datasct and the noise ones.
ACCURACY
PAvPU
ORIGINAL
NOISY
ORIGINAL
NOISY
BASE
66.7
63.58
71.96
68.29
AA-GAN
66.92
64.28
72.17
69.80
AA-CT
67.01-0.02
64.57-20.03
72.21=10.03
69.98:10.04

---

Results. The results are summarized in Table 5 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 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.
We report the accuracy and uncertainty of different
0.22
Did
0.60
the batter hit the
0.20
the
0.45
0.18
batter
0.16
hit
0.30
0.14
the
0.15
ball
0.12
Did
the
batter
hit the
bail
Iho ball
Bux Accetton
(s) Aligmment Allention
Figure 3:
For one question from VQA, we visualize the attention weights of AA (a) and soft attention (b). Rows represent queries, and columns represent keys. For example, the left plot, when the is 'Did' and the column is 'hit', the color represents the average attention weight from the query 'Did' to the key 'hit'.
Results Analysis. Visualizations. We plot the attention weights of both alignment attention and soft attention on one question for VQA. Figure attention weight represents the average importance of each query and key pair. For example, in (b), when the row is 'Did' and the column '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 we provide the parameter sizes and step time for different attention types combined with MCAN where the attention module constructs the main model. shows that alignment attention (AA) keeps the parameter size the very similar level as soft attention while moderately increasing the step time compared to soft attention.
Table 6: Efficiency on VQA task.
ATTENTION
PARAMS
S/STEP 1
BASE
43.3M
0,25
AA-GAN
43.4M
0.31
AA-CT
43.4M
0.36
Ablation Study. We conduct ablation study with AA+CT to exam the role the weigh hyperparameter À in Equation [ 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 X. Any number from 0.01 to 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 in Appendix.
5 Conclusion
Our proposed alignment attention aims to match the key and query distributions. We leverage different alignment methods with generic and efficient architecture design which requires surprisingly few modifications to standard soft attention and enables us to casily convert existing soft attention models, including pretrained ones, alignment attention. Our experiments on a variety of language under- standing 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 healtheare 774 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 plug and and-play component to many existing attention models.
Acknowledgements
The authors acknowledge the support of Grant IIS-1812699 from the U.S. National Science Foun- dation, 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 rescarch results reported within this paper.
References
[1] Ilya Sutskever, Oriol Vinyals, and Quoc V Le. Sequence sequence learning with neural networks. In Advances neural information processing systems, pages 3104-3112, 2014.
[2] Dzmitry Bahdanau, Kyung Hyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In 3rd International Conference on Learning Representations, ICLR 2015, 2015.
[3] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is you need. In Advances neural information processing systems, pages 5998 6008, 2017.
[4] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 188.021810 2018.
[5] Zhenzhong Lan, Mingda Chen. Sebastian Goodman, Kevin Gimpel, Piyush Sharma, Radu Sori- cut. ALBERT: A lite BERT for self-supervised learning of language representations. arXiv preprint arXiv: 1942, 2019.
[6] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv e-prints, pages arXiv-1907, 2019.
[7] Mandar Joshi, Danqi Chen, Yinhan Liu, Daniel S Weld, Luke Zeutlemoyer and Omer Levy. Span- bert: Improving pre-training by representing and predicting spans. Transactions of the Association for Computational Linguistics, 8:64-77, 2020.
[8] Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language understanding by generative pre-training. 2018.
[9] Zhilin Yang. Zihang Dai, Yiming Yang. Jaime Carbonell, Ruslan Salakhutdinov, and Quoe V Le. Xlnet: Generalized autoregressive pretraining for language understanding. arXiv preprint 1906::900..822 2019.
[10] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words:
Transformers for image recognition at seale. arXiv preprint vrXiv:2010.11922 2020.
[11] Mark Chen, Alec Radford, Rewon Child, Jeffrey Wa, Heewoo Jun, David Luan, and 1lya Sutskever. Generative pretraining from pixels. In International Conference on Machine Learning, pages 1691-1703. PMLR, 2020.
[12] Yen-Chun Chen, Linjie Li, Licheng Yu, Ahmed El Kholy, Faisal Ahmed, Zhe Gan, Cheng, and Jingjing Liu. Uniter: Learning universal image-text representations. 2019.
[13] Jiasen Lu, Dhruv Batra, Devi Parikh, and Stefan Lee. Vilbert: Pretraining task-agnostic visiolinguistic representations for rision-and-language tasks. arXiv preprint (tXXxv:1908.022000 02060, 2019.
[14] Kelvin Xu, Jimmy Ba. Ryan Kiros, Kyunghyun Cho, Aaron Courville, Ruslan Salakhudinov, Rich Zemel, and Yoshua Bengio. Show, attend and tell: Neural image caption generation with visual attention. In International conference machine learning, pages 2048-2057 2015.
[15] Steven J Rennie, Etienne Marcheret, Youssef Mrouch, Jerret Ross, and Vaibhava Goel. Self-critical sequence training for image captioning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 7008 7024, 2017.
10

---

[16] Dzmitry Bahdanau. Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint urXiv:1409.0473, 2014.
[17] Jian Li. Zhaopeng Tu, Baosong Yang. Michael R Lyu, and Tong Zhang. Multi-heat attention with disagreement regularization. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 2897-2903, 2018.
[18] David Berthelot, Nicholas Carlini, Ian Goodfellow, Nicolas Papernot, Avital Oliver, and Colin Raffel. Mixmatch: A holistic approach to semi-supervised learning. arXiv preprint (Xix::905.022499 2019.
[19] Shujjan Zhang, Chengyue Gong, and Eunsol Choi. Knowing more about questions can help: Improving calibration in question answering. arXiv preprint arXiv:2106.01 01494, 2021.
[20] Zhibin Duan, Yishi Xu, Bo Chen, Dongsheng Wang. Chaojje Wang, and Mingyuan Zhou. Topicnet: Semantic graph-guided topic discovery. In NeurlPS 2021: Neural Information Processing Systems, Dec. 2021.
[21] Shujian Zhang, Chengyue Gong, and Funsol Choi. Learning with different amounts of annotation: From zero to many lahels. arviv preprint 100x:2200,04408,, 2021.
[22] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008.
[23] Zhibin Duan, Dongsheng Wang. Bo Chen, Chaojie Wang. Wenchao Chen, Yewen Li, Jie Ren, and Mingyuan Zhou. Sawtooth factorial topic embeddings guided gamma belief network. In International Conference on Machine Learning, pages 2903-2913. PMLR, 2021.
[24] lan Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Fard Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in neural information processing systems, pages 2672-2680, 2014.
[25] Jianhua Lin. Divergence measures based on the shannon entropy. IEEE Transactions on Information theory, 1551 1991.
[26] Cédric Villani. Optimal Transport: Old and New, volume 338. Springer Science & Business Media, 2008.
[27] Gabriel Peyré and Marco Cuturi. Computational optimal transport: With applications to data science. Foundations and Trendst in Machine Learning, (56-6))35566077 2019.
[28] Huangjie Zheng and Mingyuan Zhou, Comparing probability distributions with conditional transport. arXiv preprint arXiv:2012. 14100, 2021.
[29] Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adverrarril training of neural networks. The journal of machine learning research, 17(1):2096-2330 2016.
[30] Eric Tzeng, Judy Hoffman, Kate Saenko, and Trevor Darrell. Adversarial discriminative domain adaptation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 7167-7177 2017.
[31] Rupesh Kumar Srivastava, Klaus Greff, and Jürgen Schmidhuher. Training very deep networks. In NIPS, 2015.
[32] Yaroslav Ganin and Victor Lempitsky. Unsupervised domain adaptation by backpropagation. In Interna- tional conference on machine learning, pages 1180
1 1189. PMLR, 2015.
[33] Yogesh Balaji, Rama Chellappa, and Soheil Feizi. Robust optimal tramsport with applications in generative modeling and domain adaptation. In Advances in Neural Information Processing Systems, 2020.
[34] Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26:2292-2300 2013.
[35] Jean Feydy, Thibault Séjouméé, François-Xavinn Xavier Vialard, Shun-ichi Amari, Alain Trouvé, and Gabriel Peyré. Interpolating between optimal transport and mmd using sinkhorn divergences. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 2681 2681-2690. PMLR, 2019.
[36] Percy Liang, Ben Taskar, and Dan Klein. Alignment by agreement. In Proceedings of the Human Language Technology Conference of the NAACL, Main Conference, pages 104- 111, 2006.
11

---

[37] Tomer Levinboim, Ashish Vaswani, and David Chiang, Model invertibility regularization: Sequence alignment with without parallel data. In Proceedings of the 2015 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 609-618, 2015.
[38] Tamer Alkhouli, Gabriel Bretschner, and Hermann Ney. On the alignment problem in multi-head attention- based neural machine translation. In Proceedings of the Third Conference on Machine Translation: Research Pupers. pages 177-185, 2018.
[39] Kevin P Murphy. Machine learning: a probabilistii perspective. MIT press, 2012.
[40] Solomon Kullback and Richard A Leibler. On information and sufficiency. The annals of mathematical statistics, 22(1):79-86, 1951.
[41] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein gan. arXiv preprint arXiv:701.07888 2017.
[42] Mare G Bellemare, Ivo Danihelka, Will Dabney, Shakir Mohamed, Balaji Lakshminarayanan, Stephan Hoyer, and Rémi Munos. The cramer distance a solution biased wasserstein gradients. arXiv preprint arXiv: 10743, 2017.
[43] Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and policy considerations for deep learning in nlp. In PurPre Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 3645-3650, 2019.
Linguitti
[44] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. GLUE:
A multi-task benehmark and analysis platform for natural language understanding. arXiv preprint arXiv: 1804.0760) 2018.
[45] Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. arXiv preprint 16Xiv 2016.
[46] Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don't 't know: Unanswerable questions for 2018
SQuAD. arXiv preprint 1806.
[47] Thomas Wolf, Lysandre Dehut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Transformers: State-of-the-at natural language processing. arXiv preprint arXiv: 1910.03771, 2019.
[48] Xilun Chen, Yu Sun, Ben Athiwaratkun, Claire Cardie, and Kilian Weinberger. Adversarial deep averaging networks for cross lingual sentiment classification. Transactions of the Association for Computational Linguistics, 6:557-570, 2018
[49] Minlong Peng, Qi Zhang, Yu-gang Jiang, and Xuan-Jing Huang. Cross-domain sentiment classification with target domain specific information. In Proceedingy of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pupers), pages 2505-2513, 2018
[50] Timothy Miller. Simplified neural unsupervised domain adaptation. In Proceedings of the conference. Association for Computational Linguistics. North American Chapter. Meeting, volume 2019, page 414. NIH Public Access, 2019.
[51] Shrey Desai and Greg Durrett. Calibration of pre-trained transformers. arXiv preprint x/Xiv:2003.07892. 2020.
[52] Samuel R Bowman, Gahor Angeli, Christopher Potts, and Christopher D Manning. A large annotated corpus for learning natural language inference. In EMNLP, 2015.
[53] Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long Papers),
pages 1112 1122, 2018.
[54] Shankar Iyer, Nikhil Dandekar, and Koréel Csernai. First quora dataset release: Question pairs. data. quora. com, 2017.
[55] Wuwei Lan, Siyu Qiu, Hua He, and Wei Xu. A continuously growing dataset of sentential paraphrases. In Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, pages 2017.
12

---

[56] Rowan Zellers, Yonatan Bisk, Roy Schwartz, and Yejin Choi. Swag: A large-scale adversarial dataset for grounded commonsense inference. In EMNLP, 2018.
[57] Ankur P Parikh. Oscar Täckström, Dipanjan Das, and Jakob Uszkoreit. A decomposable attention model for natural language inference. In EMNLP, 2016.
[58] Qian Chen, Xiaodan Zhu, Zhen-Hua Ling, Si Wei, Hui Jiang, and Diana Inkpen. Enhanced Istm for natural language inference. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pages 1657-1666 2017.
[59] lan Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. arXiv preprint arXiv: /4l::44122655 26722, 2014.
[60] John Morris, Eli Lifland, Jin Yong Yoo, Jake Grigsby, Di Jin, and Yanjun Qi. Textattack: A framework for adversarial attacks, data augmentation. and adversarial training in nip. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 119-126, 2020.
[61] Di Jin, Zhijing Jin, Joey Tianyi Zhou, and Peter Szolovits. Is bert really robust? a strong baseline for natural language attack on text classification and entailment. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 8018-8025, 2020.
[62] Li, S Ji, Du, B Li, and T Wang. Textbugger: Generating adversarial text against real-world applications. In 26th Annual Network and Distributed System Security Symposium, 2019.
[63] Siddhant Garg and Goutham Ramakrishnan. Bae: Bert-hased adversarial examples for text classification. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 6174-6181 6181, 2020.
[64] Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint 1eXiv:1710.1000 2017.
[65] Prithviraj Sen. Galileo Namata, Mustafa Bilgic. Lise Getoor, Brian Galligher, and Tina Eliassi-Ran Rad. Collective classification in network data. Al magazine, 29(3)993-93, 2008
[66] Zhilin Yang. William W Cohen, and Ruslan Salakhutdinov. Revisiting semi semi-super learning with graph embeddings. arXiv preprint 3tfi::60000888/ 2016.
[67] Yash Goyal, Tejas Khot, Douglas Summers Stay. Dhruy Batra, and Devi Parikh. Making the V in VQA matter: Elevating the role of image understanding in visual question answering. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 6904 6913,,1,,000 2017.
[68] Zhou Yu, Jun Yu, Yuhao Cui, Dacheng Tao, and Qi Tian. Deep modular co-attention networks for visual question answering. In Proceedings of the IEFE Conference on Computer Vision and Pattern Recognition, pages 0281-62900 2019.
[69] Hugo Larochelle, Dumitru Erhan, Aaron Courville, James Bergstra, and Yoshua Bengio. An empirical evaluation of deep architectures on problems with many factors of variation. In Proceedings of the 24th international conference on Machine learning, pages 473-480, 2007.
[70] Xinjie Fan, Shujian Zhang. Bo Chen, and Mingyuan Zhou. Bayesian attention modules. Advances in Neural Information Processing Systems, 33, 2020.
[71] Shujian Zhang, Xinjie Fan, Bo Chen, and Mingyuan Zhou. Bayesian attention belief networks, arXiv preprint 1xXiv:2006.05251,, 2021.
[72] Jishnu Mukhoti and Yarin Gal. Evaluating Bayesian deep learning methods for semantic segmentation. arXiv preprint 1rxiv:1811.127111 2018.
[73] Jinkyu Kim and John Canny. Interpretable learning for self-driving cars by visualizing causal attention. In Proceedings of the IEEE international conference on computer vision, pages 2942-2950, 2017.
[74] Edward Choi, Mohammad Taha Bahadori, Jimeng Sun, Joshua Kulas, Andy Schuetz, and Walter Stewart. Retain: An interpretable predictive model for healthcare using reverse time attention mechanism. In Advancey in Neural Information Processing Systems, pages 3504-3512, 2016.
[75] Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, D Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanan. and Jasper Snoek. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift, arXiv preprint arXiv: 02530. 2019.
13

---

[76] William B Dolan and Chris Brockett. Automatically constructing corpus of sentential paraphrases. In Proceedings of the Third International Workshop on Paruphrasing (IWP2005), 2005.
[77] Alex Warstadt, Amanpreet Singh, and Samuel R Bowman. Neural network acceptability judgments. Transactions the Association for Computational Linguistics, 7:625-644 2019.
[78] Ido Dagan, Oren Glickman, and Bernardo Magnini. The pascal recognising textual entailment challenge. In Machine Learning Challenges Workshop, pages 177-190. Springer, 2005.
[79] Adina Williams, Nikita Nangia, and Samuel R Bowman. A broad-coverage challenge corpus for sentence understanding through inference. arXiv preprini arXiv:1004.05426, 2017.
[80] Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Y Ng. and Christopher Potts. Recursive deep models for semantic compositionality sentiment treebank. In Proceedings of the 2013 conference on empirical methods in natural language processing, pages 1631-1642, 2013.
[81] Daniel Cer, Mona Diab, Eneko Agirre, Inigo Lopez-Gazpio, and Lucia Specia. Semeval-2017 task I: Semantic textual similarit-mullilnnuaal and cross-lingual focused evaluation. arXiv preprint arXiv:1708,00055. 2017.
[82] Matt Gardner, Joel Grus, Mark Neumann, Oyvind Tafjord, Pradeep Dasigi, NH Liu, Matthew Peters, Michael Schmitz, and Luke S Zettlemoer. A deep semantic natural language processing platform. arXiv preprint arXiv: 006.07640, 2017.
[83] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Reprexentations, 2018
[84] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks, In Proceedings of the thirteenth international conference artificial intelligence and statistics, pages 249-256, 2010.
[85] Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint vrXiv:1412.6988 2014.
[86] Amm Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint arX1v:151,.0728,, 2015.
[87] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way prevent neural networks from overfitting. The Journal of Machine Learning Research, 15 (1):1929-1958, 2014.
[88] Xinjie Fan, Shujian Zhang, Korawat Tanwisuth, Xiaoning Qian, and Mingyuan Zhou. Contextull dropout: An efficient sample-dependent dropout module. In International Conference on Learning Representations. 2021.
[89] Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft COCO: Common objects in context. In European conference on computer vision, pages 740-755. Springer, 2014.
[90] Yuntian Deng, Yoon Kim, Justin Chiu, Demi Guo, and Alexander Rush. Latent alignment and variational attention. In Advances in Neural Information Processing Systems, pages 9712-9724 2018.
[91] Damien Teney. Peter Anderson. Xiaodong He, and Anton Van Den Hengel. Tips and tricks for visual question answering: Learnings from the 2017 challenge. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 4223 4232,, 2018
14


Reviews:
[{"review_id": "WpMqUpLg-v", "reviewer": "Reviewer_qa5u", "summary": "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. \nCorrespondingly, 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. ", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I recommend accepting the paper. In general, this paper is well-written and easy to follow. It aims at studying the interactions between the keys and queries for attention modules. A variety of experiments on different problem settings with different proposed alignment methods are performed, and, the proposed method consistently outperforms the baseline, which proves that explicitly regularizing the key and query distributions can help to improve the feature learning for the attention-based methods.  I summarize the strengths and weaknesses as follows:\n\nStrengths:\n- This paper proposes a plug-and-play alignment module that can be applied to various attention-based models. \n- This paper proves that explicitly regularizing the distributions of the key and query feature can help the feature learning of the self-attention module. \n- Three alignment methods are introduced to regularizing the distributions of the keys and queries. \n- The experimental results are convincing. Experiments are performed in different language understanding settings. It shows that the proposed alignment can bring consistent improvements no matter which alignment method is adopted. \n\nWeaknesses:\n- It lacks the experimental results showing that after being trained with the alignment loss, the distribution of the keys and queries are well-aligned. \n- In the baseline attention module, the keys and queries will be dot-producted, which means they should be of the same feature domain. Thus, it has implicitly introduced the regularization of their alignments. In comparison, the novel part of the proposed method is explicitly modeling the regularization. However, it lacks a detailed explanation on why explicitly align the keys and queries can bring improvements over the baselines. \n- Typos: Line-104 the tensor dimension should be [B, H, w, d]", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3h"}, "extra_text": {}}, {"review_id": "WbDxO2aJTRA", "reviewer": "Reviewer_UFmi", "summary": "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.\n\n> After Rebuttal:\n\nThanks the authors for the response. I'd like to keep my initial score.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1. The paper is well written. \n2. The paper gives extensive experiments to prove the effectiveness of the proposed approach in terms of various aspects, i.e., accuracy, uncertainty estimation, generalization across domains, and robustness to adversarial attacks.\n3. The proposed approach achieves state-of-the-art results on several tasks while remaining efficient in memory and computation cost.\n\nWeaknesses:\n1. The contributions are limited.\n- The idea of alignment learning has been explored in literature. Thus, the main contribution of this paper is the methodological contribution, i.e., the alignment attention model, which can align the query and key distributions. However, the motivation is unclear.\n\n2. The motivation and the studied problem are unclear.\n- Personally, it is unclear what problems of the conventional self-attention model can be solved by the proposed approach? In other words,  what are the problems that will cause if the distributions between query and key are not aligned?\n- Specifically, in L44-45, the authors claim that \"misalignment could lead to an undesired dependence between the attention weights of a token and its position in the input feature space\", however, there is no evidence/experimental analysis that can support this claim. How to prove it? Although the extensive experiments have indeed demonstrated the effectiveness of the proposed approach, I would like to know why performing the alignment of query and key distributions can improve the performance. \n\n3. Some detailed analyses are missing.\n- Can the proposed alignment attention be incorporated into the source-to-target attention component?\n- In the discriminator-based modules, i.e., L137, why do you adopt the complex highway architecture to construct the discriminator? What are the performances of a single MLP (i.e., FC-Relu-FC) or a fully connected layer (i.e., FC)? Besides, to train the discriminator, what are the output dimension in the discriminator and the adopted loss, e.g., BCE/MSE?\n- In Section 2.3, what are the performances of adopting the widely-used distance losses, e.g., KL Loss, to perform the alignment.\n- In Equation (2), what are the performances of exchanging query and key distributions?\n- If the approach brings errors, and what type of errors does it bring, and why? More importantly, are there errors caused by the approach when the model well aligns the distributions of query and key?\n- The margins of the gains in all Tables are small and there appears to be no statistical significance testing. Statistical significant test and error ranges are highly appreciated.\n\n4. Some parts of the paper can be made clearer for easier reading. \n- In all Tables, I encourage the authors to list several performances of recently published models to help the readers better understand the proposed approach.\n- For the results of all base models in all experiments, it is better to clarify whether the results are reproduced by yourself or copied directly from the original paper.\n- L48: 'emprical' -> 'empirical'; L180: 'is' -> 'are'; L269: 'with maximum' -> 'with a maximum'.\n\nI highly encourage the authors to address these concerns and I am glad to increase my score if I see satisfying answers, thanks!\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No.\n\n- If the approach brings errors, and what type of errors does it bring, and why? More importantly, are there errors caused by the approach when the model well aligns the distributions of query and key?\n- The margins of the gains in all Tables are small and there appears to be no statistical significance testing. Statistical significant test and error ranges are highly appreciated.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7 hours"}, "extra_text": {}}, {"review_id": "Ln751qD3nuA", "reviewer": "Reviewer_E8gR", "summary": "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.", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Main Review:\n\nThe paper rethinks the dominant self-attention mechanism in both deep neural networks and points out the alignments of query and key distributions in the current attention mechanism are largely unexplored. The motivation of align the projection of both query and key tokens is strong and meaningful as decreasing the risk of misaligned interactions between queries and keys. To conclude, the paper attempts to take distribution matching into consideration of attention mechanism.\nBased on that motivation, the authors proposed three alignment methods attempting to align the attention and finally adopted adversarial training and conditional transport methods to align the distributions. The analyses of methods are theoretical and insightful.\nThe authors conducted extensive experiments in natural language understanding, GNN and VQA domains. The evaluations are not only in-domain dataset, but also out-domain ones to illustrate the generalization across domains. Adversarial attacks are taken into consideration to show the robustness of alignment attention. The experiments are very meticulous and rational to show the effectiveness of alignment attention as a plugin-and-play component.\n\nLimits:\n\nAs is known, the multi-head attention mechanism is widely applied in vision tasks like classification, detection and segmentation. The authors have conducted detailed ablations to show effectiveness in NLP, GNN and VQA covering language, graph and multi-modal domains. I suggest supplements of some vision tasks to finalize the ablations would make the paper more completed \n\nReason for weak rejection:\n\nMy concern about the paper is that do we really need alignment between key and query? As key and query are linear projections of input, the distribution of key and query should be well aligned. Thus I do not think the alignment story is reasonable and suggest weak rejection.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No", "review_text": "", "extra_scores": {"ethical_concerns": "No", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "10 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: I recommend acceptance of the paper for the following reasons. As pointed out by reviewer E8gR, the premise of aligning key and value distributions might not be as important as it sounds at first for the purposes of improved learning. In addition, the analysis and ablations studies could be more thorough and convincing. Yet, experimentally the module works well. So despite potential limitations on the "intuition" and "analysis" of the module, there is value in knowing that such a module has empirical validation, especially for its potential impact on future analyses. There is also agreement from reviewers about this point and general support from two reviewers for the paper.

For these reasons, I recommend to accept the paper.
Decision: Accept (Poster)