Mask Attention Networks: Rethinking and Strengthen Transformer

Transformer has two sublayers: Self-Attention Network (SAN) and Feed-Forward Network (FFN).

As discussed in Vaswani et al. (2017), an attention function maps a query and a set of key-value pairs to an output shown in Equation 1.

$$\begin{gathered}\mathcal{A}(Q,K,V)=\mathcal{S}(Q,K)V\\ \mathcal{S}(Q,K)=\Bigg{[}\frac{\text{exp}\big{(}Q_{i}K_{j}^{T}/\sqrt{d_{k}}\big{)}}{\sum_{k}\text{exp}\big{(}Q_{i}K_{k}^{T}/\sqrt{d_{k}}\big{)}}\Bigg{]}\end{gathered}$$

(1)

where the queries $Q$, keys $K$ and values $V\in\mathbb{R}^{T\times d_{k}}$ are all matrices.

SAN produces representations by applying attention function to each pair of tokens from the input sequence. It is beneficial to capture different contextual features with multiple individual attention functions. Given a text representation sequence $H^{l}\in\mathbb{R}^{T\times d}$. in the $l$-the layer.

$$\begin{gathered}H^{l}=\big{[}A^{1},\cdots,A^{I}\big{]}W_{H}\\ A^{i}=\mathcal{A}\big{(}H^{l}W^{i}_{Q},H^{l}W^{i}_{K},H^{l}W^{i}_{V}\big{)}\end{gathered}$$

(2)

where $\{W^{i}_{Q},W^{i}_{K},W^{i}_{V}\}\in R^{d\times d_{k}}$ are trainable parameters, $i$ denotes the attention head and $d$ is the hidden size.

In FFN, the computation of each $h^{l}_{t}$ in $H^{l}$ is independent of others. It consists of two affine transformations with a pointwise non-linear function:

$$\begin{gathered}H^{l+1}=\text{ReLU}\big{(}H^{l}W_{1}\big{)}W_{2}\end{gathered}$$

(3)

where $W_{1}$ and $W_{2}$ are matrices of dimension $d\times d_{f}$ and $d_{f}\times d$, respectively. Typically, $d_{f}$ is set to be 4 times larger than $d$.

On the basis of attention function in Equation 1, we define a new mask attention function:

$$\begin{gathered}\mathcal{A}_{M}(Q,K,V)=\ {S}_{M}(Q,K)V\\ \mathcal{S}_{M}(Q,K)=\Bigg{[}\frac{M_{i,j}\text{exp}\big{(}Q_{i}K_{j}^{T}/\sqrt{d_{k}}\big{)}}{\sum_{k}M_{i,k}\text{exp}\big{(}Q_{i}K_{k}^{T}/\sqrt{d_{k}}\big{)}}\Bigg{]}\end{gathered}$$

(4)

where $M\in\mathbb{R}^{T\times T},M_{i,j}\in[0,1]$ is a mask matrix and can be static or dynamic. Intuitively, the value in each position of $M$ can be viewed as the color shade in Figure 1.

With the knowledge of mask attention function, we introduce Mask Attention Networks(MANs), in which each network can be written as Equation 5.

$$\begin{gathered}H^{l+1}=\mathcal{F}\big{(}\big{[}A^{1}_{M^{1}},\cdots,A^{I}_{M^{I}}\big{]}\big{)}W_{H}\\ A^{i}_{M^{i}}=\mathcal{A}_{M^{i}}\big{(}H^{l}W^{i}_{Q},H^{l}W^{i}_{K},H^{l}W^{i}_{V}\big{)}\end{gathered}$$

(5)

where $\mathcal{F}$ is the activation function, $M^{i}$ is the mask matrix for the $i$-th attention head.

Next, we show that SAN and FFN both belong to the Mask Attention Networks.

For SAN, let $M=[1]\in\mathbb{R}^{T\times T}$ be an all-ones matrix and $\mathcal{F}=\mathcal{F}_{id}$ be the identity function, its mask attention function would be formalized:

$$\begin{gathered}\mathcal{S}_{[1]}(Q,K)=\Bigg{[}\frac{1\cdot\text{exp}\big{(}Q_{i}K_{j}^{T}/\sqrt{d_{k}}\big{)}}{\sum_{k}\text{exp}\big{(}Q_{i}K_{k}^{T}/\sqrt{d_{k}}\big{)}}\Bigg{]}=\mathcal{S}(Q,K)\\ \mathcal{A}_{[1]}(Q,K,V)=\mathcal{S}_{[1]}(Q,K)V=\mathcal{A}(Q,K,V)\end{gathered}$$

(6)

Then, the MAN degenerates into SAN.

	$\displaystyle H^{l+1}$	$\displaystyle=\mathcal{F}_{id}\Big{(}\big{[}A^{1}_{[1]},\cdots,A^{h}_{[1]}\big{]}\Big{)}W_{H}$		(7)
		$\displaystyle=\big{[}A^{1},\cdots,A^{h}\big{]}W_{H}$

For FFN, let $M=\mathbb{I}\in\mathbb{R}^{T\times T}$ be the identity matrix, $\mathcal{F}=\text{ReLU}$ and head number $I=1$.

$$\begin{gathered}\mathcal{S}_{\mathbb{I}}(Q,K)=\Bigg{[}\frac{1_{i}(j)\cdot\text{exp}\big{(}Q_{i}K_{j}^{T}/\sqrt{d_{k}}\big{)}}{\sum_{k}1_{i}(k)\cdot\text{exp}\big{(}Q_{i}K_{k}^{T}/\sqrt{d_{k}}\big{)}}\Bigg{]}=\mathbb{I}\\ \mathcal{A}_{\mathbb{I}}(Q,K,V)=\mathcal{S}_{\mathbb{I}}(Q,K)V=\mathbb{I}V=V\end{gathered}$$

(8)

where $1_{i}(x)$ is an indicator function that equal to 1 if $x=i$, otherwise 0.

The MAN degenerates into FFN.

$\displaystyle H^{l+1}$

$\displaystyle=\text{ReLU}\Big{(}\big{[}A^{1}_{M}\big{]}\Big{)}W_{H}=\text{ReLU}\big{(}H^{l}W^{1}_{V}\big{)}W_{H}$

(9)

In summary, SAN and FFN are two special cases in MANs with different static mask matrices.

The mask matrix of SAN is an all-ones matrix and that of FFN is an identity matrix, they are two extreme cases in MANs. We analyze that these two static MANs are deficient in localness modeling. Intuitively, through blocking other tokens in advance, FFN focuses on its own information and is unable to perceive the information except itself, let alone its neighbors. In SAN, each token is equally accessible to any other ones. As the example in Introduction shows, we find that tokens not in neighborhood are also likely to attend to each other with relatively large scores. Therefore, SAN might introduce noises to semantic modeling and overlook the relation of neighboring signals.

We demonstrate the issue of self-attention. Generally assuming that $\big{[}a,b,c\big{]}$ appear in sequence, and $(a,b),(b,c)$ are two neighbor pairs, but $a,c$ are not neighbors.

First, to explicitly define the relationship of tokens, we introduce ${U_{\delta}(h)}$ as the set of tokens at the distance of $\delta$ from $h$ with key and query linear transformation in SAN, in other words, $u\in U_{\delta}(h)\Leftrightarrow||hW_{Q}-uW_{K}||_{2}^{2}\leq\delta$. For example, if $(a,b)$ is a neighbor pair, there would exist some small $\delta\geq 0$ such that $a\in{U_{\delta}(b)}$ and $b\in{U_{\delta}(a)}$.

Second, we know that the larger the inner product is, the smaller the Euclidean distance is, and vice versa. With the awareness of the relationships between $\big{[}a,b,c\big{]}$, we have $a,b\in{U_{\delta}(a)}$, $b,c\in{U_{\delta}(c)}$ and $a,b,c\in{U_{\delta}(b)}$ for some small $\delta\geq 0$.

Third, we are able to estimate the semantic distance between $a$ and $c$ as the Equation 10 shows.

		$\displaystyle||aW_{Q}-cW_{K}||_{2}^{2}$		(10)
	$\displaystyle=$	$\displaystyle||aW_{Q}-bW_{K}+bW_{K}-bW_{Q}+bW_{Q}-cW_{K}||_{2}^{2}$
	$\displaystyle\leq$	$\displaystyle 3||aW_{Q}-bW_{K}||_{2}^{2}+3||bW_{K}-bW_{Q}||_{2}^{2}$
	$\displaystyle+$	$\displaystyle 3||bW_{Q}-cW_{K}||_{2}^{2}\big{)}\leq 9\delta$

Thus, though $a$ and $c$ are not neighbors, no matter how irrelevant the semantics of $a$ and $c$, $c\in{U_{9\delta}(a)}$ that $c$ would play an important role in modeling semantics of $a$.

The upper phenomenon illustrates following normal attention function in Equation 1, some tokens not in neighborhood not are still likely to occupy an important position in attention weight that can not be ignored.

With the knowledge of MANs, we propose to mask other tokens that not in neighborhood of the target token for better local semantic modeling.

For example, we build a distance-dependent mask matrix SM. If each token only model the relationship with those tokens within $b$ units of itself, we can set

$\displaystyle\text{SM}[t,s]=\left\{\begin{array}[]{cc}0,&|\ t-s\ |>b\\ 1,&|\ t-s\ |\leq b\end{array}\right.$

(11)

where $t,s$ are the positions of query and key, and $\text{SM}[t,s]$ is the value of the $t$-th row and $s$-th column of SM .

By means of SM, we take those tokens within $b$ units into account and ignore others. The static mask does assign more weights to a specific neighborhood, but lacks flexibility. Considering the neighborhood size varies with different query tokens, number of tokens that benefit for different query tokens’ local semantic representation are different. Moreover, their mask matrices should match different attention heads and layers in MANs.

We propose Dynamic Mask Attention Network (DMAN) that replaces the static mask matrix. Incorporating query tokens, relative distance, attention head and layer, we build a dynamic mask function which replaces the hard $0/1$ mask gate in Equation 11 with a soft one through sigmoid activation function in Equation 12.

$\displaystyle\text{DM}_{i}^{l}[t,s]=\sigma\Big{(}h^{l}_{t}W^{l}+P^{l}_{t-s}+U^{l}_{i}\Big{)}{}$

(12)

where $s,t$ are the positions of query and key, $i$ is the attention head, $l$ is the layer. ${P}^{l}_{t-s}$ is parameterized scalar for the positions $t$ and $s$, $U^{l}_{i}$ is for the $i$-th head, and $W^{l}\in\mathbb{R}^{d\times 1}$. $W^{l}$, ${P}^{l}_{t-s}$ and $U^{l}_{i}$ are trainable parameters.

Until here, we have three sub-networks of MANs, namely, SAN, FFN and DMAN. SAN that does not mask any tokens and specializes in global semantic modeling. FFN that masks all tokens except itself and focuses on self-processing. DMAN masks the tokens not in neighborhood and is able to model local structure more effectively.

Transformer is composed of SAN and FFN that achieves positive results in various NLP tasks, the stacking method of Transformer inspires us to stack DMAN, SAN and FFN to incorporate their advantages. We insert DMAN in the manner of DMAN$\to$SAN$\to$FFN, which is shown in Figure 2. With this architecture, we first model the localness then globalness, and take the step for self-evolution in the end.

Here, we investigate different collaboration mechanisms of the elements in MANs. Under our design principles, there are three elements: FFN, SAN, and DMAN. For the convenience of comparison, we take FFN as the last component in the sequential layered structure. We try different collaboration methods and test them on IWSLT2014 German-to-English (De-En). The results are shown in the Table 3. We conclude that:

1. 1.

   Our proposed $C$#5 achieves the best performance that verify the effectiveness of our proposed sequential layered structure.

2. 2.

   All of $C$#3, $C$#4 and $C$#5 outperform $C$#1 and $C$#2, and the least improvement in BLEU is 0.2. This shows that no matter what collaboration method, models with the participation of DMAN perform better than models without DMAN, which validates the capability of DMAN.

3. 3.

   Both $C$#5 and $C$#4 are better than $C$#3 and $C$#2. This indicates that models without DMAN or SAN are not comparable to models with all three modules. This shows that DMAN and SAN have their own strengths, namely, localness modeling and globalness modeling, and are able to make up for each other’s defects through collaboration.

4. 4.

   $C$#5 is better than $C$#4. This indicates that first modeling the localness and then globalness would be better than the inverse order.

In this section, we compare the performance of Static Mask Attention Network (SMAN) and Dynamic Mask Attention Network (DMAN). Both of them follow the collaboration strategy of DMAN(SMAN)$\to$SAN$\to$FFN. In SMAN, we set a fixed mask boundary which has been determined in advance following Equation 11. Empirically, we propose two static mask strategies: (a) SMAN₁, the boundary $b$ depends on sentence length $L$, $b=\sqrt{L}/2$; (b) SMAN₂, $b$ is set to 4, which is chosen from 2, 4, 6, 8 through validation.

The results in IWSLT2014 De-En are shown in Table 4. The performance of SMAN₁ and SMAN₂ are very close. They both outperform the Transformer but fall behind our proposed DMAN. This indicates that our proposed DMAN is superior to SMAN. SMAN fails to manage various neighborhood for different query tokens, but DMAN can model localness with more flexibility according to these factors.

In this section, we analyse the behavior of DMAN and SAN in localness modeling through attention scores in Equation 4. To quantify the role of neighbors in semantic modeling, we compute the sum of attention scores within some particular window size. Generally, if the attention score from $a$ to $c$ is bigger than $b$ to $c$, we consider that $a$ contributes more to the semantic modeling of $c$ compared to $b$, in other words, model utilizes more information of $a$ than $b$ to learn the semantic representation of $c$. Therefore, larger attention scores mean that model utilizes more information of the corresponding tokens to learn the semantic representation of query token.

For each sentence in dataset $X_{i}=(x_{i,1},\cdots,$ $x_{i,T_{i}})\in\mathcal{D}$, we utilize $\bar{s}^{l}_{i,\emph{DMAN}}$ and $\bar{s}^{l}_{i,\emph{SAN}}\in\mathbb{R}^{T_{i}\times T_{i}}$ to denote the average attention scores $\mathcal{S}_{M}(Q,K)$ in Equation 4 across different heads in the $l$-th layer for DMAN and SAN, respectively. We sum the attention scores of these tokens $x_{i,k}$ within the window size $w$ of the query $x_{i,j}$ in the $l$-th layer, and average the sum across $X_{i}$ and dataset $\mathcal{D}$ following Equation 13.

$$\begin{gathered}attn\_s_{w,l,*}=\frac{1}{|\mathcal{D}|}\sum_{X_{i}\in\mathcal{D}}\frac{1}{T_{i}}\sum_{x_{i,j}\in X_{i}}\sum_{|k-j|\leq w}\bar{s}^{l}_{i,*}\big{[}j,k\big{]}~{}\end{gathered}$$

(13)

where $*\in\{\emph{DMAN},\emph{SAN}\}$, and $\bar{s}^{l}_{i,*}\big{[}j,k\big{]}$ is the value of the $j$-th row and $k$-th column of $\bar{s}^{l}_{i,*}$. $attn\_s_{w,l,*}$ measures the overall contribution of these neighbor tokens within the window size $w$ to the query tokens’ semantic modeling. We take $\mathcal{D}$ as the test set of IWSLT14 De-En and compute $attn\_s_{w,l,*}$ with $w=1,2,4$ and $l=1,3,6$.

The result is shown in Table 5. We see that in layer#1, #3 and #6, the sum attention scores of DMAN within the window size $2$ are 50% more than those of SAN, especially in layer#1 where the gap is as much as five times between SAN and DMAN. This phenomenon validates that the attention scores of DMAN in neighbors are larger than those of SAN, thus DMAN is more specialized in localness modeling than SAN.