Alignment Attention by Matching Key and Query Distributions

1 Introduction . 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 . 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 in 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 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 ] . 2.2 Alignment attention . 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 . 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 [ 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 . 2.3.1 Adversarial training-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 , 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 .

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