N arrow BERT : Accelerating Masked Language Model Pretraining and Inference (2023)

Abstract: Large-scale language model pretraining is a very successful form of self-supervised learning in natural language processing, but it is increasingly expensive to perform as the models and pretraining corpora have become larger over time. We propose NarrowBERT, a modified transformer encoder that increases the throughput for masked language model pretraining by more than 2x. NarrowBERT sparsifies the transformer model such that the self-attention queries and feedforward layers only operate on the masked tokens of each sentence during pretraining, rather than all of the tokens as with the usual transformer encoder. We also show that NarrowBERT increases the throughput at inference time by as much as 3.5x with minimal (or no) performance degradation on sentence encoding tasks like MNLI. Finally, we examine the performance of NarrowBERT on the IMDB and Amazon reviews classification and CoNLL NER tasks and show that it is also comparable to standard BERT performance.

## introduction pretrained masked language models, such as bert @xcite , roberta @xcite , and deberta @xcite , have pushed the state-of-the-art on a wide range of downstream tasks in natural language processing. at their core is the transformer architecture @xcite that consists of interleaved self-attention and feedforward sublayers. since the former sublayer implies quadratic time complexity in the input sequence length @xcite , many have proposed methods to make the self-attention computation more efficient @xcite @xcite @xcite . in this work, we explore an orthogonal approach to efficiency: can we make masked language models efficient by reducing the length of the input sequence that each layer needs to process? in particu-lar, pretraining by masked language modeling only involves prediction of masked tokens (typically, only 15% of the input tokens; @xcite . despite this sparse pretraining objective, each transformer layer computes a representation for every token. in addition to pretraining, many downstream applications only use a single vector representation (i.e., only the [cls] token) for prediction purposes, which is much smaller than the number of input tokens (e.g., sequence classification tasks as in glue/superglue; @xcite . by narrowing the input sequence for transformer layers, we can accelerate both pretraining and inference. we present narrowbert, a new architecture that takes advantage of the sparsity in the training objective. we present two narrowbert methods in the sections that follow (figure 1 ). we provide the code to reproduce our experiments at https:// github.com/lihaoxin2020/narrowbert . the first method reduces the input sequence for the feedforward sublayers by reordering the interleaved self-attention and feedforward sublayers in the standard transformer architecture @xcite : after two standard, interleaved transformer layers, self-attention sublayers are first applied, followed only by feedforward sublayers. this way, the feedforward sublayer computations are only performed for masked tokens, resulting in a 1.3× speedup in pretraining ( §3). the second approach reduces the input length to the attention sublayers: queries are only computed for masked tokens in the attention mechanism @xcite , while the keys and values are not recomputed for non-masked tokens, which leads to a greater than 2× speedup in pretraining. we extensively evaluate our efficient pretrained models on well-established downstream tasks (e.g., @xcite we find that our modifications result in almost no drop in downstream performance, while providing substantial pretraining and inference speedups ( §3). while efficient attention variants are a promising research direction, this work presents a different and simple approach to making transformers efficient, with minimal changes in architecture. ## narrowbert in figures 1b and 1c , we illustrate two variations of narrowbert. we define some notation to describe the configuration of our models. s refers to a single self-attention layer and f refers to a single feedforward layer. the colon : refers to the 'narrowing' operation, which gathers the masked positions from the output of the previous layer. the first variation ('contextfirst' in fig. 1b ) uses attention to contextualize all-at-once at the beginning of the model. in short, the transformer layers have been rearranged to frontload the attention components. the example given in the figure specifies the model as sf{5,s}:{5,f}, which means that the input sentence is encoded by a selfattention layer, a feedforward layer, and 5 consecutive self-attention layers. at that point, the masked positions from the encoded sentence are gathered into a tensor and passed through 5 feedforward lay-ers, thereby avoiding further computations for all unmasked tokens. finally, the masked positions are unmasked and the mlm loss is computed. the second variation ('sparsequeries' in fig. 1c ) does not reorder the layers at all. instead, the sf:{5,sf} model contextualizes the input sentence in a more limited way. as shown in figure 2 , the input sentence is first contextualized by a s and a f layer, but the non-masked tokens are never contextualized again afterwards. only the masked tokens are contextualized by the remaining {5,sf} layers. since the masked tokens are only about 15% of the total sentence length, the potential speedup is ~6.6× for every feedforward or attention layer downstream of a narrowing : operation. the memory usage can also decrease by ~6.6× for those layers since the sequence length has decreased, which allows us to use larger batch sizes during training. ## experiments we focus on 2 models in our experiments: contextfirst (sfsf{10,s}:{10,f}) and sparse-queries ({1,sf}:{11,sf}, • • • , {4,sf}:{8,sf}). our narrowbert models all contain 12 selfattention and 12 feedforward layers in total, with the narrowing operation used at different points in the model. we compare narrowbert with the baseline bert model and the funnel transformer model @xcite , which is a pretrained encoder-decoder transformer model where the encoder goes through a sequence of length bottlenecks. in our experiments, we use 15% masking in masked language model (mlm) training. following @xcite , we do not use next sentence prediction as a pretraining task. we use large batch sizes and high learning rates to fully utilize gpu memory, as suggested in @xcite . batches are sized to be the largest that fit in gpu memory. we use a learning rate of 0.0005. models are trained for 70k steps, where each step contains 1728 sequences of 512 tokens, and gradient accumulation is used to accumulate the minibatches needed per step. models were trained on hosts with 8 nvidia a100 gpus. we used the hugging face implementations of the baseline bert and funnel transformer models. we pretrained the baseline bert, funnel transformer, and narrow-bert models using the same wikipedia and books corpora and total number of steps. in figure 3 , we see the evolution of the development mlm loss over the course of model training. the bert and narrowbert models all converge to similar values, with the narrowbert models reaching a slightly higher mlm loss near the end of training. spearman correlation for sts-b @xcite . for the amazon reviews corpus, we consider both the usual 5-star prediction task and the binarized (i.e., 1-2 stars versus 4-5 stars) task. in in ## discussion and conclusion we have explored two straightforward ways of exploiting the sparsity in the masked language model loss computations: rearranging the layers of the transformer encoder to allow the feedforward components to avoid computations on the non-masked positions, and sparsifying the queries in the attention mechanism to only contextualize the masked positions. the narrowbert variants can speed up training by a factor of ~2× and inference by a factor of ~3×, while maintaining very similar performance on glue, imdb, amazon, and conll ner tasks. based on the favorable trade-off between speed and performance seen in section 3, we recommend that practitioners consider using the sparsequeries narrowbert model with 2 or 3 layers before narrowing.