MLM — Masked Language Modeling: Complete Guide

Summary

MLM (Masked Language Modeling) is a fundamental self-supervised technique in natural language processing (NLP). Popularized by Google’s BERT model in 2018, this approach revolutionizes the way machines learn to understand human language. Unlike traditional language models that read text in a single direction, MLM allows the model to access the complete context — left and right — to predict deliberately masked words.

This complete guide explores the mathematical foundations, pedagogical intuition, practical implementation in PyTorch, essential hyperparameters, as well as the advantages, limitations, and concrete use cases of masked language modeling. If you’re looking to understand how BERT and its derivatives acquire their deep understanding of language, this guide is for you.

Mathematical Principle

Formalizing the masking

Given an input sequence $x = [x_1, x_2, \ldots, x_n]$ composed of $n$ tokens, the principle of MLM consists of randomly masking a subset of these tokens according to a precise strategy. The standard masking rate is 15%, but this operation is not a simple replacement with a [MASK] token.

More precisely, for each selected token among the 15%:

• 80% of the time: the token is replaced by the special [MASK] token
• 10% of the time: the token is replaced by a random token from the vocabulary
• 10% of the time: the token is kept unchanged

This strategy, called three-way masking, is crucial. If we systematically replaced all masked tokens with [MASK], the model would never see this token during fine-tuning (since [MASK] does not appear in real data). By keeping some tokens intact and introducing random noise, we create a more robust and more generalizable version of the learning task.

Loss function

The model learns by maximizing the likelihood of the masked tokens. The loss function is formally written as:

$$\mathcal{L} = -\mathbb{E}\left[\sum_{i \in \mathcal{M}} \log P(x_i \mid x_{\setminus \mathcal{M}})\right]$$

where:

• $\mathcal{M}$ denotes the set of masked token indices
• $x_{\setminus \mathcal{M}}$ represents the observed sequence (all tokens except the masked ones)
• $P(x_i \mid x_{\setminus \mathcal{M}})$ is the probability that the model assigns to the original token $x_i$ given all other visible tokens

In practice, this loss is computed via cross-entropy between the predicted distribution and the true labels, averaged over all masked positions in the training batch.

Bidirectionality: the MLM revolution

This is the fundamental difference from unidirectional language models like GPT. A classical language model predicts the next token $P(x_t \mid x_1, \ldots, x_{t-1})$ by only seeing the left context. This approach is natural for text generation, but it limits semantic understanding: the model cannot use clues that appear after the word it is analyzing.

MLM, on the other hand, is inherently bidirectional. When it needs to predict a masked token at position $i$, it has access to the entire context — the tokens on the left and on the right. This two-directional viewing capability allows the model to learn much richer and more nuanced representations.

Illustrative example: Consider the sentence “The banker went to the ___ to withdraw money.” A classical unidirectional model can only predict the missing word based on what precedes it. MLM, however, would see the entire sentence with a gap, and use both “banker” (before the gap) and “withdraw money” (after the gap) to deduce that the missing word is most likely “bank.” It is this global, contextual vision that makes MLM so powerful.

Intuition

Imagine a student preparing for a French exam by practicing with cloze tests. They’re given a paragraph with certain words erased, and they must guess the missing words using the rest of the text. The more they practice, the finer their understanding of grammar, vocabulary, and the semantic relationships between words becomes.

This is exactly what MLM does. Each training sentence becomes a large-scale “fill-in-the-blank” exercise:

“The [MASK] of the cat is soft and warm.”

To guess the missing word, the model observes “of,” “the cat,” “is,” “soft,” “warm.” It progressively learns that the missing word must be a noun, probably a physical attribute of the cat, and eventually proposes “fur.” By repeating this exercise billions of times on varied corpora, the model acquires an extremely rich contextual representation of each word.

This approach is remarkably similar to the way human beings learn their mother tongue. A child does not learn words in isolation: they hear them in complete sentences, with rich context that helps them guess the meaning of words they don’t yet know. MLM reproduces this mechanism at computational scale.

Another useful parallel is crossword puzzles. To guess a word in a grid, you use all available clues horizontally and vertically. In the same way, MLM uses all available context — before and after the gap — to make its prediction. This bidirectional approach is what fundamentally distinguishes MLM from causal language models like GPT.

Python Implementation with PyTorch

Model architecture

We will implement a simplified BERT-like model composed of the following elements:

1. Token embeddings: transformation of tokens into dense vectors
2. Position embeddings: injection of positional information
3. Segment embeddings: distinction between sentences (for pair tasks)
4. Transformer layers (encoder): bidirectional context processing
5. MLM Head: masked token prediction layer

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from torch.utils.data import Dataset, DataLoader
import random

# ============================================================
# 1. Token masking
# ============================================================

def mask_tokens(inputs, vocab_size, tokenizer_pad_id, mask_token_id,
                mask_ratio=0.15, device='cpu'):
    """
    Applies strategic MLM masking:
    - 80% of selected tokens -> replaced by [MASK]
    - 10% of selected tokens -> replaced by a random token
    - 10% of selected tokens -> kept unchanged

    Args:
        inputs: tensor of tokens [batch_size, seq_len]
        vocab_size: size of the vocabulary
        tokenizer_pad_id: ID of the padding token
        mask_token_id: ID of the [MASK] token
        mask_ratio: proportion of tokens to mask (default: 0.15)

    Returns:
        masked_inputs: tensor with masked tokens
        labels: tensor of labels (=-100 except at masked positions)
    """
    labels = inputs.clone()
    masked_inputs = inputs.clone()

    # Create a probability matrix for each token
    prob_matrix = torch.full(inputs.shape, mask_ratio, device=device)

    # Do not mask special tokens [PAD], [CLS], [SEP]
    prob_matrix[inputs == tokenizer_pad_id] = 0.0

    # Random selection of tokens to mask
    mask_selector = torch.bernoulli(prob_matrix).bool()

    # Initialize labels
    labels[~mask_selector] = -100  # -100 is ignored by CrossEntropyLoss

    # Apply the three-way masking strategy
    for i in range(inputs.size(0)):
        for j in range(inputs.size(1)):
            if mask_selector[i, j]:
                rand = random.random()
                if rand < 0.8:
                    # 80%: replace with [MASK]
                    masked_inputs[i, j] = mask_token_id
                elif rand < 0.9:
                    # 10%: replace with a random token
                    random_token = random.randint(0, vocab_size - 1)
                    masked_inputs[i, j] = random_token
                else:
                    # 10%: keep unchanged
                    masked_inputs[i, j] = inputs[i, j]

    return masked_inputs, labels


# ============================================================
# 2. Positional embeddings
# ============================================================

class PositionalEmbedding(nn.Module):
    """Sinusoidal positional embeddings, as in the original attention paper."""

    def __init__(self, d_model, max_len=512):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2, dtype=torch.float)
                             * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)  # [1, max_len, d_model]
        self.register_buffer('pe', pe)

    def forward(self, x):
        """x: [batch_size, seq_len, d_model]"""
        return x + self.pe[:, :x.size(1), :]


# ============================================================
# 3. Simplified BERT-like model
# ============================================================

class SimpleBERTMLM(nn.Module):
    """
    Simplified BERT-like model for MLM.
    Includes embeddings, Transformer encoder, and prediction head.
    """

    def __init__(self, vocab_size, hidden_size=768, num_layers=6,
                 num_heads=12, ff_size=3072, max_seq_len=512,
                 dropout=0.1):
        super().__init__()

        self.vocab_size = vocab_size
        self.hidden_size = hidden_size

        # Embeddings
        self.token_embedding = nn.Embedding(vocab_size, hidden_size)
        self.position_embedding = PositionalEmbedding(hidden_size, max_seq_len)
        self.segment_embedding = nn.Embedding(2, hidden_size)
        self.layer_norm = nn.LayerNorm(hidden_size)
        self.dropout = nn.Dropout(dropout)

        # Transformer encoder (stacking layers)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_size,
            nhead=num_heads,
            dim_feedforward=ff_size,
            dropout=dropout,
            batch_first=True,
            activation='gelu'
        )
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer, num_layers=num_layers
        )

        # MLM head: projection to vocabulary space
        self.mlm_head = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.GELU(),
            nn.LayerNorm(hidden_size),
            nn.Linear(hidden_size, vocab_size)
        )

        # Weight initialization
        self._init_weights()

    def _init_weights(self):
        """Initialization inspired by BERT."""
        for module in self.modules():
            if isinstance(module, nn.Linear):
                module.weight.data.normal_(mean=0.0, std=0.02)
                if module.bias is not None:
                    module.bias.data.zero_()
            elif isinstance(module, nn.Embedding):
                module.weight.data.normal_(mean=0.0, std=0.02)
            elif isinstance(module, nn.LayerNorm):
                module.bias.data.zero_()
                module.weight.data.fill_(1.0)

    def forward(self, input_ids, segment_ids=None, attention_mask=None):
        """
        Forward pass of the MLM model.

        Args:
            input_ids: [batch_size, seq_len] — token indices
            segment_ids: [batch_size, seq_len] — segment IDs (optional)
            attention_mask: [batch_size, seq_len] — attention mask (1=visible, 0=masked)

        Returns:
            predictions: [batch_size, seq_len, vocab_size] — scores for each token
        """
        batch_size, seq_len = input_ids.size()

        # 1. Token embeddings
        x = self.token_embedding(input_ids)

        # 2. Positional embeddings
        x = self.position_embedding(x)

        # 3. Segment embeddings
        if segment_ids is not None:
            x = x + self.segment_embedding(segment_ids)

        # 4. Normalization and dropout
        x = self.dropout(self.layer_norm(x))

        # 5. Passage through the Transformer encoder
        if attention_mask is not None:
            key_padding_mask = (attention_mask == 0)
        else:
            key_padding_mask = None

        encoded = self.transformer_encoder(
            x, mask=None, src_key_padding_mask=key_padding_mask
        )

        # 6. MLM prediction head
        predictions = self.mlm_head(encoded)

        return predictions


# ============================================================
# 4. Custom dataset
# ============================================================

class MLMDataset(Dataset):
    """Dataset for MLM training from tokenized sentences."""

    def __init__(self, texts, tokenizer, max_len=128):
        self.texts = texts
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]

        # Tokenization with special tokens [CLS] and [SEP]
        tokens = self.tokenizer.tokenize(text)
        tokens = ['[CLS]'] + tokens[:self.max_len - 2] + ['[SEP]']

        # Convert to indices
        input_ids = self.tokenizer.convert_tokens_to_ids(tokens)

        # Padding
        padding_len = self.max_len - len(input_ids)
        input_ids = input_ids + [self.tokenizer.pad_token_id] * padding_len

        # Attention mask
        attention_mask = [1] * (len(input_ids) - padding_len) + [0] * padding_len

        return {
            'input_ids': torch.tensor(input_ids, dtype=torch.long),
            'attention_mask': torch.tensor(attention_mask, dtype=torch.long)
        }


# ============================================================
# 5. Training loop
# ============================================================

def train_mlm_model(model, dataloader, tokenizer, device='cuda',
                    num_epochs=10, learning_rate=5e-5, gradient_clip=1.0):
    """
    Complete training loop for the MLM model.
    """
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, eps=1e-6)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=num_epochs * len(dataloader)
    )
    criterion = nn.CrossEntropyLoss(ignore_index=-100)

    model.train()
    model.to(device)

    for epoch in range(num_epochs):
        total_loss = 0.0
        total_tokens_predicted = 0
        total_correct = 0

        for batch_idx, batch in enumerate(dataloader):
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)

            # Apply MLM masking
            masked_inputs, labels = mask_tokens(
                inputs=input_ids,
                vocab_size=model.vocab_size,
                tokenizer_pad_id=tokenizer.pad_token_id,
                mask_token_id=tokenizer.mask_token_id,
                mask_ratio=0.15,
                device=device
            )

            # Forward pass
            predictions = model(masked_inputs, attention_mask=attention_mask)

            # Compute loss
            loss = criterion(
                predictions.view(-1, model.vocab_size),
                labels.view(-1)
            )

            # Backpropagation
            optimizer.zero_grad()
            loss.backward()

            # Gradient clipping
            torch.nn.utils.clip_grad_norm_(model.parameters(), gradient_clip)

            optimizer.step()
            scheduler.step()

            total_loss += loss.item()

            # Compute accuracy
            pred_labels = predictions.view(-1, model.vocab_size).argmax(dim=-1)
            true_labels = labels.view(-1)
            mask_positions = true_labels != -100

            total_tokens_predicted += mask_positions.sum().item()
            total_correct += (
                pred_labels[mask_positions] == true_labels[mask_positions]
            ).sum().item()

            if (batch_idx + 1) % 50 == 0:
                avg_loss = total_loss / (batch_idx + 1)
                accuracy = total_correct / max(total_tokens_predicted, 1)
                print(f"Epoch {epoch+1}/{num_epochs} | "
                      f"Batch {batch_idx+1}/{len(dataloader)} | "
                      f"Loss: {avg_loss:.4f} | "
                      f"Accuracy: {accuracy:.4f}")

        epoch_loss = total_loss / len(dataloader)
        epoch_accuracy = total_correct / max(total_tokens_predicted, 1)
        print(f"\n=== End of epoch {epoch+1}/{num_epochs} ===")
        print(f"Average loss: {epoch_loss:.4f}")
        print(f"MLM accuracy: {epoch_accuracy:.4f}\n")

    return model


# ============================================================
# 6. Usage example
# ============================================================

if __name__ == "__main__":
    VOCAB_SIZE = 30000
    HIDDEN_SIZE = 256
    NUM_LAYERS = 4
    NUM_HEADS = 8
    FF_SIZE = 512
    MAX_SEQ_LEN = 128
    BATCH_SIZE = 32
    NUM_EPOCHS = 5
    LEARNING_RATE = 3e-4

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print("Device:", device)
    print(f"Vocabulary: {VOCAB_SIZE}")
    print(f"Layers: {NUM_LAYERS}")
    print(f"Max length: {MAX_SEQ_LEN}")

    model = SimpleBERTMLM(
        vocab_size=VOCAB_SIZE,
        hidden_size=HIDDEN_SIZE,
        num_layers=NUM_LAYERS,
        num_heads=NUM_HEADS,
        ff_size=FF_SIZE,
        max_seq_len=MAX_SEQ_LEN,
        dropout=0.1
    )

    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Total parameters: {total_params:,}")
    print(f"Trainable parameters: {trainable_params:,}")

Key Hyperparameters

The choice of hyperparameters profoundly determines the quality of the MLM model. Here are the most important ones:

Other important hyperparameters:

• num_heads (12 for base, 16 for large): number of multi-head attention heads. Each head learns a different type of contextual relationship.
• ff_size: internal feed-forward layer dimension (generally 4 × hidden_size).
• dropout (0.1): dropout probability applied to embeddings and feed-forward layers. Prevents overfitting.
• learning_rate (2×10⁻⁵ to 5×10⁻⁵ for fine-tuning, 1×10⁻⁴ to 3×10⁻⁴ for pre-training): learning rate. Warm-up (progressive increase during the first steps) is essential for stability.
• batch_size: generally between 256 and 8,192 sequences for pre-training. Larger batches require a higher learning rate.
• gradient_clip (1.0): gradient clipping threshold. Prevents gradient explosions, which are frequent in large Transformer models.

Advantages and Limitations

Advantages

1. Rich contextual representations: Unlike Word2Vec or GloVe which produce static embeddings (one word = one fixed vector), MLM generates dynamic representations that vary depending on context. The word “bank” will have a different representation in “the river bank” and “the central bank.”
2. Bidirectional understanding: The model accesses the complete context of each sentence, allowing it to capture complex syntactic and semantic relationships that would be inaccessible to a unidirectional model.
3. Efficient pre-training: MLM exploits massive unlabeled text corpora (Wikipedia, books, web articles) to learn generalizable representations transferable to virtually any downstream task.
4. Powerful transfer learning: Once pre-trained, an MLM model can be fine-tuned with limited labeled data for specific tasks such as classification, named entity recognition, or question answering.
5. Noise robustness: The three-way masking strategy (80/10/10) makes the model resistant to perturbations and improves its generalization ability.

Limitations

1. Unsuitability for generation: During pre-training, the model always sees [MASK] as input, but this token does not appear in real data. This pre-training/fine-tuning mismatch makes MLM poor for text generation, which is why causal models (GPT) are used for this task.
2. Conditional independence: MLM predicts each masked token independently of other masked tokens. It does not model interactions between multiple simultaneously missing words. This limitation led to the development of iterative MLM and approaches like ELECTRA.
3. Computational cost: Training a large MLM model requires dozens or hundreds of GPUs for weeks, with significant energy consumption.
4. Sensitivity to masking: The choice of masking ratio (15%) and strategy (80/10/10) is an empirical compromise. Very structured data (code, mathematical formulas) may require adapted masking strategies.
5. Bias in data: The model learns the biases present in the training corpus. If the data contains stereotypes, the model will reproduce them — a major challenge in contemporary NLP.

4 Concrete Use Cases

Use Case #1 — Text Classification

MLM provides an exceptional foundation for document classification. A model like BERT, pre-trained via MLM on Wikipedia and BookCorpus, can be fine-tuned with only a few thousand examples to classify film reviews (positive/negative sentiment), legal articles by category, or customer support tickets by priority. The model’s [CLS] layer, which aggregates information from the entire sequence, serves as a global document representation for classification.

Use Case #2 — Question Answering

Question answering systems rely heavily on MLM. Given a context (a paragraph) and a question, a fine-tuned BERT model is able to identify the text span (start and end) that contains the answer. This capability stems directly from the bidirectional understanding acquired during MLM pre-training, which allows the model to link elements of the question to corresponding elements in the context, regardless of their order of appearance.

Use Case #3 — Named Entity Recognition (NER)

Named entity recognition — identifying people, organizations, locations, dates in text — benefits particularly from MLM. Bidirectional contextual understanding allows distinguishing “Paris” (city) in “he lives in Paris” from “Paris” (mythology) in “Paris abducted Helen.” Variants like CamemBERT (French), BioBERT (biomedical), or LegalBERT (legal) are pre-trained via MLM on specialized corpora to maximize performance in their respective domains.

Use Case #4 — Semantic Search and Search Engines

Modern search engines use MLM pre-trained models to improve query understanding. Rather than searching for exact keyword matches, the MLM model encodes the query and documents into a common semantic space, allowing the retrieval of relevant documents even when they don’t use the same terms as the query. This is the principle behind Google Search, which integrated BERT in 2019 to better understand natural user queries.

See Also

• Discover ‘Factor Shuffle’ in Python: Complete Guide to Efficiently Reorganize Your Collections
• Discovering Lexicographic Neighbors in Python: Complete Guide for Developers