BPDec: Unveiling the Potential of Masked Language Modeling Decoder in BERT Model Pretraining

In the context of traditional Language Modeling, the objective is to maximize the probability of a given sequence of tokens. This can be mathematically defined as:

Masked Language Modeling (MLM) has become a fundamental technique in the field of natural language processing, particularly since the introduction of the BERT model [13, 9]. BERT represents a significant shift in model architecture by employing an encoder-only framework [1]. For autoregressive language modeling, as seen in models like GPT, the objective differs in its sequential approach to predicting each token. In autoregressive language modeling, each token $X_{i}$ is predicted based on the sequence of preceding tokens. Mathematically, this can be defined as:

The goal is to optimize the model parameters $\theta$ such that the probability of each token $X_{i}$, given its preceding context, is maximized. The encoder-only design of BERT enables bidirectional context understanding and full-text attention, unlike autoregressive decoders that are limited to unidirectional context. BERT’s training methodology evolves to masked language modeling (MLM). This approach can be mathematically represented as follows:

Here, $X_{masked}$ represents the set of tokens that have been masked, and $X_{unmasked}$ is the set of tokens that remain visible to the model. The model parameters $\theta$ are optimized to maximize the log probability of correctly predicting the masked tokens given the context of unmasked tokens. BERT’s MLM pretraining randomly masks 15% of the tokens in each sequence. Within this subset, 80% of the tokens are substituted with a [MASK] token, 10% remain as is, and the remaining 10% are replaced with a random token from the vocabulary. The model is trained to predict these tokens based on the context. This training method allows BERT to effectively grasp the bidirectional context, making it highly effective across a wide range of NLP tasks after finetuning.

Following the pretraining phase, BERT and similar models will undergo a process known as fine-tuning, which tailors them to specific tasks. In this phase, the final-layer output of the the [CLS] token, typically the first token of the sequence, is usually employed as a representative summary of the entire input sequence and is finetuned along with other model parameters to fit various downstream tasks. This fine-tuning phase is essential for adapting the generalized pretraining to more specialized applications, thereby enhancing the model’s performance on task-specific datasets.

There are variations in the finetuning process. Some approaches suggest the possibility of dropping certain layers during fine-tuning for efficiency [14]. This efficiency-focused approaches suggest the feasibility of omitting pretrained layers during fine-tuning. Notable findings in this area highlight that:

1. 1.

   Up to 50% of pretrained transformer blocks can be pruned while retaining approximately 98% efficiency in specific downstream tasks.

2. 2.

   The retention of lower layers is essential for maintaining performance in downstream applications.

3. 3.

   For certain tasks, a minimal subset of layers, as few as 3 out of 12, can sustain performance within a 1% variance of the full model’s capacity.

DeBERTa introduces an approach, where it is pretrained with an additional decoder and then this decoder is dropped during the fine-tuning phase [13]. Similarly, the T5 model can be adapted to similar finetuning tasks by utilizing only its encoder part [8]. Besides, BERT-of-Theseus presents a novel compression approach for BERT by progressively replacing original modules with compact substitutes during pretraining [12]. This method enhances interaction between original and compact models without introducing extra loss functions, ultimately outperforming existing knowledge distillation techniques on the GLUE benchmark.

DeBERTa introduces an innovative approach known as the Enhanced Masked Language Modeling Decoder (EMD), which is mentioned above. Moreover, while BERT incorporates absolute positional encodings at the input layer, DeBERTa adopts a different strategy. In DeBERTa, absolute positions are incorporated after all Transformer layers, just before the softmax layer dedicated to masked token prediction. This design choice allows the model to capture the relative positions within all Transformer layers and utilize absolute positions as supplementary information during the decoding phase. However, this disentangled attention mechanism notably introduces additional parameters and computational overhead compared to the original attention mechanism. Therefore, a direct comparison between DeBERTa and BERT with same number of layers might not be entirely equitable due to these architectural differences. As we will elaborate in the next section, our approach ensures fairness in comparison. The only difference lies in the pretraining phase compared to the original BERT, which is often just a one-time cost, whereas in practical applications, finetuning and inference are the recurring stages where computational efficiency matters more.

Our approach with BPDec preserves the core encoder structure of the original BERT while augmenting it with a specially designed decoder, exclusively active during pretraining. This additional decoder is crafted to enhance the BERT encoder’s ability in understanding and processing of language, drawing upon the strengths of the EMD approach. By incorporating this decoder, we aim to improve the overall performance of the BERT encoder without imposing significant computational costs during the subsequent phases of finetuning and deployment. The BPDec design is strategically aligned with the objectives of pretraining. This decoder is adeptly positioned right after the BERT encoder, and just before the final dense layer and softmax layer dedicated to masked token prediction. It focuses on refining the model’s ability to predict masked tokens, thereby expanding the depth and versatility of the encoded representations.

Another key aspect of our proposed methodology in BPDec involves modifying the restrictions present in the standard multi-head attention mechanism of BERT, particularly in relation to masked positions. As shown in Figure 1, the BERT architecture avoids attending to embeddings at the 15% masked positions during the multi-head attention by setting an attention mask to them during pretraining. The embeddings at masked positions are considered less informative and less accurate than those are not masked. The mechanism helps ensure that the model is not unduly influenced by those masked and noised inputs.

However, in our MLM Decoder, we try to remove the restriction in the decoder layers. This modification is rooted in the understanding that as the layers in the model become deeper, the output features become increasingly abstract and, conceivably, become sufficiently enough to guess the masked or noised information. Such information, though not fully revealed, can provide valuable context to the model. It encourages the model to formulate hypotheses and fosters a deeper alignment and understanding of the input based on a mix of noise and these formulated assumptions. However, this also brings about a drawback, which is that the model’s sudden exposure to this hypothesized information about the masked or noised input, and the noise it brings, may have adverse effects. To mitigate these risks and enhance the effects, we made two optimizations shown in Figure 2: 1) Fully disabling the attention mask only in the last several decoder layers to prevent the early involvement of hypothesized information, and 2) Progressively unmasking the attention mask across the decoder layers, starting with a randomized partial unmasking followed by a complete unmasking. This approach not only enriches the model’s interpretative capacity but also allows for a more gradual training process, leading to enhanced overall performance as shown in the experiments.

In our approach, we introduce a modest blending of the encoder’s output into the decoder’s output as shown in Figure 3. This technique serves a dual purpose. Firstly, it encourages the MLM task to not depend exclusively on the representations formed in the deeper layers of the decoder. Instead, it also takes into account the more direct and less processed information from the encoder stage. Drawing inspiration from the NEFTUNE method, which highlights the potential performance benefits of introducing controlled randomness during language model training, we adopt a similar principle in our approach [15]. Additionally, echoing the compression method utilized in BERT-of-Theseus, the authors recognize that random layer replacement is an effective approach to model compression [12]. In our case, randomly skipping decoder layers can be seen as a form of compression as well. Our experimental results reveal that using 80% of the decoder output mixed with 20% of the encoder output can significantly improve performance on downstream tasks.

Pretraining was conducted on both base and large model sizes. Additionally, we introduced a variant of the base model (base-h256) to further demonstrate the generalizability of our method on smaller models. The encoder architecture, pretraining data, pretraining tasks, and mask settings are identical with the original BERT models. Our BPDec lies in the addition of a decoder following the encoder. Details about the parameters of this decoder can be found in Table 1. This strategic augmentation aims to enhance the model’s capabilities while maintaining the foundational architecture of BERT encoder, and it also ensures a high degree of fairness in our comparative experiments.

In our study, we conducted a comparative analysis using the BERT model, BERT model with pre-LN (pre-layer normalization) [16], and DeBERTa [13] trained under identical settings as baselines. Here, we utilized the same hyperparameters, tokenizer, and pretraining dataset to isolate the impact of other factors on model performance. Our focus was solely on determining the effectiveness of the architectural change. In RoBERTa [9], authors significantly extended the pretraining duration, increasing the number of steps from 100K to 500K. This led to substantial improvements in downstream task performance. Authors also emphasized that the longest-trained model still haven’t show signs of overfitting. This observation highlights the importance of controlling for training steps and pretraining dataset when evaluating architectural modifications, as it can act as a confounding factor that influences model performance. Without careful consideration of such factors, drawing definitive conclusions about the true impact of architectural changes could be misleading.

Moreover, it is important to note that although the DeBERTa model was trained for the same number of steps as other models, its disentangled attention mechanism introduced a significant overhead comparing with other baseline models and with our BERT+BPDec models. This added complexity affected both the pretraining and fine-tuning phases, but ours BERT+BPDec models only have extra ovearhead during pretraining. The implications of this increased overhead are critical, especially when considering the balance between model performance and computational efficiency. Our analysis aims to provide insights into how these different architectures impact training and computation efficiency, while also maintaining or enhancing language processing capabilities.

We summarize the results on seven Natural Language Understanding (NLU) tasks from the GLUE benchmark. These tasks and their corresponding metrics are as follows:

• •

  AX (GLUE Diagnostic Task): Tests ability to understand linguistic phenomena such as logic, predicate-argument structure, and lexical semantics. The evaluation metric is typically overall accuracy across all diagnostic categories.

• •

  CoLA (Corpus of Linguistic Acceptability): Measures grammatical correctness. The metric used is Matthew’s correlation.

• •

  MNLI-m/mm (Multi-Genre NLI - Matched/Mismatched): Assesses the ability to predict textual entailment across different genres. Accuracy is the metric for both the matched (MNLI-m) and mismatched (MNLI-mm) sections.

• •

  MRPC (Microsoft Research Paraphrase Corpus): Focuses on identifying whether two sentences are paraphrases of each other. Accuracy is used as the main evaluation metric.

• •

  QNLI (Question NLI): Involves determining whether a context sentence contains the answer to a question. Accuracy is the evaluation metric.

• •

  QQP (Quora Question Pairs): Focuses on determining if two questions asked on Quora are semantically equivalent. The metrics for evaluation are accuracy and F1 score.

• •

  RTE (Recognizing Textual Entailment): Sentence entailment task. Accuracy is used as the metric.

• •

  STS (Semantic Textual Similarity Benchmark): Evaluates the degree of semantic similarity between two sentences. The primary metric is Pearson and Spearman correlation coefficients.

These tasks provide a broad evaluation of the model’s performance across various dimensions of NLU. The experiment results are shown in Table 2

In our finetuning experiments, each model was trained on the designated training set and subsequently evaluated on the development set. Table 2 illustrates that under identical conditions and with the same number of pretraining samples, BERT with BPDec enhancement consistently outperforms both BERT and DeBERTa models in terms of many tasks and the averaged score. Notably, MNLI, the largest dataset among these tasks, serves as a robust indicator of model performance. Our model competes closely with the similarly sized DeBERTa model, achieving slightly higher scores in the matched condition and comparable results in the mismatched one. However, it’s crucial to consider that the DeBERTa model’s encoder is approximately 17% larger in parameter size than BERT encoder, and it incurs a 32% higher cost in finetuning cost and a 29% increase in inference latency. Despite these disparities, BERT+BPDec matches or even surpasses DeBERTa in various benchmarks in base-h256, base and large level sizes. Moreover, BERT+BPDec achieves significant improvement over the original BERT at exactly the same finetuning and serving costs. The results tell us that this decoder only enhancement significantly boosts performance without incurring further additional costs. Building on this, we extended our testing to question answering tasks, further solidifying the conclusion that our work can yield substantial performance gains.

The Stanford Question Answering Dataset (SQuAD) serves as a benchmark for evaluating question answering systems in NLP. It comprises two versions: SQuAD v1, focused on answerable questions from given passages, and SQuAD v2, which includes both answerable and unanswerable questions. Performance is measured using Exact Match (EM) and F1 Score, reflecting the precision and accuracy of the model’s responses.

The experiments conducted on both the v1 and v2 tasks, as presented in Table 3, followed a similar trend to those performed on the GLUE tasks: all models were trained on the training set and evaluated on the development set. While the results for BERT+BPDec on SQuAD v1 did not achieve the highest scores in the table, its performance on SQuAD v2, particularly for the base and base-h256 size model, was notably stronger than the original BERT. Especially, when considering that BERT+BPDec competes closely with, or even surpasses, the DeBERTa model, the improvement is significant. These results underscore BPDec’s efficiency and efficacy, showcasing its ability to deliver robust performance in question answering tasks, a key area in NLP, without additional overhead.

The BPDec involves multiple modifications of architecture and randomization which introduce many parameters in the process. It is essential to clearly articulate whether each change is necessary and how much it contributes to the final result. Furthermore, understanding the influence of each major parameter on the outcome is also crucial. To further demonstrate the effectiveness of our method, we conducted several ablation experiments on selected GLUE tasks and SQuAD tasks with base-level sized models.

To further demonstrate the effectiveness of BPDec, we applied the same methodology to ALBERT [17] and BERT-PreLN [16] models. For ALBERT, the BPDec does not share parameters with the ALBERT encoder, but the decoder layers within BPDec share parameters. For BERT-preLN, we followed the encoder’s structure and used pre-LayerNorm in BPDec. Experiments on the GLUE benchmark and SQuAD dataset showed consistent performance improvements across both models.

These additional experiments shown in 7 prove the effectiveness, adaptability, and generalizability of the BPDec method, highlighting its potential for broader application in various BERT-like architectures. The consistent improvements observed across different model types suggest that BPDec’s principles can be successfully transferred to enhance the performance of a wider range of models.