4.2.7Tokenization & Language Modeling

Masked language modeling (BERT)

2,195 words10 min readdifficulty · medium4 backlinks

Overview

Masked Language Modeling (MLM) is a pre-training objective where we randomly mask some tokens in the input sequence and train the model to predict the original masked tokens based on their bidirectional context. This is the core innovation behind BERT (Bidirectional Encoder Representations from Transformers).

WHY does this matter? Traditional language models (like GPT) are unidirectional — they only see words to the left. But understanding language requires seeing BOTH directions. "The bank was steep" vs "The bank was closed" — you need words on BOTH sides to know if we mean riverbank or financial institution.

Figure — Masked language modeling (BERT)

Unlike next-word prediction (which is causal/autoregressive), MLM is a denoising autoencoder approach — we corrupt the input and learn to reconstruct it.


The Masking Strategy

WHY 15%? Balance between:

  • Too low → not enough training signal
  • Too high → too much context removed, harder to predict

WHY not 100% [MASK]? The [MASK] token never appears during fine-tuning (downstream tasks), so we'd have a train-test mismatch. By using random tokens 10% and unchanged 10%, the model learns to be robust.

The MLM loss is: LMLM=iMlogP(xix\M)\mathcal{L}_{\text{MLM}} = -\sum_{i \in M} \log P(x_i \mid \mathbf{x}_{\backslash M})

where x\M\mathbf{x}_{\backslash M} represents the corrupted sequence (with masks).

DERIVATION from first principles:

  1. Start with likelihood: We want to maximize the probability of the original tokens given the corrupted input. P(xMx\M)=iMP(xix\M)P(\mathbf{x}_M \mid \mathbf{x}_{\backslash M}) = \prod_{i \in M} P(x_i \mid \mathbf{x}_{\backslash M})

    WHY product? Assuming conditional independence of predictions given the context (each masked token predicted independently).

  2. Take log for numerical stability: logP(xMx\M)=iMlogP(xix\M)\log P(\mathbf{x}_M \mid \mathbf{x}_{\backslash M}) = \sum_{i \in M} \log P(x_i \mid \mathbf{x}_{\backslash M})

    WHY log? Products of probabilities → underflow. Sums are stable.

  3. Negate for minimization: LMLM=iMlogP(xix\M)\mathcal{L}_{\text{MLM}} = -\sum_{i \in M} \log P(x_i \mid \mathbf{x}_{\backslash M}) WHY negative? Gradient descent minimizes loss. Maximizing log-likelihood = minimizing negative log-likelihood.

  4. Implementation: The probability P(xix\M)P(x_i \mid \mathbf{x}_{\backslash M}) comes from: P(xix\M)=exp(wxihi)vVexp(wvhi)P(x_i \mid \mathbf{x}_{\backslash M}) = \frac{\exp(\mathbf{w}_{x_i}^\top \mathbf{h}_i)}{\sum_{v \in \mathcal{V}} \exp(\mathbf{w}_v^\top \mathbf{h}_i)}

    where:

    • hi\mathbf{h}_i = contextualized embedding from Transformer at position ii
    • wv\mathbf{w}_v = output embedding for vocabulary token vv
    • This is a softmax over the entire vocabulary

Architecture Details

WHY encoder-only? We're not generating sequences autoregressively. We need bidirectional context for understanding, not left-to-right generation.

Input Representation

The input to BERT combines three embedings (added element-wise):

Einput(i)=Etoken(xi)+Eposition(i)+Esegment(si)\mathbf{E}_{\text{input}}(i) = \mathbf{E}_{\text{token}}(x_i) + \mathbf{E}_{\text{position}}(i) + \mathbf{E}_{\text{segment}}(s_i)

WHY three types?

  • Token embedding: The word/subword itself
  • Position embedding: Which position in sequence (learned, not sinusoidal like original Transformer)
  • Segment embedding: Which sentence (A or B) — needed for Next Sentence Prediction (NSP), the second pre-training task

After WordPiece tokenization: [CLS] the cat sat on the mat [SEP]

Masking (15% → here we mask "cat"):

[CLS] the [MASK] sat on the mat [SEP]

Token embeddings: Look up in vocab Position embeddings: [0, 1, 2, 3, 4, 5, 6, 7] Segment embeddings: [A, A, A] (all sentence A)

The model outputs hidden states H=(h0,h1,,h7)\mathbf{H} = (\mathbf{h}_0, \mathbf{h}_1, \ldots, \mathbf{h}_7)

At position2 (where [MASK] is), we compute: P(x2="cat"context)=softmax(Wh2+b)[cat_id]P(x_2 = \text{"cat"} \mid \text{context}) = \text{softmax}(\mathbf{W} \mathbf{h}_2 + \mathbf{b})[\text{cat\_id}]

where W\mathbf{W} projects back to vocabulary size.


Training Process

Input sentence: "The quick brown fox jumps"

Step 1 - Tokenize:

[CLS] the quick brown fox jumps [SEP]

Step 2 - Random selection (15%): Suppose we select tokens at positions {3, 5} (brown, jumps)

Step 3 - Apply masking protocol:

  • Position 3: 80% chance → [MASK]
  • Position 5: 10% chance → random token, say "happy"

Result:

[CLS] the quick [MASK] fox happy [SEP]

Step 4 - Forward pass: Transformer encoder produces h3\mathbf{h}_3 and h5\mathbf{h}_5

Step 5 - Predict originals:

  • At position 3: Compute P(h3)P(\cdot \mid \mathbf{h}_3) over vocab. True label: "brown"
  • At position 5: Compute P(h5)P(\cdot \mid \mathbf{h}_5) over vocab. True label: "jumps"

Step 6 - Compute loss: L=logP("brown"h3)logP("jumps"h5)\mathcal{L} = -\log P(\text{"brown"} \mid \mathbf{h}_3) - \log P(\text{"jumps"} \mid \mathbf{h}_5)

Step 7 - Backprop and update weights

WHY this works: The model CANNOT cheat by memorizing position-token mappings because:

  1. Only 15% masked (not every position)
  2. Random corruption (10% random, 10% unchanged)
  3. Bidirectional attention forces using ALL context

Differences from Other Language Models

Aspect GPT (Causal LM) BERT (Masked LM) T5 (Seq2Seq)
Architecture Decoder-only Encoder-only Encoder-Decoder
Attention Causal (left-to-right) Bidirectional Enc: bidirectional, Dec: causal
Objective Predict next token Predict masked tokens Span corruption
Use case Generation Understanding/Classification Both
Context Unidirectional Full bidirectional Bidirectional encoding

WHY can't GPT use bidirectional attention? Because it's designed for autoregressive generation. If it could "see the future", it would cheat during generation.

WHY can BERT use bidirectional? Because we're not generating — we're encoding for understanding. During fine-tuning (classification, QA), we don't generate, we just need rich representations.


Next Sentence Prediction (NSP)

BERT uses a second pre-training objective alongside MLM:

The loss is: LNSP=logP(label[CLS])\mathcal{L}_{\text{NSP}} = -\log P(\text{label} \mid [\text{CLS}])

where the [CLS] token's final hidden state h[CLS]\mathbf{h}_{[\text{CLS}]} is passed through a binary classifier.

WHY NSP? To help BERT learn sentence-level relationships for tasks like Question Answering and Natural Language Inference.

TOTAL BERT LOSS: LBERT=LMLM+LNSP\mathcal{L}_{\text{BERT}} = \mathcal{L}_{\text{MLM}} + \mathcal{L}_{\text{NSP}}

Input to BERT:

[CLS] I went to the store [SEP] I bought some milk [SEP]

NotNext pair:

  • Sentence A: "I went to the store."
  • Sentence B: "The moon orbits Earth."
  • Label: 0

Input:

[CLS] I went to the store [SEP] The moon orbits Earth [SEP]

The model learns to use h[CLS]\mathbf{h}_{[\text{CLS}]} as a sentence-pair representation.


Fine-tuning BERT

After pre-training with MLM+NSP on massive text corpora, BERT is fine-tuned on downstream tasks:

For classification tasks: y=softmax(Wh[CLS]+b)y = \text{softmax}(\mathbf{W} \mathbf{h}_{[\text{CLS}]} + \mathbf{b})

For token-level tasks (NER, POS tagging): yi=softmax(Whi+b)y_i = \text{softmax}(\mathbf{W} \mathbf{h}_i + \mathbf{b})

For Question Answering (span extraction): Learn start/end pointers: P(start=i)=softmax(Shi)P(\text{start}=i) = \text{softmax}(\mathbf{S}^\top \mathbf{h}_i) P(end=j)=softmax(Ehj)P(\text{end}=j) = \text{softmax}(\mathbf{E}^\top \mathbf{h}_j)

WHY this works: The pre-trained bidirectional representations already encode rich contextual information. We just add a thin task-specific layer on top.

Key insight: Same pre-trained model → many different tasks. The "knowledge" is in the encoder weights.


Common Mistakes

Why it feels right: We trained with [MASK], so we should test with [MASK].

The fix: Never use [MASK] during fine-tuning. Downstream tasks have real text, no masks. This is why BERT masks only 80% with [MASK] and uses random/unchanged for the other 20% — to reduce train-test mismatch.

Steel-man: The concern about distribution shift is valid! But the10% random + 10% unchanged in pre-training creates robustness. The model learns: "sometimes there's a [MASK], sometimes not — I need to use context regardless."

Why it feels right: BERT predicts tokens, so it should generate text.

The fix: BERT is an encoder, not a generator. It predicts tokens in parallel given bidirectional context. It has no notion of left-to-right autoregressive generation. For generation, use GPT-style models.

Steel-man: You could technically generate by iteratively masking and predicting, but it's inefficient and incoherent because BERT wasn't trained for sequential generation. Use the right tool for the job.

Why it feels right: If we mask everything, we train on everything.

The fix: The 10% random + 10% unchanged is crucial. If we only saw [MASK], the model would learn a shortcut: "if I see [MASK], predict based on context; if I see real token, ignore it." During fine-tuning (no [MASK]), this fails catastrophically.


Mathematical Properties

## 🖼️ Concept Map ```mermaid flowchart TD MLM[Masked Language Modeling] BERT[BERT Bidirectional Encoder] Bidir[Bidirectional Context] Uni[Unidirectional GPT] Denoise[Denoising Autoencoder] Mask15[Mask 15% of Tokens] Split[80% MASK / 10% Random / 10% Keep] Predict[Predict Original Tokens] Loss[MLM Loss] Mismatch[Train-Test Mismatch] MLM -->|core of| BERT BERT -->|uses| Bidir Bidir -->|contrasts with| Uni MLM -->|is a| Denoise Denoise -->|corrupts via| Mask15 Mask15 -->|split into| Split Split -->|avoids| Mismatch Mask15 -->|then| Predict Predict -->|optimized by| Loss Loss -->|derived from| Predict ``` ## 🔊 Hinglish (regional understanding) > [!intuition]- Hinglish mein samjho > Chalo is concept ko simple tareeke se samajhte hain. Masked Language Modeling ka core idea ekdum fill-in-the-blank test jaisa hai. Jaise school mein aapko sentence milta tha jisme kuch words missing hote the, aur aapko blanks bharne hote the — bilkul waisa hi BERT karta hai. Hum input sentence ke kuch tokens (around 15%) ko chupa dete hain (mask kar dete hain), aur model ko train karte hain ki wo original word predict kare. Sabse important baat yeh hai ki BERT dono directions dekh sakta hai — left aur right dono side ke words ka context use karta hai. Isliye ise "bidirectional" kehte hain, jabki GPT jaise models sirf left side dekhte hain. > > Ab yeh matter kyun karta hai? Language samajhne ke liye dono directions ka context zaroori hai. Jaise "The bank was steep" aur "The bank was closed" — dono mein "bank" word hai, lekin meaning alag hai (riverbank vs paisa wala bank). Sirf ek side ke words se aap decide nahi kar sakte. BERT ko yeh puzzle solve karne ke liye majboor kiya jaata hai dono sides pe attention dene ke liye, isse wo deep aur rich representations seekhta hai. Ek chhoti si clever trick bhi hai: masked tokens mein se 80% ko [MASK] se replace karte hain, 10% random word se, aur 10% wahi rehne dete hain. Yeh isliye kyunki fine-tuning ke time [MASK] token appear nahi karta, toh train-test mismatch se bachne ke liye model ko robust banate hain. > > Maths side pe, loss function basically negative log-likelihood hai — model jo probability deta hai correct original word ke liye, usko maximize karna hai. Log isliye lete hain kyunki chhote probabilities ka product underflow kar deta hai (numbers itne chhote ho jaate hain ki computer handle nahi kar paata), aur log lene se yeh sum ban jaata hai jo stable hota hai. Negative isliye kyunki gradient descent minimize karta hai, toh maximize karne ke bajaye hum negative bana ke minimize karte hain. Final prediction ek softmax se aati hai jo poore vocabulary pe probability distribution deta hai. Bas yaad rakho — BERT encoder-only hai kyunki ise sequence generate nahi karni, sirf samajhni hai, aur yahi uska sabse bada strength hai. ![[audio/4.2.07-Masked-language-modeling-(BERT).mp3]]

Go deeper — visual, from zero

Test yourself — Tokenization & Language Modeling

Connections