Mask More and Mask Later: Efficient Pre-training of Masked Language Models by Disentangling the [MASK] Token

TL;DR

This paper proposes a more efficient pre-training method for masked language models by delaying the appending of [MASK] tokens to later layers, significantly reducing computational costs without sacrificing performance.

Contribution

It introduces a novel approach to disentangle [MASK] tokens from initial embeddings and append them later, enabling faster pre-training of MLMs.

Findings

Pre-training with 78% and 68% of original compute achieves comparable results.

The method outperforms RoBERTa on 6 out of 8 GLUE tasks.

Masking rate can be increased to 50% without performance loss.

Abstract

The pre-training of masked language models (MLMs) consumes massive computation to achieve good results on downstream NLP tasks, resulting in a large carbon footprint. In the vanilla MLM, the virtual tokens, [MASK]s, act as placeholders and gather the contextualized information from unmasked tokens to restore the corrupted information. It raises the question of whether we can append [MASK]s at a later layer, to reduce the sequence length for earlier layers and make the pre-training more efficient. We show: (1) [MASK]s can indeed be appended at a later layer, being disentangled from the word embedding; (2) The gathering of contextualized information from unmasked tokens can be conducted with a few layers. By further increasing the masking rate from 15% to 50%, we can pre-train RoBERTa-base and RoBERTa-large from scratch with only 78% and 68% of the original computational budget without any degradation on the GLUE benchmark. When pre-training with the original budget, our method outperforms RoBERTa for 6 out of 8 GLUE tasks, on average by 0.4%.