Training Objective Self-Supervised

Masked Language Modeling

A self-supervised pretraining paradigm where models learn bidirectional representations by reconstructing hidden parts of corrupted inputs.

Aliases: masked autoencoding, self-supervised pretraining

In One Sentence

Instead of predicting the next token sequentially, Masked Language Modeling (MLM) hides random tokens in an input sequence and tasks the model with predicting those specific missing tokens using both left and right context.

Intuition

Standard autoregressive language models (like GPT) read text left-to-right, predicting the next word. While excellent for generation, this creates a blind spot: the model cannot look ahead. For tasks like classification, question answering, or named entity recognition, we want representations that understand the entire context simultaneously.

Masked Language Modeling solves this by turning pretraining into a "fill-in-the-blank" exercise. By corrupting a small percentage of the input tokens (typically 15%) with a special [MASK] token, we force the network to synthesize information from both directions to reconstruct the original token.

Bidirectional Context

Because the model doesn't know which tokens are masked beforehand, it must build deep, contextualized representations across all layers for every token in the sequence.

Formal Definition

Let $\mathbf{x} = (x_1, x_2, \dots, x_n)$ be an input sequence of tokens. We define a subset of indices to corrupt as $M \subset \{1, \dots, n\}$, where each index is selected independently with probability $p$ (typically $p = 0.15$).

We construct a corrupted sequence $\tilde{\mathbf{x}}$ where:

$$\tilde{x}_i = \begin{cases} \texttt{[MASK]} & \text{with probability } 0.80 \\ \text{a random token} & \text{with probability } 0.10 \\ x_i & \text{with probability } 0.10 \end{cases} \quad \forall i \in M$$

Note: This 80-10-10 rule (introduced in BERT) mitigates the discrepancy between pretraining and downstream fine-tuning, where the [MASK] token never appears.

The training objective is to minimize the negative log-likelihood of the original tokens at the corrupted positions $M$:

$$\mathcal{L}_{\text{MLM}}(\theta) = - \sum_{i \in M} \log P(x_i \mid \tilde{\mathbf{x}}; \theta)$$

where $P(x_i \mid \tilde{\mathbf{x}}; \theta)$ is modeled by a neural network (typically a Transformer encoder) followed by a softmax layer over the vocabulary.

When to use & When to avoid

Ideal For

Feature Extraction: Learning rich, transferrable representations for downstream tasks.
NLU Tasks: Classification, Named Entity Recognition (NER), and Question Answering.
Dense Retrieval: Building dual-encoders for semantic search.
High-Masking Vision: Reconstructing raw pixels (MAE) to learn robust visual features.

Avoid For

Autoregressive Generation: Open-ended text generation, chatbots, or story writing.
Low-Compute Budgets: MLM training can converge slower than causal modeling because loss is only calculated on a small fraction (e.g., 15%) of tokens per forward pass.

History & Evolution

October 2018

BERT: Bidirectional Pre-training

Devlin et al. introduce BERT, establishing Masked Language Modeling as the standard for self-supervised pretraining of deep bidirectional Transformers, sweeping NLP benchmarks.

November 2021

Masked Autoencoders (MAE) for Vision

He et al. adapt MLM to computer vision. By masking a high proportion (75%) of image patches and using an asymmetric encoder-decoder, they scale vision Transformers efficiently.

April 2023

DINOv2: Robust All-Purpose Features

Oquab et al. scale self-supervised pretraining (including masked patch modeling) on curated datasets to produce robust visual features that generalize without fine-tuning.

Interactive Walkthrough

Toggle between Text (BERT style) and Vision (MAE style) to see how masking ratios and strategies change across modalities.

Masking Probability (p) 15%

Original Sentence

The quick brown fox jumps over the lazy dog.

Masked Input to Model

PyTorch Implementation

masked_language_modeling.py PyTorch

import torch
import torch.nn as nn

class MaskedLanguageModelingLoss(nn.Module):
    def __init__(self, vocab_size, mask_token_id, mask_prob=0.15):
        super().__init__()
        self.vocab_size = vocab_size
        self.mask_token_id = mask_token_id
        self.mask_prob = mask_prob
        self.loss_fn = nn.CrossEntropyLoss(ignore_index=-100) # Only compute loss on masked positions

    def forward(self, input_ids, logits):
        # input_ids: [batch_size, seq_len]
        # logits: [batch_size, seq_len, vocab_size]
        
        # Create labels (clone input, fill unmasked with -100)
        labels = input_ids.clone()
        
        # Sample positions to mask
        probability_matrix = torch.full(labels.shape, self.mask_prob)
        masked_indices = torch.bernoulli(probability_matrix).bool()
        
        # We only compute loss on the selected indices
        labels[~masked_indices] = -100  
        
        # Flatten for loss calculation
        loss = self.loss_fn(logits.view(-1, self.vocab_size), labels.view(-1))
        return loss

Common Variants

Span Masking

Instead of masking individual tokens, contiguous spans of tokens are masked (popularized by T5). This forces the model to learn longer-range dependencies.

Asymmetric MAE

Used in computer vision. The encoder only processes visible patches, while a lightweight decoder processes both visible representations and learnable mask tokens.

Permutation LM

XLNet-style pretraining that avoids the use of explicit [MASK] tokens by permuting factorization orders, eliminating pretrain-finetune mismatch.

Representative Flash Papers

BERT

Grounding References

• BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding (arXiv:1810.04805)
• Masked Autoencoders Are Scalable Vision Learners (arXiv:2111.06377)
• DINOv2: Learning Robust Visual Features without Supervision (arXiv:2304.07193)