BERT & Masked Language Modeling

A traditional autoregressive language model factors the probability of a sequence as a product of conditionals, $P(x) = \prod_i P(x_i \mid x_{<i})$, and can only condition each prediction on preceding tokens. This unidirectionality is a training necessity, not a design choice: if you let the model see the whole sequence while training it to predict each token in place, the model could trivially copy the answer through the self-attention mechanism, since every position attends to every other position including itself and the very token it is supposed to predict.

Masked language modeling (MLM) breaks this trivial-solution problem by replacing a subset of input tokens with a special `[MASK]` token (or another substitute) before the model ever sees them, and training the model to predict the original identity of only those masked positions. Because the true token is removed from the input, attending to that position no longer leaks the answer, and the model is forced to infer the missing word from the surrounding bidirectional context. This is sometimes called a cloze task, after the classic reading-comprehension exercise of filling in blanks. The second pretraining objective, next sentence prediction (NSP), trains the model to judge whether one sentence B plausibly follows another sentence A in the original text, versus being a random sentence sampled from elsewhere in the corpus. NSP was intended to teach sentence-pair relationships useful for tasks like question answering and natural language inference, though later research (RoBERTa and successors) showed NSP contributes little and can be dropped or replaced with cheaper objectives, provided training still exposes the model to pairs of contiguous text segments.

Input representation. Before any token reaches the Transformer, BERT builds an input embedding as the sum of three signals:

• Token embeddings — a learned vector per WordPiece subword, from a fixed vocabulary (tying back to Tokenization at Scale).
• Segment embeddings — one of two learned vectors indicating whether a token belongs to sentence A or sentence B.
• Position embeddings — a learned vector per absolute position, since self-attention has no inherent notion of order.

A special `[CLS]` token is prepended to every input; its final hidden state is used as an aggregate sequence representation for classification tasks. A `[SEP]` token separates sentence A from sentence B (and terminates a single-sentence input).

Masking procedure. During pretraining, 15% of WordPiece tokens are selected uniformly at random (excluding `[CLS]`/`[SEP]`) to participate in the masked-prediction loss. Of those selected positions:

• 80% are replaced with the literal `[MASK]` token,
• 10% are replaced with a random token from the vocabulary,
• 10% are left unchanged.

This 80/10/10 split exists to close a train/fine-tune mismatch: `[MASK]` never appears during fine-tuning or inference, so if the model only ever saw `[MASK]` during pretraining it would learn representations conditioned on an artificial token that never recurs later. Occasionally leaving the original token, or substituting a random wrong one, forces the encoder to maintain a good contextual representation for every input token (not just masked ones), because it can never be sure whether a given token is "trustworthy" or secretly a corrupted target it must correct.

The Transformer encoder. The embedded, positionally-augmented sequence is passed through $L$ stacked Transformer encoder layers, each composed of multi-head self-attention followed by a position-wise feed-forward network, with residual connections and layer normalization around each sublayer (assume this stack as given from Transformer Encoder Stack). Because encoder self-attention is fully bidirectional — every position attends to every other position in the sequence with no causal mask — each output vector is a function of the entire input, subject to the masking-out described above.

Output heads. At the masked positions, the final hidden vectors are fed through a linear layer plus softmax over the vocabulary to predict the original token, and the cross-entropy loss is computed only over the masked positions. Separately, the final hidden vector of `[CLS]` is fed through a small feed-forward classifier and softmax to predict `IsNext` versus `NotNext` for the NSP loss. The two losses are summed and backpropagated jointly. ``` function BUILD_PRETRAINING_EXAMPLE(sentence_A, sentence_B, is_actual_next): tokens = ["[CLS]"] + tokenize(sentence_A) + ["[SEP]"] + tokenize(sentence_B) + ["[SEP]"] segment_ids = [0](len(sentence_A_tokens)+2) + [1](len(sentence_B_tokens)+1) candidate_positions = all indices i where tokens[i] not in {"[CLS]", "[SEP]"} shuffle(candidate_positions) num_to_mask = round(0.15 * len(candidate_positions)) mask_positions = candidate_positions[0 : num_to_mask] labels = array of size len(tokens), default = IGNORE for i in mask_positions: labels[i] = tokens[i] # ground truth to predict r = random_uniform(0, 1) if r < 0.8: tokens[i] = "[MASK]" else if r < 0.9: tokens[i] = random_vocab_token()

# else: leave tokens[i] unchanged (10% of the time) return tokens, segment_ids, labels, is_actual_next

function PRETRAIN_STEP(model, batch): total_loss = 0 for (tokens, segment_ids, mlm_labels, nsp_label) in batch: input_embeds = TokenEmbed(tokens) + SegmentEmbed(segment_ids) + PositionEmbed(0..len(tokens)-1) hidden = TransformerEncoderStack(input_embeds) # bidirectional, L layers

# --- Masked LM loss (only over masked positions) --- mlm_loss = 0 for i where mlm_labels[i] != IGNORE: logits_i = OutputProjection(hidden[i]) # vocab-sized logits mlm_loss += CrossEntropy(softmax(logits_i), mlm_labels[i])

# --- Next sentence prediction loss --- cls_vector = hidden[0] nsp_logits = NSPHead(cls_vector) # 2-way logits nsp_loss = CrossEntropy(softmax(nsp_logits), nsp_label) total_loss += mlm_loss + nsp_loss gradients = backprop(total_loss) update_parameters(model, gradients) return total_loss / len(batch) ```

Fine-tuning. After pretraining, the same encoder is reused for a downstream task by attaching a small task-specific head — a linear classifier on `[CLS]` for sentence classification, two linear layers predicting span start/end logits for extractive question answering, a token-wise classifier for sequence labeling — and all parameters (encoder plus head) are updated together for a small number of epochs on the labeled task data. Because the encoder already encodes rich contextual and syntactic structure, this stage typically needs orders of magnitude less labeled data than training from scratch.

Time complexity. Each Transformer encoder layer computes self-attention over a sequence of length $n$ with hidden dimension $d$. Computing the pairwise attention scores costs $O(n^2 d)$ (the $n \times n$ attention matrix, each entry requiring an $O(d)$ dot product), and the position-wise feed-forward sublayer costs $O(n d^2)$. Stacking $L$ layers gives a total forward-pass cost of

per sequence. This is identical to the cost of a plain Transformer encoder — MLM changes what is predicted, not the architecture's asymptotic cost. The MLM output projection adds $O(m \cdot d \cdot |V|)$ for $m$ masked positions and vocabulary size $|V|$, which is typically small relative to the encoder cost since $m \approx 0.15n \ll n$ in absolute count but the projection is applied only at those $m$ positions rather than all $n$ (a deliberate efficiency choice: computing softmax over $|V|$ at every position would cost $O(n|V|d)$ instead).

Space complexity. Storing the attention matrices for backpropagation costs $O(L n^2)$, and activations across layers cost $O(L n d)$; parameters cost $O(L d^2)$ independent of sequence length. For BERT-base ($L=12$, $d=768$) and BERT-large ($L=24$, $d=1024$), parameter counts are roughly 110M and 340M respectively, dominated by the $O(Ld^2)$ feed-forward and attention projection weights.

Why masking (rather than naive bidirectionality) is necessary — a correctness argument. Consider training a bidirectional encoder to predict every token $x_i$ directly from a hidden state $h_i$ that was computed by letting position $i$ attend to all positions including itself. Since self-attention, absent any masking, includes the identity mapping among its accessible information (the query at position $i$ can trivially attend most strongly to the key at position $i$, and the residual connection passes the input embedding of $x_i$ straight through), the model can achieve zero loss by learning the identity function, never learning any genuine contextual dependency. This degenerate optimum is a global minimizer of the naive objective, so gradient descent has no pressure to avoid it. MLM breaks this degenerate optimum by removing the target token from the input entirely at masked positions (replacing it with `[MASK]`, a random token, or leaving it — but computing loss only at the subset that was selected for masking). The identity shortcut is unavailable at a masked position because the true $x_i$ is not present anywhere in the input for the model to copy; the residual stream at that position carries information derived only from the corrupted/absent token plus attention to other positions. The only way to reduce loss at masked positions is to build a representation that infers the missing token from bidirectional context — which is exactly the representation-learning signal MLM is designed to elicit. The 10%-unchanged rule reintroduces a small amount of genuine self-information at a controlled rate specifically to prevent the encoder from underusing the current token's own identity at non-masked positions during fine-tuning, where no masking occurs at all.

Tradeoffs. MLM's central cost is statistical efficiency at training time: because loss is computed at only ~15% of positions per example, each training example yields less gradient signal per token than an autoregressive model's per-token loss over the entire sequence, requiring more total pretraining steps or data to match. BERT also cannot generate text autoregressively — there is no well-defined way to sample a full sequence left-to-right from a model trained only to fill in blanks — which is why the encoder-only, MLM-pretrained family (BERT, RoBERTa, ALBERT) is used for representation learning and classification/extraction tasks, while decoder-only, next-token-pretrained models (GPT family) dominate open-ended generation.

Consider the sentence pair used for pretraining: > Sentence A: "the cat sat on the mat" > Sentence B: "it purred softly" (the true next sentence, so `IsNext`) After WordPiece tokenization and adding special tokens: ``` [CLS] the cat sat on the mat [SEP] it purred softly [SEP] 0 1 2 3 4 5 6 7 8 9 10 11 ``` Suppose the random 15% selection picks positions 2 (`cat`) and 9 (`purred`) to mask. Applying the 80/10/10 rule with independent draws:

The final model input becomes: ``` [CLS] the [MASK] sat on the mat [SEP] it purred softly [SEP] ``` with `mlm_labels[2] = "cat"` and `mlm_labels[9] = "purred"` (the label at position 9 is still recorded and scored even though the visible token was left unchanged — the model must learn that even an apparently-normal token might be a scored target). Passing this through the encoder, the hidden vector at position 2 is computed with full bidirectional self-attention: it can attend to `the` (position 1), `sat`, `on`, `the`, `mat` (positions 3–6), and even the second sentence, `it purred softly` (positions 8–10) through the shared attention mechanism across the `[SEP]` boundary. Given "the ___ sat on the mat" plus "it purred softly," the model should assign high probability to `cat` at the MLM head applied to hidden[2]. Simultaneously, the hidden vector at position 9 is scored against the label `purred`, forcing the model to maintain an accurate representation of tokens even when they were not corrupted. Separately, the `[CLS]` hidden state feeds the NSP head. Because sentence B genuinely follows sentence A in the source text, the target label is `IsNext`; had sentence B been swapped for an unrelated sentence sampled elsewhere in the corpus (e.g., "stock prices fell sharply today"), the target would be `NotNext`, and the model would need to use the `[CLS]` representation to detect the topical and discourse discontinuity. Total loss for this example is `CrossEntropy(predicted[2], "cat") + CrossEntropy(predicted[9], "purred") + CrossEntropy(NSP_prediction, IsNext)`, backpropagated through the shared encoder so that every parameter is updated to improve both objectives jointly.

BERT-style pretraining reshaped how production NLP systems are built. Search engines use bidirectional encoder representations to better match query intent to document content, since understanding a query word often requires seeing words on both sides of it. Question-answering systems fine-tune a pretrained encoder with a span-extraction head to locate answers within a passage, reusing the same encoder weights across many domains. Content moderation, sentiment analysis, and named-entity recognition pipelines commonly fine-tune a compact BERT variant (DistilBERT, ALBERT, or a domain-specific checkpoint like BioBERT or SciBERT) rather than training from scratch, since the pretrained encoder already encodes syntax and much of the world knowledge needed for the downstream task. The masked-token idea also generalizes beyond text: masked modeling objectives inspired vision transformers trained with masked patch prediction and masked audio modeling for speech representation learning, showing that the cloze-style self-supervision principle is architecture- and modality-agnostic. Understanding BERT is a natural bridge to the broader Pretraining-Finetuning Paradigm and to Transfer Learning in NLP, both of which generalize the idea that a single large model pretrained on unlabeled data can be efficiently specialized for many downstream tasks with modest labeled data and compute.

Ask the tutor: How would you modify the masking procedure or objective to make pretraining more sample-efficient, the way ELECTRA does?