NarrowBERT: Accelerating Masked Language Model Pretraining and Inference
This addresses the problem of expensive large-scale language model training for NLP researchers and practitioners, offering an incremental improvement in efficiency.
The paper tackles the high computational cost of masked language model pretraining by proposing NarrowBERT, a modified transformer encoder that sparsifies self-attention and feedforward layers to operate only on masked tokens, achieving over 2x throughput acceleration in pretraining and up to 3.5x in inference with minimal performance loss on tasks like MNLI.
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 $2\times$. 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.5\times$ 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.