TLM : Token-Level Masking for Transformers (2023)

Abstract: Structured dropout approaches, such as attention dropout and DropHead, have been investigated to regularize the multi-head attention mechanism in Transformers. In this paper, we propose a new regularization scheme based on token-level rather than structure-level to reduce overfitting. Specifically, we devise a novel Token-Level Masking (TLM) training strategy for Transformers to regularize the connections of self-attention, which consists of two masking techniques that are effective and easy to implement. The underlying idea is to manipulate the connections between tokens in the multi-head attention via masking, where the networks are forced to exploit partial neighbors’ information to produce a meaningful representation. The generality and effectiveness of TLM are thoroughly evaluated via extensive experiments on 4 diversified NLP tasks across 18 datasets, including natural language understanding benchmark GLUE, ChineseGLUE, Chinese Grammatical Error Correction, and data-to-text generation. The results indicate that TLM can consistently outperform attention dropout and DropHead, e.g., it increases by 0.5 points relative to DropHead with BERT-large on GLUE. Moreover, TLM can establish a new record on the data-to-text benchmark Rotowire (18.93 BLEU). Our code will be publicly available at https://github.com/Young1993/tlm.

## introduction in recent years, a variety of pre-trained language models based on the transformer @xcite architecture have been presented, such as bert, @xcite , and t5 @xcite . these models push state-of-the-art forward in numerous nlp tasks. with the rapid growth of model parameters, deep neural networks are highly likely to encounter overfitting challenges because supervised data is usually expensive and insufficient for large language models. this problem could cause the degradation of model generalization. to address this issue, regularization methods, such as dropout @xcite and subsequent research @xcite @xcite have been developed from a structural perspective, which trained with "thinned" subnetworks. the feature of dropout is that it randomly drops units from the neural networks during training, preventing units from co-adapting too much. to further mitigate overfitting for transformers, structured methods such as drophead @xcite are proposed to drop entire attention heads in the attention mechanism with the purpose of pre-venting a small subset of heads from dominating the whole model. dropping entire attention heads may result in losing a significant amount of feature information. attention dropout is the application of the dropout technique to the attention mechanism. it arbitrarily drops some attention weights in the matrix calculation of self-attention. however, experiments in drophead and our preliminary trials (shown in in this work, we introduce a novel regularization scheme based on token-level instead of a structural perspective. this method, token-level masking (tlm), is a training technique to regularize the connections among tokens during the attention calculation in each layer of transformer blocks. specifically, considering the example shown in fig. 1 and 2, tlm contains two techniques: 1) siblingsmasking. the first step of this method is that we use random function 1 to select a percentage of the tokens in the k-th layer, e.g., the masked token 'i' (the gray block in fig. 2 ) is excluded from the calculation of attention weights among the sibling tokens but copies itself, and its neighboring tokens considers other siblings in addition to 'i'. then, we feed the feed-forward network with the attention output to obtain the new hidden state as the input of the next layer. 2) for self-masking, we borrow the idea from cbow @xcite where its attention score is entirely contributed by others. the difference with siblings-masking is the masked token 'i' is forbidden to attend to attention computation. in the training phase, we randomly invoke one of two masking strategies at each batch with a 50-50 chance 2 . in this manner, the networks are forced to utilize partial neighbors' attention information, not the whole (i.e. the connections between the masked tokens and their neighboring tokens are invalid, which are implemented by assigning a large negative number in the matrix of attention weights). this scheme introduces a bottleneck that the nets should work hard to become robust and produce a meaningful representation. to confirm the effectiveness of our approach, we conducted extensive experiments on 18 popular datasets. the tasks range from english natural language understanding benchmark glue, our main contributions are summarized as follows: the experiments on 18 popular datasets indicate that tlm can lead to consistency improvements compared to the strong baselines. attention mask. attention mask in transformers, which is originally used to ignore those invaluable tokens, called padding tokens, in order to ensure that all the sequences are the equal length for each batch; or to prevent the networks from seeing the tokens' subsequent position in the auto-regressive decoder, such as unilm @xcite , opt @xcite , and llama @xcite . in this work, we also employ the attention mask with the purpose of mitigating overfitting, i.e., we remove some attention links among tokens during training and keep the same as the standard transformer in inference. regularization. common regularization methods include l1 and l2 regularization, dropout, early stopping, and data augmentation. dropout @xcite is a very popular training technique that aims to avoid overfitting by arbitrarily dropping some units. layerdrop @xcite randomly drops some entire substructures or components. data dropout @xcite is performed on the input data level as data augmentation. compared to the previous methods, our approach aims to carefully control the connections between tokens in the multi-head attention mechanism. ## approach in this section, we first review the self-attention computing workflow of the transformer and then describe our tlm in more detail. ## experiments to verify the generality and effectiveness of our proposed tlm, we perform extensive experiments on a wide variety of tasks with 18 benchmark datasets, including the english natural language understanding benchmark glue (10 datasets), chineseglue (6 datasets), chinese grammatical error correction (1 dataset), and data-to-text generation (1 dataset). note that our method can be utilized both in the pre-training and fine-tuning, but we only estimate tlm during the fine-tuning in this work due to limited computational resources. in the following, we present the key findings, and more details can be found in the appendix.