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 [8, 4]. BERT represents a significant shift in model architecture by employing an encoder-only framework [9]. 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.

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.

Some approaches suggest the possibility of dropping certain layers during fine-tuning for efficiency [10]. 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 [8]. Similarly, the T5 model can be adapted to similar finetuning tasks by utilizing only its encoder part [1]. Besides, BERT-of-Theseus presents a novel compression approach for BERT by progressively replacing original modules with compact substitutes during pretraining [7]. This method enhances interaction between original and compact models without introducing extra loss functions, ultimately outperforming existing knowledge distillation techniques on the GLUE benchmark.

Our approach with BPDec maintains the fundamental encoder architecture of the original BERT, but supplements it with a specially designed decoder, which is used exclusively during pretraining. This additional decoder, consisting of transformer blocks similar to those in the encoder, is crafted to enhance the BERT encoder’s ability in understanding languages. By incorporating this pretraining-only 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 is strategically positioned after the encoder and just before the output layers for masked token prediction. It specifically sharpens the model’s ability to predict masked tokens, enriching the encoded representations with greater depth.

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.

In the development of our new pretraining strategy, we introduce the Gradual Unmasking Attention (GUA) within our BPDec. This pioneering modification challenges the conventional constraints of the attention mechanism used in BERT model pretraining. Our approach is based on the insight that as the model’s layers progress deeper, the resulting output features become more abstract. We hypothesize that these abstracted features at masked positions are sufficiently accurate and detailed to serve as effective hypothesized contexts. GUA progressively unveils the attention mechanism to those hypothesized contexts of the masked inputs, fostering a more nuanced and robust formulation of language understanding. This method begins with partially unmasking the attention mask in the initial layers of the decoder. For instance, in a decoder with two transformer blocks as shown in Figure 2, the Gradual Unmasking Attention (GUA) in the first decoder layer would randomly unmask 50% of the masked positions. This introduces a controlled amount of noise and hypothesized contexts. In the second GUA, the remaining 50% of the positions are unveiled, allowing the model to integrate the full context. This phased approach of exposure ensures that the model gradually adapts to the changes in the inputs without being overwhelmed by sudden floods of unmasked data. More importantly, allowing the model to progressively form a complete picture of the input data significantly improves its interpretative capabilities, enhancing its ability to generate accurate and contextually relevant predictions.

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 [11]. Additionally, echoing the compression method utilized in BERT-of-Theseus, the authors recognize that random layer replacement is an effective approach to model compression [7]. 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 comparativeness in our comparative experiments.

In our study, we conducted a comparative analysis using the BERT model, BERT model with pre-LN (pre-layer normalization) [12], and DeBERTa [8] 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 [4], 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 introduces a significant overhead comparing with other baseline models and with our BERT+BPDec models. The comparisons can be found in Table 2 This added complexity affects 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.

• •

  CoLA (Corpus of Linguistic Acceptability): Measures grammatical correctness.

• •

  MNLI-m/mm (Multi-Genre NLI - Matched/Mismatched): Assesses the ability to predict textual entailment across different genres.

• •

  MRPC (Microsoft Research Paraphrase Corpus): Focuses on identifying whether two sentences are paraphrases of each other.

• •

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

• •

  QQP (Quora Question Pairs): Focuses on determining if two questions asked on Quora are semantically equivalent.

• •

  RTE (Recognizing Textual Entailment): Sentence entailment task.

• •

  STS (Semantic Textual Similarity Benchmark): Evaluates the degree of semantic similarity between two sentences.

In our finetuning experiments, each model was trained on the designated training set and subsequently evaluated on the development set. Table 3 illustrates that under identical conditions and with the same number of pretraining samples, BERT with BPDec enhancement outperforms both BERT and DeBERTa models. 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 32% additional costs in finetuning 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 with consuming the exact same finetuning and serving costs. The results tell us that this decoder only enhancement significantly boosts performance without incurring 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, focusing 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 4, following 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, is notably stronger than the original BERT and even surpases the DeBERTa model. 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. For these experiments, we strategically selected GLUE tasks that have larger datasets to ensure robustness and credibility in our results and to fit the data representations within a single table effectively. Tasks like RTE and SST, which have relatively smaller datasets were ommited. In contrast, tasks such as MNLI, with its substantial dataset of nearly 400,000 examples, were included. This selection criteria is designed to illustrate the impact of our modifications more clearly and designed to be fitted into a single table. By focusing on tasks with larger datasets, we can provide more significant insights, but can still validate the efficacy of our approach comprehensively.

To further demonstrate the effectiveness of BPDec, we implemented and applied the same methodology to ALBERT [13], BERT-PreLN [12] and DeBERTa [8] models. For ALBERT implementation, 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, ensuring the structural parity between encoder and decoder. Experiments on the GLUE benchmark and SQuAD dataset shows consistent performance improvements across the models. For DeBERTa model, because the original model already has a similar decoder structure, we added GUA and applied random mix upon the original decoder.

These additional experiments shown in 8 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.