Masked Language Modelling

Next-token prediction forces a decision every model designer would rather not make: each token may only see its predecessors. That is required for generation and it is a real cost for understanding — the representation of a word cannot depend on the sentence it ends up in.

MLM removes the constraint by removing generation. Corrupt some tokens, let every position attend everywhere, and predict what was taken out.

Note the sum’s range. Only masked positions contribute; the other 85% are context and produce no gradient of their own. That single fact decided the outcome.

Sample efficiency

A 512-token sequence yields 512 training signals under next-token prediction and about 77 under MLM. Same forward pass, same cost, a sixth of the supervision.

The 80/10/10 rule

BERT does not simply replace the chosen tokens with [MASK]. Of the 15% selected: 80% become [MASK], 10% become a random token, and 10% are left alone with the loss still applied.

The reason is a train/test mismatch. [MASK] exists only during pre-training; no fine-tuning input contains it. A model that keyed on the token would have learnt a feature that is never again present. The random-token share forces the model to maintain a representation of every position — it cannot assume an unmasked token is correct — and the unchanged share keeps the identity prediction in distribution.

It is a patch, and the ratios are not derived from anything.

The variants that fixed the efficiency

Span corruption (T5) masks contiguous runs and predicts them as a sequence, so a single training example can be several tokens long. This unifies the encoder-decoder setup with generation and is why T5 could do both.

ELECTRA is the sharper answer. A small generator fills the masks; the main model, a discriminator, predicts for every position whether it was replaced. The loss now covers 100% of positions rather than 15%, and ELECTRA matches RoBERTa’s quality at roughly a quarter of the compute — the same objective reformulated so that no token is wasted.

import torch
from torch import Tensor


def mask_tokens(
    tokens: Tensor, mask_id: int, vocab: int, rate: float = 0.15, special: Tensor | None = None
) -> tuple[Tensor, Tensor]:
    """Returns the corrupted input and the labels (-100 where no loss applies)."""
    probs = torch.full(tokens.shape, rate)
    if special is not None:
        probs.masked_fill_(special, 0.0)                      # never mask [CLS]/[SEP]
    chosen = torch.bernoulli(probs).bool()

    labels = tokens.masked_fill(~chosen, -100)                # loss on chosen only
    x = tokens.clone()

    # 80% [MASK]
    swap = torch.bernoulli(torch.full(tokens.shape, 0.8)).bool() & chosen
    x[swap] = mask_id

    # 10% random; the remaining 10% keep their original token.
    rand = torch.bernoulli(torch.full(tokens.shape, 0.5)).bool() & chosen & ~swap
    x[rand] = torch.randint(vocab, (int(rand.sum()),), device=tokens.device)

    return x, labels

The -100 is not arbitrary: it is cross_entropy’s default ignore_index, so unmasked positions drop out of the loss without a separate mask tensor.

Every serious embedding and reranking model is an encoder trained this way — retrieval wants a bidirectional representation of a whole passage, and there is nothing to generate. ModernBERT and its contemporaries are MLM models with the last decade of architecture work folded in, and they beat decoder-based embeddings at a fraction of the size.

The correct summary is not that MLM lost. It is that generation and representation turned out to be different jobs, and the field spent five years discovering that one objective is better at each.