Cascaded Head-colliding Attention (2021)

Abstract: Transformers have advanced the field of natural language processing (NLP) on a variety of important tasks. At the cornerstone of the Transformer architecture is the multi-head attention (MHA) mechanism which models pairwise interactions between the elements of the sequence. Despite its massive success, the current framework ignores interactions among different heads, leading to the problem that many of the heads are redundant in practice, which greatly wastes the capacity of the model. To improve parameter efficiency, we re-formulate the MHA as a latent variable model from a probabilistic perspective. We present c ascaded head-c o lli d ing a ttention (CODA) which explicitly models the interactions between attention heads through a hierarchical variational distribution. We conduct extensive experiments and demonstrate that CODA outperforms the transformer baseline, by 0.6 perplexity on Wikitext-103 in language modeling, and by 0.6 BLEU on WMT14 EN-DE in machine translation, due to its improvements on the parameter efficiency.

## introduction transformers @xcite have advanced the field of natural language processing (nlp) on a variety of important tasks, including language modeling @xcite , language understanding @xcite , and machine translation @xcite @xcite . it has also found its place in computer vision @xcite , and in intelligent agents @xcite where sequence modeling plays a key role as well. the cornerstone of the transformer architecture is the multi-head attention (mha) mechanism which models pairwise interactions between the elements of the sequence. an attention function can be described as mapping a query and a set of key-value pairs to an output, where the query, keys, values, and output are all vectors. the output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key. a multi-head attention (mha) mechanism extends the idea through performing multiple separately parameterized attention functions acting in parallel to contextualize the input representations. their outputs are then gathered by an affine transformation, allowing the model to jointly attend to information from different representation subspaces at different positions. despite its massive success, the current framework ignores the interactions among different heads, leading to the problem that many of the heads are redundant in practice (i.e., attending to the same regions of the sequence), which underutilizes the capacity of the model @xcite . at the same time, recent research @xcite @xcite demonstrates that heads in mha have the potential to capture distinct information from input sequences, ranging from syntactic and semantic features to alignment information between source and target sentence pairs. these observations suggest that multiple heads should be encouraged to extract complementary information. therefore, it is highly appealing to take into account the interactions among different attention heads from the perspective of parameter efficiency and the expressiveness of the model. in this work, we introduce head-colliding attention ( §3). we formulate mha as a probabilistic model, where each attention head is represented by a latent variable and all of them collide into the observed sequence data (figure 1a ). in this probabilistic graphical model structure, attention heads work as individual factors to explain the data. although each factor is independent of each other a priori, they interact with each other automatically, conditioning on observations, thanks to the explaining-away effects @xcite . the head-colliding attention mechanism introduces new computational challenges in training the model. we will discuss how we tackle these using variational methods @xcite . we propose cascaded head-colliding attention (coda, figure 1b ). as our main model, coda adopts a hierarchical variational distribution @xcite to allow both rich head interactions and effective computations ( §4). we validate our method in language modeling and machine translation experiments ( §5). coda outperforms the vanilla mha transformer on both tasks, on wikitext-103 by 0.6 perplexity and on wmt14 en-de by 0.6 bleu. further analysis shows that coda learns to encourage diversity in different heads (figure 2 ) and to promote parameter efficiency when increasing the number of heads ( §5.3). ## background multi-head attention (mha) mechanism plays an important role in modern transformer architecture @xcite . it extends the classical attention mechanism by running multiple attention function heads in parallel. an mha module is composed of h identical blocks (usually referred to as attention heads). each head will generate a hidden state h i based on the input query, key and value matrices, denoted as q, k, and v respectively. the hidden states from different heads are then aggregated as the output of the mha module: n @xmath0, where w o i are model parameters. in the i-th head, the input matrices q, k and v are first linearly projected into different subspace representations q i , k i , and v i , based on different learnable parameters. after that, we compute the inner product over all projected queries and keys as the attention logits z i , which are then passed through a row-wise softmax 2 to obtain head attention weights a i : the final output of a single attention block is the weighted sum of v i : as we can see, the core of mha is to calculate a i in each head. we thus refer to a i as the i-th attention head. in sequence prediction tasks, the model takes as input a source sequence of length m and outputs a target sequence of length n in an autoregressive manner. it predicts each token y within the target sequence through a categorical distribution p vanilla (y|x), where x includes the source sequence as well as a previously generated prefix. with respect to an mha block a 1 , . . . , a h , the model predicts target tokens y by first feeding these heads into a complex non-linear transformation foot_0 denoted by φ(•), and then passing it through a softmax function over the entire vocabulary. therefore, the output probability can be written as p vanilla (y|x) = f (a 1 , . . . , a h ), where ## head-colliding attention in this section, we introduce head-colliding attention. specifically, we formulate mha as a probabilistic model, where each attention head is represented by a latent variable. the name reflects a "collider" in the context of probabilistic graphical models (figure 1a ). we will first explain how head-colliding attention permits the modeling of interactions among different heads and then discuss how vanilla mha can be viewed as a marginalized version of head-colliding attention, which ignores any head interactions. considering a single mha block, we cast each attention head a i as a latent variable. the probability of target y conditioning on input x can be obtained by marginalizing over all heads a (we denote a := {a 1 , . . . , a h }): a|x) is the joint prior distribution. the corresponding directed graphical model is demonstrated a 1 a 2 a h y . . . (a) head-colliding attention. a 1 1 a 1 2 a 1 h a 2 1 a 2 2 a 2 h a l 1 a l 2 a l h y . . . . . . . . . (b) cascaded head-colliding attention (coda). figure 1: (a) left: probabilistic graphical model (pgm) diagram of head-colliding attention. although each head variable is independent a priori, they interact with each other after observing targets y, which is referred as explaining-away effect. (b) right: pgm diagram of a 3-layer cascaded head-colliding attention (coda). a l i denotes the i-th attention head at transformer layer l. note that all dependencies from x are omitted in these diagrams for simplicity. in figure 1a , where the links from different heads collide on the observation variable y. a crucial property of this graphical model is the "explainingaway" effect @xcite of attention heads a when observing the output y. in other words, if a head a i attends to part of the input which accords well with observation, it immediately discourages other heads from attending to the same part of the input but encourages them to look into complementary information. 4 this mechanism effectively reduces head redundancy and in turn improves parameter efficiency. vanilla vs. head-colliding attention we now take a closer look at the vanilla mha ( §2). recall that in vanilla mha, all attention heads are deterministic. from the perspective of latent variable models, this is computationally equivalent to taking expectations of latent head variables. the output probability distribution p vanilla (y|x) can then be expressed as: this means we are only interested in the individual expectations when using the attention heads in vanilla mha for predictions. on the contrary, in head-colliding attention the distribution of y is defined as: note the inherent difference of when to take the expectation in vanilla and head-colliding attention. since f (•) is a complex non-linear function ( §2), these two formulations are not equivalent in general and may have a large gap between the two distributions. concretely, vanilla mha ignores any possible interactions among different heads. as indicated in equation 2, it first marginalizes out every single head before observing targets -one head will not learn what other heads are attending to despite the fact y is observed. this is why vanilla mha is prone to redundancy as many previous studies @xcite , inter alia) discovered. head-colliding attention, on the other hand, permits rich head interactions due to the expressive non-linear function f (•) inside the expectation over different latent variables a 1 , . . . , a h . however, the complexity of head interactions also leads to intractability in training the model, which we will discuss in the next section. ## training head-colliding attention we train the model by performing maximum likelihood estimation. here, the log marginal likelihood can be expressed as: unfortunately, this is intractable in general because it requires marginalizing over all possible configurations of attention heads. the standard technique is to use variational inference, which optimizes the log marginal by maximizing its evidence lower bound (called elbo) @xcite : a straight-forward choice of q(a|x) is to use the mean-field approximation @xcite : however, it has similar drawbacks as the vanilla mha. 6 the mean-field approximation assumes the independence of different heads and hence the interactions are greatly limited. alternatively, one could parameterize q(a|x) using an auto-regressive model. 7 although this is much more expressive, its sequential nature severely slows down training, making it infeasible in practice. cascaded head-colliding attention our solution to this problem is to employ hierarchical structures for head-colliding attention, where interactions among heads could be effectively incorporated into the model @xcite . conveniently, the hierarchical nature of the transformer architecture offers an effective way of constructing such proposal distributions. given a transformer with l layers, we denote the set of all attention heads at layer l -1 and l as a l-1 and a l , respectively. following the bottom-up computation of the transformer, the distribution of a l must rely on the instantiated values of a l-1 . in this sense, a l-1 can be seen as the common variables that govern a l (figure 1b ). formally, we have: despite the fact that each attention head a l @xmath1, they become dependent when we marginalize a l-1 out. in particular, the marginal distribution of each a l becomes: this corresponds to an infinite mixture of the meanfield distributions q(a l |x, a l-1 ) and is able to capture rich head interactions @xcite . our main model adopts this cascaded proposal distribution in figure 1b , and therefore we name it cascaded head-colliding attention (coda). the only problem left now is how to specify the conditional distribution q(a l |x, a l-1 ) for all @xmath2, 2, . . . , l. we first impose the basic constraints on head values as in vanilla mha, that is, all head values must range within a simplex ∆ n-1 : here a l i,:k is the k-th column of the i-th attention head at layer l and 1 denotes the vector of all 1's. for efficient training and inference, we adopt gaussian-logistic distributions @xcite , which not only satisfy the constraints above but also benefit from the effective reparameterization trick @xcite @xcite . in particular, recall that in vanilla mha, a @xmath3. we also denote the attention logits at l-th layer as z l := {z l 1 , . . . , z l h }. for head i at layer l, we first sample from a multivariate gaussian distribution q(z l i,j: |z l-1 i,j: ) foot_5 and pass the samples into a row-wise softmax function to yield head values: to explicitly model hierarchical structures among attention heads, we propose to add a direct connection between attention heads at adjacent layers (figure 1b ). such connections offer direct access to the information of attention in the previous layer. specifically, for each head i at layer l we set the mean µ i l as the sum of two parts: ## experiments we conduct experiments on language modeling and machine translation tasks. ## related work attention mechanisms were first applied to recurrent networks in @xcite . it was then extended to multi-head attention (mha) and became the key component in transformer architectures @xcite . to study the utility of multiple attention heads, @xcite focused on identifying individual contributions of each attention head. @xcite conducted extensive experiments to demonstrate that pruning out most heads after training does not lead to a drop in performance during inference. @xcite further revealed that replacing learnable attention heads with samples from fixed gaussian distributions can achieve almost the same performance as original models. additionally, @xcite proposed to iteratively prune attention heads during training based on the lottery ticket hypothesis. these works indicate that there is a lot of head redundancy in the mha transformer architectures. instead of pruning unnecessary parameters and down-sizing transformer models, there are also works that propose to improve parameter efficiency in transformers. for instance, @xcite introduced a regularization term to explicitly promote diversity among different heads. @xcite proposed to use convolutional kernels to capture correlations among not only local windows of sequences, but also different heads. @xcite considered each head as a sample from the same distribution, and presented a sampling algorithm that avoids samples from collapsing into local modes. it hence explicitly encouraged the repulsiveness in mha. besides, mae @xcite converted a vanilla mha to a mixture-of-experts model, where each expert component activates only a subset of attention heads. with learned probabilities, different experts could be specialized on different inputs. different from these works, coda does not explicitly promote head diversity nor specialize different heads. instead, we focus on studying head interactions from a probabilistic perspective, which reveals the close connection between vanilla mha and coda. another research line relating to our work is to incorporate latent variables into attention modules. @xcite investigated the connection between vanilla deterministic single-head attention and its stochastic counterpart. @xcite explored this further and proposed to use variational inference techniques for training the model. they considered both cases of discrete and continuous latent variables. bayesian attention modules @xcite introduced continuous latent distributions for attention that are amenable to reparameterization tricks. our work is different from them in that we mainly investigate the mha mechanism and aim to improve parameter-efficiency by recovering potential interactions among different heads, which are ignored in vanilla mha. concurrently, @xcite proposed to add residual connections between attention scores at adjacent layers, similar to our cascading connections. nevertheless, our motivation for using the cascaded structure is quite different: we aim to construct direct hierarchical dependencies for latent variable models, while @xcite is mainly motivated to improve transformer architectures and obtain performance gains. ## conclusion and future work we present coda by re-formulating the multi-head attention (mha) as a latent variable model from a probabilistic perspective. coda explicit models of the interactions among attention heads through a hierarchical variational distribution. we conduct extensive experiments and demonstrate that coda outperforms the transformer baseline in language modeling and machine translation. the analysis shows that coda learns to encourage the diversity in different heads and to promote parameter efficiency when increasing the number of heads. in this framework, we will be able to impose explicit constraints or regularization on different attention heads in a principal way (e.g. informative priors that promote diversity). besides, we can also consider more expressive (data-driven) variational distributions. we leave these as the future work. our code is publicly available at https://github.com/lzhengisme/coda .