Mask Attention Networks: Rethinking and Strengthen Transformer (2021)

Abstract: Transformer is an attention-based neural network, which consists of two sublayers, namely, Self-Attention Network (SAN) and Feed-Forward Network (FFN). Existing research explores to enhance the two sublayers separately to improve the capability of Transformer for text representation. In this paper, we present a novel understanding of SAN and FFN as Mask Attention Networks (MANs) and show that they are two special cases of MANs with static mask matrices. However, their static mask matrices limit the capability for localness modeling in text representation learning. We therefore introduce a new layer named dynamic mask attention network (DMAN) with a learnable mask matrix which is able to model localness adaptively. To incorporate advantages of DMAN, SAN, and FFN, we propose a sequential layered structure to combine the three types of layers. Extensive experiments on various tasks, including neural machine translation and text summarization demonstrate that our model outperforms the original Transformer.

## introduction recently, transformer @xcite has been widely applied in various natural language processing tasks, such as neural machine translation @xcite ) and text summarization @xcite . to further improve the performance of the text representation, transformer-based variants have attracted a lot of attention @xcite @xcite . each building block of transformer has two sublayers: self-attention network (san) and feed-forward network (ffn). @xcite presents an extension to san which incorporates the relative positional information for the sequence. @xcite proposes attention span to control the maximum context size used in san and scales transformer to long-range (∼ 8192 tokens) language modeling. recently, some works targeting on ffn have been proposed. @xcite gives a new understanding of transformer from a multi-particle dynamic system point of view and designs a macaron architecture following strang-marchuk splitting scheme. @xcite regards the ffn as the persistent memory in san to augment san. these works focus on enhancing san or ffn, but neglect the inner relationship between san and ffn that hinders further improvement. in this work, we present a more systematic analysis for both san and ffn to reveal their connections. we introduce mask attention networks(mans), in which each network has a mask matrix that element-wise multiplies a key-query attention matrix. we show that san and ffn are two special cases in mans with static mask matrices. the mask matrix of san is an all-ones matrix, while that of ffn is an identity matrix, which is shown as (a) and (c) in figure 1 . since the mask matrix of san has no restriction on relationship modeling with other tokens, san is expert in longrange dependency modeling and capture the global semantics. in contrast, mask of ffn disables it to perceive the information of other tokens and forces it into self-evolution. we believe that these two specialties endowed by two mask matrices make the success of transformer in text representation. although positive results of transformer have been reported, recent works @xcite @xcite have shown that modeling localness would further improve the performance through experiments. we argue that deficiency of transformer in local structure modeling is caused by the attention computation with static mask matrix. in the framework of mans, we find a problem that irrelevant tokens with overlapping neighbors incorrectly attend to each other with relatively large attention scores. for example "a black dog jump to catch the frisbee", though "catch" and "black" are neither relevant nor neighbors, for the reason that both of them are highly related to their common neighbor "dog" in attention, we demonstrate that the attention score from "catch" to "black" would be large, which also decreases the attention score from "catch" to "frisbee". the issue in self-attention not only introduces noise to the semantic modeling, but also mislead query tokens to overlook these neighbor tokens. this reveals that self-attention is insufficient in localness modeling and inspires us to mask tokens that not appear in neighborhood. to strengthen transformer in localness modeling with better keeping the advantage of san and ffn, we propose a dynamic mask attention network (dman) as shown in figure 1 (b), which originates from mans. observations reveal that tokens have different ranges of neighbors, for example, that of "dog", which is also connected with "frisbee", is larger than "black" and "catch". instead of being static that determined in advance, the mask matrix of dman is dependent on the query context and relative distance. in dman, the tokens in a specific neighborhood are able to receive more attention beyond the normal self-attention mechanism. the dynamic endows dman with text representation in different scales, and we validate the superiority through experiments. in transformer @xcite , san and ffn cooperate in a sequential layered structure @xmath0. considering san, ffn, and dman all belong to mans and have different advantages in text representation, instead of directly replacing san in previous works @xcite @xcite , we propose to incorporate them with the architecture @xmath1. the main contributions of this work are threefold: ## model in § 2.1, we review the transformer architecture. we introduce mask attention networks and reformulate san and ffn to point out they are two special cases in § 2.2, and analyze their deficiency in localness modeling in § 2.3. then, in § 2.4, we describe dynamic mask attention network (dman) in detail. at last, in § 2.5, we discuss the collaboration of dman, san and ffn. ## experiments in this section, we introduce our experiments. we first describe the experimental details in § 3.1. then we show the experimental results in § 3.2. ## further analysis in this section, we conduct further analysis for our model. we first investigate stacking methods for different sublayers in § 4.1. then we compare strategies of static mask and dynamic mask in § 4.2. finally, we analyse the behavior of san and dman in localness modeling through attention scores in § 4.3. ## related work recently, there is a large body of work on improving transformer @xcite for various issues. for recurrence modeling, @xcite introduces a novel attentive recurrent network to leverage the strengths of both attention and recurrent networks. for context modeling, @xcite focuses on improving self-attention through capturing the richness of context and proposes to contextualize the transformations of the query and key layers. @xcite introduces dynamic convolutions to predict separate convolution kernels solely based on the current time-step in order to determine the importance of context elements. in order to adjust attention weights beyond san, @xcite extends the self-attention mechanism to efficiently consider representations of the relative positions or distances between sequence elements through adding a relative position embedding to the key vectors; @xcite transfers the distance between two nodes in dependency trees with a pre-defined gaussian weighting function and multiply the distance with the key-query inner product value; @xcite presents a relative position encoding scheme that adds additional relative position representation to the key-query computation. @xcite proposes a parameterized linear function over self-attention to learn the optimal attention span in order to extend significantly the maximum context size used in transformer. to merge ffn to san, @xcite proposes a new model that solely consists of attention layers and augments the self-attention layer with persistent memory vectors that play a similar role as the feedforward layer. as for the collaboration of @xcite introduces macaron layer that split the ffn into two half-steps based on strang-marchuk splitting scheme in ode. for localness modeling, yang et al. (2018) casts localness modeling as a learnable gaussian bias according to relative distance to external energy in softmax function as a new self-attention network. zhao et al. (2019) explores parallel multi-scale representation learning to capture both long-range and short-range language structures with combination of convolution and self-attention. in our work, dman, san and ffn are unified in mask attention networks, where dman is a supplement of san and ffn that specializes in localness modeling. moreover, we investigate different collaboration mechanisms. ## conclusion in this paper, we introduce mask attention networks and reformulate san and ffn to point out they are two special cases with static mask in mans. we analyze the the deficiency of san and ffn in localness modeling. dynamic mask attention network is derived from mans for better local structure modeling. considering the different specialities of san, ffn, and dman, we investigate a sequential layered structure @xmath2. compared with original transformer, our proposed model achieves better performance in neural machine translation and abstract summarization. for future work, we consider adding structure information or external knowledge, e.g., dependency tree, with mask matrices in mans.