Masked Language Modeling

This note revisits Masked Language Modeling (MLM) — already introduced as BERT's pretraining objective in BERT — specifically as a canonical example of the pretext-task pattern from this category, worth understanding in that broader context.

MLM as a Pretext Task, Precisely

Recall the mechanics: randomly mask roughly 15% of input tokens, then train the model to predict the original tokens from bidirectional context. This is a textbook instance of the pretext-task pattern from Pretext Tasks: the "label" (the masked-out word) is generated automatically from unlabeled raw text — no human ever needs to annotate anything — and the pretext task itself (filling in blanks) isn't the actual end goal; the rich, bidirectional contextual representations the model must learn to solve it well are what actually get reused for downstream tasks.

Why Masking, Specifically, Forces Useful Learning

To correctly fill in a masked word, the model must integrate information from the entire surrounding sentence — grammar, semantics, sometimes world knowledge. A weaker pretext task (like just predicting whether a sentence is grammatically well-formed) would provide a much less rich training signal, since it doesn't require the same depth of contextual synthesis at every single masked position.

Code — The Core Masking Mechanic

import torch
import random

def mask_tokens(input_ids, mask_token_id, vocab_size, mask_prob=0.15):
    labels = input_ids.clone()
    mask_positions = torch.rand(input_ids.shape) < mask_prob

    labels[~mask_positions] = -100   # -100 is ignored by cross-entropy -- only masked positions count

    for i in range(input_ids.shape[0]):
        for j in range(input_ids.shape[1]):
            if mask_positions[i, j]:
                r = random.random()
                if r < 0.8:
                    input_ids[i, j] = mask_token_id     # 80%: replace with [MASK]
                elif r < 0.9:
                    input_ids[i, j] = random.randint(0, vocab_size - 1)   # 10%: replace with a random token
                # 10%: leave the original token unchanged -- forces the model not to over-rely on [MASK] itself

    return input_ids, labels

Notice the 80/10/10 split — this specific detail exists because [MASK] tokens never actually appear at real fine-tuning or inference time, only during pretraining. Occasionally replacing with a random token or leaving it unchanged forces the model to maintain a useful, context-sensitive representation for every token, not just ones that happen to be visibly masked.

Common Mistakes

• Masking every token equally likely to appear in downstream fine-tuning data as [MASK] — since fine-tuning data never actually contains real [MASK] tokens, the 10%/10% random-token/unchanged-token strategy specifically reduces this train/fine-tune mismatch.
• Computing loss over every position instead of only masked positions — the -100 ignore-index pattern in the code above (or an equivalent mechanism) is essential; MLM's loss should only be computed where a genuine prediction target exists.

Interview Relevance

Q: "Why does BERT's masking scheme replace masked tokens with [MASK] only 80% of the time, rather than always?" Since the special [MASK] token never appears in real downstream fine-tuning or inference data, always using it during pretraining would create a mismatch — the model might over-specialize its representations around detecting [MASK] specifically, in a way that doesn't transfer to real text without it. Occasionally substituting a random token (10%) or leaving the original token unchanged (10%) forces the model to maintain meaningful, context-aware representations for every token position, not just visibly masked ones.

Practice Question

Why is masking specifically well-suited to a bidirectional (encoder-only) architecture like BERT, rather than a causal, left-to-right decoder-only architecture?

Related DL Notes

• BERT
• Pretext Tasks
• Masked Image Modeling