Representation Deficiency in Masked Language Modeling

TL;DR

It is demonstrated empirically and theoretically that MLM pretraining allocates some model dimensions exclusively for representing real tokens, resulting in a representation deficiency for real tokens and limiting the pretrained model's expressiveness when it is adapted to downstream data without $\texttt{[MASK]}$ tokens.

Abstract

Masked Language Modeling (MLM) has been one of the most prominent approaches for pretraining bidirectional text encoders due to its simplicity and effectiveness. One notable concern about MLM is that the special $\texttt{[MASK]}$ symbol causes a discrepancy between pretraining data and downstream data as it is present only in pretraining but not in fine-tuning. In this work, we offer a new perspective on the consequence of such a discrepancy: We demonstrate empirically and theoretically that MLM pretraining allocates some model dimensions exclusively for representing $\texttt{[MASK]}$ tokens, resulting in a representation deficiency for real tokens and limiting the pretrained model's expressiveness when it is adapted to downstream data without $\texttt{[MASK]}$ tokens. Motivated by the identified issue, we propose MAE-LM, which pretrains the Masked Autoencoder architecture with MLM where $\texttt{[MASK]}$ tokens are excluded from the encoder. Empirically, we show that MAE-LM improves the utilization of model dimensions for real token representations, and MAE-LM consistently outperforms MLM-pretrained models across different pretraining settings and model sizes when fine-tuned on the GLUE and SQuAD benchmarks.

Figures (6)

• Figure 1: In an MLM-pretrained model, (a) some model dimensions are exclusively used for representing [MASK] tokens, resulting in a representation deficiency for modeling inputs without [MASK], especially in deeper layers; (b) the effective rank of [MASK] token representation space increases throughout Transformer layers.
• Figure 2: Overview of MAE-LM. Masked positions are omitted from encoder inputs so that the encoder purely models real tokens. A shallow decoder takes the encoder's output representations and masked positions to predict the original tokens. After pretraining, only the encoder (but not the decoder) is fine-tuned for downstream tasks.
• Figure 3: MNLI dev set accuracy by fine-tuning intermediate MAE-LM$_{\text{base}}$ checkpoints at different time steps. We also mark the pretraining time and final performance of RoBERTa (Ours).
• Figure 4: GLUE average scores and SQuAD EM scores when different fractions of [MASK] tokens are included in the input sequences to the encoder of MAE-LM$_{\text{base}}$.
• Figure 5: (a) MAE-LM effectively closes the rank gap in vanilla MLM with inputs containing or not containing [MASK]. (b) During fine-tuning, the advantage in effective rank of MAE-LM over vanilla MLM still holds.

Theorems & Definitions (8)

• Lemma 2.1: Rank increase of [MASK] token representations in Transformer encoder
• proof
• Remark
• Theorem 2.2: Rank deficiency of real token representations
• proof
• Remark
• Theorem A.1: Rank deficiency of real token representations
• proof