The Unmasking Game

A masked language model is a game. Take a sentence, hide one or more tokens, and ask the model to guess the hidden tokens from context. Train on the game long enough and the model learns the structure of the language.

Masked language modeling produced BERT, RoBERTa, and a generation of encoders. The objective differs from autoregressive generation, and the difference remains useful in systems that have moved on. The structure below starts from a naive fill-in-the-blank guesser, derives the bidirectional objective, builds the masking and training loop step by step, and ends with the measurement questions.

The naive fill-in-the-blank: counting#

The simplest guesser uses co-occurrence counts. Given a sentence with one blank, the guesser ranks candidate words by how often each candidate appears near the visible words in a training corpus. The implementation is a scoring function over a dictionary.

from collections import Counter, defaultdict
 
def naive_guess(context, corpus):
    # context: list of visible tokens with one blank slot.
    # corpus: list of token sequences.
    #
    # Score each candidate by counting the visible tokens that
    # co-occur with it in the corpus.
    visible = [t for t in context if t is not None]
    scores = Counter()
    for seq in corpus:
        for i, tok in enumerate(seq):
            if tok in visible:
                for j, other in enumerate(seq):
                    if j != i and other not in visible:
                        scores[other] += 1
    if not scores:
        return None
    return scores.most_common(1)[0][0]
 
corpus = [
    ["the", "cat", "sat", "on", "the", "mat"],
    ["the", "dog", "sat", "on", "the", "rug"],
    ["the", "cat", "slept", "under", "the", "table"],
]
 
print(naive_guess(["the", "cat", "sat", None, "the", "mat"], corpus))

The counter has two failures. Sparsity dominates: a rare but correct candidate appears in few windows, so the count is zero. Coverage fails: a word never seen in training receives no score at all. The model needs a representation of words as vectors so similarity can transfer from frequent words to rare words. That need leads to the next step.

The representation problem: words as vectors#

A masked model replaces the counter with learned embeddings. Each token gets a vector. Context positions produce vectors through attention layers. The hidden position receives a classifier over the vocabulary. Training adjusts vectors so the correct token scores highest.

def softmax(scores: list[float]) -> list[float]:
    m = max(scores)
    exps = [s - m for s in scores]
    exps = [2.718281828 ** e for e in exps]
    total = sum(exps)
    return [e / total for e in exps]
 
def cross_entropy(probs: list[float], target: int) -> float:
    import math
    return -math.log(max(probs[target], 1e-9))

The loss is cross-entropy between the predicted distribution and the true token. Self-supervision supplies the label: the hidden token is the label. Bidirectionality supplies the signal: the model sees both sides of the blank. Representation learning is the purpose: the hidden state must encode enough language structure to predict the missing piece.

The naive bidirectional variant picks the most frequent continuation from a fixed window:

def most_common_neighbour(context, corpus):
    left = None
    for i, tok in enumerate(context):
        if tok is None and i > 0:
            left = context[i - 1]
    nxt = Counter()
    for seq in corpus:
        for i, tok in enumerate(seq):
            if tok == left and i + 1 < len(seq):
                nxt[seq[i + 1]] += 1
    return nxt.most_common(1)[0][0] if nxt else None

A fixed window cannot weigh distant context against local context. Attention solves the weighing by assigning each position a learned weight. The loss still requires a masking policy, which is the next bottleneck.

The masking policy: how to hide tokens#

A masking policy decides which tokens hide, how many, and what replaces them. The standard policy in BERT is: select roughly 15% of tokens, replace 80% with a mask token, 10% with a random token, and 10% with the original token. The random replacement forces the model to use context rather than memorizing that a mask always means predict.

import random
 
MASK = "[MASK]"
 
def mask_tokens(tokens: list[str], rate: float = 0.15,
                seed: int = 0) -> tuple[list[str], list[int]]:
    rng = random.Random(seed)
    out = list(tokens)
    masked_idx = []
    for i in range(len(tokens)):
        if rng.random() < rate:
            masked_idx.append(i)
            r = rng.random()
            if r < 0.8:
                out[i] = MASK
            elif r < 0.9:
                j = rng.randrange(len(tokens))
                out[i] = tokens[j]
            else:
                out[i] = tokens[i]
    return out, masked_idx

The bottleneck is mismatch. Fine-tuning data rarely contains mask tokens. The model trains on a distribution that downstream tasks do not show. The 80/10/10 split is a mitigation, not a cure. Alternative policies mask spans, mask whole words, or mask at a higher rate for larger models. Each policy trades representation quality against downstream mismatch.

def mask_span(tokens, span_len, seed=0):
    rng = random.Random(seed)
    start = rng.randrange(max(1, len(tokens) - span_len + 1))
    out = list(tokens)
    for i in range(start, start + span_len):
        out[i] = MASK
    return out, list(range(start, start + span_len))

Span masking hides phrase-level structure. Whole-word masking hides all subword pieces of a word together, which matters for morphologically rich languages. Both are direct fixes to the same bottleneck: the model must guess the unit that carries meaning, not a predictable fragment.

The training loop: masking plus prediction#

Training combines masking, encoding, and classification. Each step picks a batch, masks it, runs the encoder, and computes loss only at masked positions.

def training_step(batch: list[list[str]], vocab_size: int):
    losses = []
    for tokens in batch:
        masked, idx = mask_tokens(tokens)
        # Pseudo-forward pass: hidden states per position,
        # logits over the vocabulary at each masked position.
        hidden = encode(masked)          # list of vectors
        for i in idx:
            logits = classifier(hidden[i])  # vocab_size scores
            probs = softmax(logits)
            losses.append(cross_entropy(probs, vocab_index(tokens, i)))
    return sum(losses) / len(losses) if losses else 0.0

The cross-entropy at one masked position is:

where $y_i$ is the true token and $x_{\setminus i}$ is the sequence with position $i$ hidden. The total loss averages over masked positions and over the batch. No labels beyond the text are required. The labels are the text.

The bottleneck is context width. Bidirectional attention lets every visible position influence every masked position. Quadratic attention cost limits sequence length. Longer documents need truncation, sliding windows, or sparse attention. Each fix trades global dependency for compute.

Why the objective worked#

BERT matched or exceeded prior methods on classification, question answering, and natural language inference with one architecture and fine-tuning. The reason is transfer. Masked prediction forces the hidden state to encode which words fit which contexts. That knowledge is a good basis for downstream tasks that need a sentence representation.

Bidirectionality matters for understanding tasks. An autoregressive model predicts the next token from the left. A masked model predicts from both sides. For tasks that require the whole sentence, the bidirectional signal is more informative. The gain shows up in embeddings: pooled encoder outputs cluster by meaning without any label supervision.

Why generation needed a different objective#

A masked model cannot generate text in the usual sense. It predicts missing tokens rather than producing tokens in sequence. Unmasking one token at a time and conditioning on the rest is possible but slow and awkward. Call-order masking produces incoherent sequences because each step ignores the order the text will be read in.

For open-ended generation, autoregressive models took over, and the field followed. For understanding tasks where a single label or embedding suffices, masked models remain competitive. The choice is task-shaped: generation wants left-to-right prediction, representation wants bidirectional prediction.

Where the objective still appears#

• Encoders for retrieval. Pooled masked-model embeddings support semantic similarity and search.
• Protein and genome models. Masked prediction on biological sequences learns structure without labels.
• Multimodal models. Masked image modeling predicts masked patches and learns visual representations.
• Mixed pre-training. Generative models often include masked prediction as an auxiliary objective because it improves representations.

What to measure#

Masked-model quality needs three separate measurements. Prediction accuracy measures the game itself: fraction of masked tokens ranked first by the model. Representation quality measures transfer: downstream task scores after fine-tuning or with frozen features. Consistency measures stability: whether performance holds across masking rates and span configurations.

def accuracy_at_mask(pairs, predict):
    correct = 0
    for tokens, idx in pairs:
        masked, held_out = mask_tokens(tokens)
        for i, true_tok in zip(idx, held_out):
            if predict(masked, i) == true_tok:
                correct += 1
    return correct / max(sum(len(p) for p in pairs), 1)

The second axis is efficiency: tokens processed per second, loss reduction per compute unit, and accuracy against masking rate. A policy that masks more tokens trains faster per example but risks predicting fragments instead of structure.

A checklist#

• Report masking rate, span policy, and the replacement distribution alongside accuracy.
• Evaluate representation transfer separately from masked-token accuracy.
• Check downstream mismatch: fine-tuning data contains no mask tokens.
• Compare span masking against token masking on morphologically rich languages.
• Budget attention cost against sequence length before extending context.
• Choose the objective by task: bidirectional for understanding, autoregressive for generation.

The unmasking game taught the field that prediction is a powerful training signal. The lesson survived the game. Masked prediction remains the cleanest way to learn representations from unlabeled text when the downstream task is understanding rather than generation.