Fine-Tuning BERT with Masked Language Modeling: A Comprehensive Guide

Since its introduction in 2018, BERT (Bidirectional Encoder Representations from Transformers) [1] has revolutionized the field of natural language processing (NLP). Developed by researchers at Google, BERT is a powerful language model that can be pre-trained on large amounts of unlabeled text data and then fine-tuned for a variety of downstream NLP tasks with minimal additional training.

At the core of BERT‘s success is the transformer architecture [2] and the masked language modeling (MLM) pre-training technique. In this article, we‘ll take an in-depth look at the mathematical formulation of MLM, explore variations like static vs dynamic masking, walk through code for fine-tuning a pre-trained BERT model, and analyze BERT‘s performance on benchmark datasets. Whether you‘re an NLP researcher or practitioner, read on to gain a comprehensive understanding of BERT and MLM.

A Primer on BERT and the Transformer Architecture

BERT is a deep learning model that processes a sequence of input tokens x=(x_1,…,x_n) to produce a sequence of contextualized vector representations h=(h_1,…,h_n). Unlike directional language models that read the input sequentially from left-to-right or right-to-left, BERT captures bidirectional context by attending to the entire input at once.

Internally, BERT is composed of L stacked transformer encoder layers. Each layer applies multiple self-attention heads to aggregate information from different positions and linear transformations to produce hidden states that are propagated to the next layer:

h_i^l = Transformer(h_i^(l-1))

where h_i^l is the hidden state of the ith token at the lth layer. The first layer takes input word embeddings and positional encodings. The final hidden states h^L capture rich, contextual representations that can be used for downstream tasks.

The power of the transformer lies in its self-attention mechanism [2], which allows each token to attend to all other tokens in the sequence when computing its representation. The attention weights a_ij between tokens i and j are calculated based on the similarity of their query, key, and value vectors Q, K, V:

a_ij = softmax(Q_i * K_j / sqrt(d_k))

where d_k is the dimension of the key vectors. The final output is a weighted sum of the value vectors:

Attention(Q,K,V) = sum(a_ij * V_j)

Multi-head attention applies this operation in parallel with different learned projections to capture different types of relationships between tokens.

Masked Language Modeling

The key innovation of BERT is its use of masked language modeling (MLM) for pre-training. In contrast to traditional language modeling which predicts the next token given the previous ones, MLM randomly masks out a percentage of input tokens and trains the model to predict the original vocabulary id of the masked tokens based on their context.

Formally, given an input sequence x, a random sample of tokens m
∈
(
1
,
.
.
.
,
n
)
is selected for masking with probability 0.15. The selected tokens x_m are replaced with either a [MASK] token 80% of the time, a random token 10%, or the original token 10%. The MLM objective is to predict the original token x_m from the masked input x_corrupt:

L_MLM = -sum(log p(x_m|x_corrupt))

Intuitively, by masking different tokens across many training examples, the model learns to leverage both left and right context to predict each token. This allows it to build bidirectional representations useful for numerous language understanding tasks.

Static vs Dynamic Masking

In the original BERT implementation, the set of masked tokens is chosen once in advance and kept fixed across all training epochs. However, subsequent work [3] found that dynamic masking, where a fresh set of masks is generated each time a sequence is fed to the model, can moderately improve downstream performance.

The hypothesized benefit is that dynamic masking exposes the model to more varied masks for each example, potentially enhancing its ability to learn useful representations. The tradeoff is increased training time since a new masking pattern must be generated for each input.

Fine-Tuning BERT with MLM

After pre-training on a large unlabeled corpus, the resulting BERT model can be fine-tuned on a smaller supervised dataset for a specific task. The fine-tuning process largely follows the pre-training procedure with a few key differences:

1. The model is initialized with pre-trained weights instead of random values
2. A task-specific output layer is added on top of the final encoder hidden states
3. The training objective is the supervised loss for the target task (e.g. cross-entropy for classification)
4. A smaller learning rate is used to avoid catastrophic forgetting of the pre-trained weights

Here‘s a PyTorch code snippet demonstrating the key steps of loading a pre-trained BERT model, masking tokens, and fine-tuning:

from transformers import BertTokenizer, BertForMaskedLM

tokenizer = BertTokenizer.from_pretrained(‘bert-base-uncased‘)
model = BertForMaskedLM.from_pretrained(‘bert-base-uncased‘)

def mask_tokens(inputs, tokenizer):
    """Prepare masked tokens inputs/labels for MLM."""

    # Create labels - use input ids as labels b/c we predict original tokens
    labels = inputs["input_ids"].detach().clone()

    # Mask 15% of tokens in each sequence at random
    mask = torch.rand(labels.shape) < 0.15
    labels[mask] = -100  # set labels of masked tokens to -100
    inputs["input_ids"][mask] = tokenizer.mask_token_id

    return inputs, labels

docs = [
    "Paris is the [MASK] of France.", 
    "The Eiffel [MASK] is a famous landmark.",
    "La Tour Eiffel est située à [MASK]."]

inputs = tokenizer(docs, max_length=20, padding=‘max_length‘, truncation=True, return_tensors="pt")
masked_inputs, labels = mask_tokens(inputs, tokenizer)

optim = torch.optim.AdamW(model.parameters(), lr=5e-5)

for epoch in range(100):
    optim.zero_grad()
    input_ids = masked_inputs["input_ids"]
    attention_mask = masked_inputs["attention_mask"]
    outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
    loss = outputs.loss
    loss.backward()
    optim.step()

This example showcases fine-tuning BERT on a toy dataset of three sentences about Paris. After tokenization, 15% of tokens are randomly masked, with their original ids serving as labels for the MLM loss. The model is fine-tuned for 100 epochs with a learning rate of 5e-5 using the AdamW optimizer.

Of course, real-world datasets will be much larger. Fine-tuning BERT on GPU for a few epochs can often achieve competitive performance with previous state-of-the-art models that took days or weeks to train from scratch.

Evaluating Fine-Tuned BERT Models

To assess the effectiveness of MLM fine-tuning, it‘s important to benchmark the resulting models on relevant downstream tasks. The GLUE benchmark [4] is a popular collection of 9 tasks for evaluating language understanding, including sentiment analysis, paraphrase detection, and natural language inference.

On the GLUE leaderboard, fine-tuned BERT models have achieved state-of-the-art results, demonstrating the power of transfer learning from self-supervised pre-training:

Model	MNLI	QQP	QNLI	SST-2	CoLA	STS-B	MRPC	RTE
BiLSTM+ELMo+Attn [5]	76.4	87.4	87.1	93.2	44.1	70.3	84.9	63.8
OpenAI GPT [6]	82.1	88.5	88.1	91.3	45.4	80.0	82.3	56.0
BERT-base [1]	84.6	89.2	90.5	93.5	52.1	85.8	88.9	66.4
BERT-large [1]	86.7	91.1	92.7	94.9	60.5	86.5	89.3	70.1

BERT models outperform previous baselines by a significant margin on most tasks. The gains are especially pronounced on smaller datasets like CoLA and RTE, highlighting the data efficiency of fine-tuning.

Ablation studies in the original BERT paper also demonstrate the importance of the MLM objective. A model trained only on next sentence prediction, the other pre-training task used in BERT, achieves much lower accuracy than the full model. This suggests that the bidirectional representations learned through MLM are crucial for downstream performance.

Conclusion and Future Directions

Fine-tuning with masked language modeling has proven to be a highly effective technique for adapting pre-trained language models like BERT to new domains and tasks. By learning to predict masked tokens, the model builds rich, contextualized representations that can be readily fine-tuned for state-of-the-art performance on a wide range of language understanding benchmarks.

However, scaling BERT to ever-larger model sizes and pre-training corpora remains a challenge due to the quadratic memory complexity of self-attention. Subsequent works have proposed more efficient transformer variants like Sparse Transformers [7] and Longformer [8] to tackle longer sequences. Models like RoBERTa [9], ALBERT [10], and ELECTRA [11] have also introduced improved pre-training techniques to reduce compute and achieve even better performance.

The emergence of giant language models like GPT-3 [12] and PaLM [13], with up to billions of parameters trained on massive web-scale corpora, have pushed the capabilities of language models even further. These models exhibit remarkable abilities like few-shot learning and open-ended generation that are expanding the boundaries of what‘s possible with NLP.

Nonetheless, BERT remains an important and influential model that helped popularize the pre-train/fine-tune paradigm in NLP. Mastering the skills of fine-tuning BERT equips practitioners with a powerful tool for a wide range of language tasks. And the core ideas behind BERT continue to shape the development of more advanced language models that are driving the field forward.

References:

[1] Devlin et al. "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding." arXiv (2018).

[2] Vaswani et al. "Attention is All You Need." NeurIPS (2017).

[3] Liu et al. "RoBERTa: A Robustly Optimized BERT Pretraining Approach." arXiv (2019).

[4] Wang et al. "GLUE: A Multi-Task Benchmark and Analysis Platform for
Natural Language Understanding." EMNLP (2018).

[5] Peters et al. "Deep contextualized word representations." NAACL-HLT (2018).

[6] Radford et al. "Improving Language Understanding by Generative Pre-Training." OpenAI Blog (2018).

[7] Child et al. "Generating Long Sequences with Sparse Transformers." arXiv (2019).

[8] Beltagy et al. "Longformer: The Long-Document Transformer." arXiv (2020).

[9] Liu et al. "RoBERTa: A Robustly Optimized BERT Pretraining Approach." arXiv (2019).

[10] Lan et al. "ALBERT: A Lite BERT for Self-supervised Learning of Language Representations." ICLR (2020).

[11] Clark et al. "ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators." ICLR (2020).

[12] Brown et al. "Language Models are Few-Shot Learners." NeurIPS (2020).

[13] Chowdhery et al. "PaLM: Scaling Language Modeling with Pathways." arXiv (2022).