How a model that reads in both directions becomes a search engine, and how search engines become the memory of language models. BERT, contrastive learning, vector indexes, BM25, and RAG, built from first principles.
Chapter 7 · Section 7.1 · 50 min

Encoders and BERT

Fill in the blank. "The ___ barked all night, so the neighbors complained." You knew the answer the moment you read "barked", which comes after the blank. A GPT-style model reading left to right gets to the blank having seen only "The", so the best it can do is guess. For generating text that's the right design, because the future doesn't exist yet. But a huge share of what we ask models to do is understanding text that already exists in full (classify this review, find the answer span in this paragraph, embed this document for search). Throwing away the right half of the sentence there is a strange handicap.

This section is about the models that read in both directions, called encoders, and about the training trick that made them work, which is BERT (Devlin et al., 2019). BERT and its descendants are not a historical curiosity. They are the backbone of most production search, classification, and retrieval systems today, and every embedding model in the rest of this chapter starts life as one of them.

Why a language model can't just look both ways

Recall from §4.4 how a decoder learns. Every position predicts the next token, and a causal mask stops position tt from attending to anything after it. Remove the mask and the task becomes trivial. In the very first layer, position tt can attend straight to token t+1t+1 and copy it. The model would reach zero loss and learn nothing.

So bidirectional attention needs a different question, one whose answer is not sitting in the input. The old cloze test from reading comprehension (Taylor, 1953) gives it. Hide some words and ask the model to fill them in from everything else. The hidden word can't leak, because it isn't there.

Before we build it, let's check that the right-hand context really carries that much information. Below we generate a tiny corpus of sentences like "the dog barked all night", then predict each middle word from counts, once from the word on its left only, and once from the words on both sides.

Python · cell 1

The left-only model can't name the animal at all. After "the", its best guess is whatever word most often follows "the" in the corpus, which here is a place noun. With one word of right context, it's right more than 90 percent of the time, and the misses come from the 10 percent of sentences we made noisy. Real text is less extreme, but the lesson holds. For understanding a sentence that's already written, both directions matter.

Masked language modeling

Here is the objective, written for one sequence of tokens x=(x1,…,xT)\mathbf{x} = (x_1, \dots, x_T) drawn from a vocabulary V\mathcal{V} of size VV. We pick a random set of positions M⊂[T]\mathcal{M} \subset [T], corrupt the input at those positions to get x~\tilde{\mathbf{x}}, run a bidirectional Transformer on x~\tilde{\mathbf{x}}, and ask it to recover the original tokens,

LMLM(θ)=−1∣M∣∑t∈Mlog⁡pθ(xt∣x~),pθ(⋅∣x~)=softmax⁡(E ht+b)∈RV.\mathcal{L}_{\text{MLM}}(\theta) = -\frac{1}{|\mathcal{M}|}\sum_{t \in \mathcal{M}} \log p_\theta\big(x_t \mid \tilde{\mathbf{x}}\big), \qquad p_\theta(\cdot \mid \tilde{\mathbf{x}}) = \operatorname{softmax}\big(\mathbf{E}\,\mathbf{h}_t + \mathbf{b}\big) \in \mathbb{R}^{V}.

Here ht∈Rd\mathbf{h}_t \in \mathbb{R}^{d} is the final hidden state at position tt (after a small dense layer with GELU and LayerNorm), E∈RV×d\mathbf{E} \in \mathbb{R}^{V \times d} is the input embedding table reused as the output layer (weight tying, as in §1.5 and §6.2), and b∈RV\mathbf{b} \in \mathbb{R}^V is a bias.

In plain words, hide some tokens, encode the damaged sentence with full attention in both directions, and score the model only on how well it reconstructs the hidden tokens.

Three details in that formula matter, and interviewers like all three.

  1. The loss covers only M\mathcal{M}. Positions that weren't selected contribute nothing. In BERT, ∣M∣|\mathcal{M}| is about 15 percent of the tokens, so each sequence gives about 6.7 times fewer prediction targets than a GPT gets from the same text. BERT's authors measured that MLM does converge slightly more slowly than left-to-right training, but that its downstream accuracy pulls ahead almost immediately and stays ahead (Devlin et al., 2019, Appendix C.1).
  2. Each masked token is predicted independently given x~\tilde{\mathbf{x}}. If "New" and "York" are both masked, the model predicts each without knowing what it predicted for the other. That's one reason MLM is not a proper generative model of text. You can't sample a paragraph from it the way you sample from a GPT.
  3. The prediction head is cheap. The heavy lifting is in the encoder. After pretraining, the head is thrown away, and ht\mathbf{h}_t is the contextual embedding you keep (exactly the "contextual token embeddings" of §1.8 and §3.4).
Core

The MLM loss and its gradient

Implement mlm_loss(logits, labels) for logits of shape (B, T, V) and labels of shape (B, T), where labels[b, t] is the original token id at selected positions and -100 everywhere else (the convention PyTorch's cross_entropy(ignore_index=-100) uses). Return (loss, dlogits), where loss is the mean negative log-likelihood over selected positions only, and dlogits has the shape of logits. Use a stable log-softmax. If no position is selected, return a loss of 0 and a zero gradient, not NaN.

Hidden tests, hints, and AI help, from $1.99 a month.

The 80-10-10 rule

Now the detail that separates people who have read the BERT paper from people who have heard of it. The simplest corruption would replace every selected token with a special [MASK] token. BERT doesn't do that. It selects 15 percent of the token positions, and then, for each selected position,

  • 80 percent of the time it replaces the token with [MASK],
  • 10 percent of the time it replaces it with a random token from the vocabulary,
  • 10 percent of the time it leaves the token unchanged.

The loss is computed at all selected positions, whatever happened to them. Why the complication? Because [MASK] exists only during pretraining. When you fine-tune BERT to classify reviews, or run it to embed documents, the input never contains a [MASK]. That's a train and test mismatch, and each of the two extra cases fixes a specific part of it.

Keeping the token unchanged. If every scored position were a [MASK], the model would never be asked to say anything useful at a position where it can see a real word. The final hidden state of a visible token could become a lazy summary, good enough to help its neighbors but not a good representation of that token in context. And visible tokens are all there is at fine-tuning time. Scoring some unchanged tokens forces the model to keep a representation at every visible position that predicts the token itself, which BERT's authors describe as biasing the representation toward the observed word.

Replacing with a random token. If the unchanged case were the only addition, the model could learn a shortcut. "If I can see a real word, copy it." The 10 percent of random replacements break the shortcut. Now a visible word might be wrong, so the model can't trust it blindly and has to check it against the context.

Put together, the model never knows which visible tokens will be scored, or whether they have been tampered with. The paper's phrase is that it is "forced to keep a distributional contextual representation of every input token". The cost is small. Random replacement touches only 0.10×0.15=1.50.10 \times 0.15 = 1.5 percent of all tokens, which BERT's authors found didn't hurt.

What does BERT actually see, and what is it graded on?
Each non-special token is selected for grading with the masking rate. Selected tokens are then replaced by [MASK], replaced by a random word, or left alone. The model is graded at every selected position, whatever happened to it.
Loading interactive figure…

How much does the rule actually matter? Less than its fame suggests, and that's worth knowing. Here are BERT's own ablations (Devlin et al., 2019, Table 8), dev-set accuracy on MNLI and F1 on named entity recognition (NER). For NER they also tried a feature-based setup where BERT's weights are frozen, which is where a mismatch should hurt most because the model can't adapt.

[MASK]samerandomMNLI (fine-tune)NER (fine-tune)NER (frozen features)
80%10%10%84.295.494.9
100%0%0%84.394.994.0
80%0%20%84.195.294.6
80%20%0%84.495.294.7
0%20%80%83.794.894.6
0%0%100%83.694.994.6

Fine-tuning is "surprisingly robust" to the choice, in the authors' words. The mismatch shows up where predicted, in frozen features, where pure [MASK] loses 0.9 F1. And using only random replacement is clearly worse. So the honest summary for an interview is that 80-10-10 is a sensible, cheap hedge against the pretrain and fine-tune mismatch, it mostly matters when the encoder is used frozen as a feature extractor, and it is not where BERT's big gains came from.

The model scores 0.15×512=76.80.15 \times 512 = 76.8 tokens on average. Of those, about 61.461.4 become [MASK], 7.77.7 become random tokens, and 7.77.7 stay as they were. So about 6969 of the 512 input tokens, 13.5 percent, are actually corrupted, and the remaining 7.7 scored tokens look exactly like the 435 unscored ones.

Interview

MLM masking with the 80-10-10 rule

Implement mlm_mask(ids, special, vocab_size, mask_id, first_regular, rng, p=0.15).

  • ids is an int array of shape (B, T) and special a boolean array of the same shape marking [CLS], [SEP], and [PAD] positions, which must never be selected.
  • Select each non-special position independently with probability p.
  • At selected positions, use [MASK] (mask_id) with probability 0.8, a random token drawn uniformly from range(first_regular, vocab_size) with probability 0.1, and the original token with probability 0.1.
  • Return (inputs, labels). inputs is the corrupted copy (do not modify ids), and labels holds the original id at selected positions and -100 elsewhere.

No Python loops over tokens. This is a real interview question, and the tests check the statistics on a large batch.

Hidden tests, hints, and AI help, from $1.99 a month.

Next sentence prediction, and why it was dropped

BERT had a second objective. Many tasks, such as question answering and natural language inference, are about the relationship between two pieces of text. To teach that, each pretraining example packs two segments, A and B, into one input,

[CLS]  the man went to the store  [SEP]  he bought a gallon of milk  [SEP]

and 50 percent of the time B really follows A in the corpus, while the other 50 percent it's a random segment from elsewhere. A binary classifier on the final hidden state of [CLS] predicts which, and that loss is added to the MLM loss. Each token's input embedding is the sum of three learned vectors, its token embedding, a segment embedding (A or B), and a learned absolute position embedding (up to 512 positions, as in §5.4).

BERT's ablation reported that removing NSP hurt QNLI, MNLI, and SQuAD, and the final model got 97 to 98 percent accuracy on the task. That number is the first clue that something was off. A task you solve at 98 percent can't be teaching much.

Three later papers took NSP apart.

  • RoBERTa (Liu et al., 2019) compared four input formats. Segment pairs with NSP (the original), single-sentence pairs with NSP, and contiguous full sentences without NSP, either crossing document boundaries or not. Removing NSP and packing full-length contiguous text "matches or slightly improves downstream task performance". Single-sentence pairs hurt, which the authors attribute to the model never seeing long-range context.
  • ALBERT (Lan et al., 2020) explained why the task is easy. The negative segment comes from a different document, so it is usually about a different topic. The model can tell the pair apart by topic alone, which MLM already teaches, without learning anything about coherence. ALBERT replaced it with sentence order prediction, where the negative is the same two consecutive segments in swapped order, so topic gives nothing away.
  • SpanBERT (Joshi et al., 2020) also found single long segments without NSP better than two half-length segments with it.

So the modern consensus is no NSP, and full-length contiguous inputs. What survived is the input format. [CLS] A [SEP] B [SEP] is still how you feed a pair of texts to an encoder, for example a query and a passage to a reranker (§7.5).

Fine-tuning, CLS pooling, and spans

BERT's selling point was that one pretrained network fine-tunes to almost any task with a tiny head. Let H∈RT×d\mathbf{H} \in \mathbb{R}^{T \times d} be the final hidden states, with row tt written ht\mathbf{h}_t, and let hCLS\mathbf{h}_{\text{CLS}} be the row for the [CLS] token at position 1.

Sequence classification. With KK classes, add W∈RK×d\mathbf{W} \in \mathbb{R}^{K \times d} and predict softmax⁡(WhCLS)\operatorname{softmax}(\mathbf{W}\mathbf{h}_{\text{CLS}}). Train everything end to end with cross-entropy. Because attention is bidirectional, [CLS] can gather information from every token, and fine-tuning teaches it what to gather. (BERT's code also passes hCLS\mathbf{h}_{\text{CLS}} through a tanh "pooler" layer trained on NSP.)

Token classification. For named entities or part-of-speech tags, put a classifier on every ht\mathbf{h}_t. When a word is split into several WordPieces, the usual convention is to predict at the first piece and ignore the others in the loss.

Span extraction. For extractive question answering (SQuAD), pack the question as segment A and the passage as segment B. Learn a start vector ws∈Rd\mathbf{w}_s \in \mathbb{R}^d and an end vector we∈Rd\mathbf{w}_e \in \mathbb{R}^d. The probability that the answer starts at token ii is a softmax over passage positions,

Pstart(i)=exp⁡(ws⊤hi)∑jexp⁡(ws⊤hj),Pend(j)=exp⁡(we⊤hj)∑j′exp⁡(we⊤hj′),P_{\text{start}}(i) = \frac{\exp(\mathbf{w}_s^\top \mathbf{h}_i)}{\sum_{j}\exp(\mathbf{w}_s^\top \mathbf{h}_j)}, \qquad P_{\text{end}}(j) = \frac{\exp(\mathbf{w}_e^\top \mathbf{h}_j)}{\sum_{j'}\exp(\mathbf{w}_e^\top \mathbf{h}_{j'})},

and the training loss is the sum of the cross-entropies of the true start and the true end. At inference you pick the span (i,j)(i, j) that maximizes ws⊤hi+we⊤hj\mathbf{w}_s^\top\mathbf{h}_i + \mathbf{w}_e^\top\mathbf{h}_j subject to j≥ij \ge i and a maximum answer length. In plain words, each token votes for being the start and for being the end, and the best span is the best valid pair of votes. For SQuAD 2.0, where some questions have no answer, BERT scores "no answer" as the span that starts and ends at [CLS] and predicts it when its score beats the best real span by a threshold tuned on dev data.

The constraint j≥ij \ge i is not decoration. Picking the start and end by separate argmaxes is a classic bug.

Python · cell 2

Check it by hand. The best start alone is token 5 (3.1) and the best end alone is token 1 (2.6), which is before it. Among valid pairs, (5,6)(5, 6) scores 3.1+0.5=3.63.1 + 0.5 = 3.6, while (2,4)(2, 4) scores 2.0+1.9=3.92.0 + 1.9 = 3.9, so the answer is tokens 2 to 4.

Interview

Decode the best answer span

Implement best_span(start, end, max_len, allowed=None) returning the pair (i, j) that maximizes start[i] + end[j] subject to j >= i, j - i + 1 <= max_len, and, if allowed is given, allowed[i] and allowed[j] both true (so the answer stays inside the passage and never lands on question tokens). Use NumPy, with no Python loops over positions. Return plain Python ints.

Hidden tests, hints, and AI help, from $1.99 a month.

The paper's recommended fine-tuning grid is small and still a sensible default for BERT-sized encoders. Batch size 16 or 32, learning rate 5×10−55\times10^{-5}, 3×10−53\times10^{-5}, or 2×10−52\times10^{-5} with Adam, and 2 to 4 epochs. Small datasets are much more sensitive to these choices than large ones, so try the whole grid.

Counting BERT

BERT-base has L=12L = 12 layers, width d=768d = 768, 12 heads, and a feed-forward width of 3072=4d3072 = 4d. BERT-large has L=24L = 24, d=1024d = 1024, 16 heads, and a feed-forward width of 4,096. The famous totals are 110M and 340M parameters. Let's reproduce them, because "count the parameters of BERT-base" is a standard warm-up and a quick check that you understand the architecture (it's the encoder-only Transformer of §5.6, with biases and post-LayerNorm).

Python · cell 3

The count lands on 109.5M for BERT-base, which rounds to the famous 110M, and on 335M for BERT-large, which the paper rounded up to 340M. (The MLM head's dense layer and output bias add another 0.6M and are dropped after pretraining.) Notice that about a fifth of BERT-base is the embedding table, V×d=23.4V \times d = 23.4M parameters, the same table from §1.2. Each layer has about 12d2=7.112d^2 = 7.1M parameters, matching the 12Ld212 L d^2 rule from §5.6. BERT was pretrained on BooksCorpus (800M words) and English Wikipedia (2,500M words) for 1M steps of 256 sequences, about 40 epochs, using length-128 sequences for the first 90 percent of steps to save compute (attention cost grows with the square of length) and length 512 for the rest.

RoBERTa is BERT trained properly

RoBERTa (Liu et al., 2019) kept BERT's architecture and asked a humbler question than most papers. What if BERT was just undertrained? Every change was to the recipe.

changeBERTRoBERTawhy it helps
maskingstatic, fixed in preprocessingdynamic, new mask every timemore varied targets for the same data
NSPyes, segment pairsremoved, full contiguous sentencesNSP was too easy, long inputs teach long-range context
batch256 sequences8K sequencesbetter perplexity and end-task accuracy at equal compute, and easier to parallelize
data16 GB160 GB (adds CC-News, OpenWebText, Stories)more and more diverse text
tokenizerWordPiece, 30Kbyte-level BPE, 50Kno unknown tokens (§2.3)
training1M stepsup to 500K steps at 8K batchfar more total tokens

The lesson generalizes far beyond BERT. Before you credit an architecture or a clever objective, check whether the baseline was trained with a comparable budget and recipe. You'll see the same story at a larger scale in §9.3, where Chinchilla showed that the large language models of its day had been trained on far too little data for their size.

After RoBERTa, from ELECTRA to ModernBERT

A few successors are worth knowing by name, because each fixes a specific weakness.

ELECTRA (Clark et al., 2020) attacks MLM's waste. Only 15 percent of tokens give a learning signal. ELECTRA trains a small MLM "generator" to fill in masked tokens with plausible guesses, and trains the main model, the "discriminator", to say for every token whether it's original or replaced. Every position now contributes to the loss, and the discriminator never sees [MASK], which removes the mismatch entirely. At equal compute, it beat MLM-trained models, by the largest margin for small models.

DeBERTa (He et al., 2021) separates content and position in attention, computing attention scores from content-to-content, content-to-position, and position-to-content terms, a cousin of the relative position ideas in §5.5.

ModernBERT (Warner et al., 2024) brings an encoder up to date with everything decoders learned since 2019. It's a strong default when you need an encoder that is fast and handles long inputs. Here's what changed, all from the paper.

  • Trained on 2 trillion tokens, including code, with a modern BPE tokenizer of 50,368 tokens (a multiple of 64 for GPU efficiency).
  • A native context of 8,192 tokens, using RoPE (§5.5). Most of the training ran at 1,024 tokens, and the context was then extended by raising the RoPE base of the global layers from 10,000 to 160,000 and training on another 300 billion tokens.
  • Alternating attention. Every third layer uses full global attention, and the others use a 128-token sliding window (§12.6). In the authors' ablation this gave the same downstream performance as global attention in every layer, at a much lower cost for long inputs.
  • GeGLU feed-forward layers (§5.3) and pre-normalization. Bias terms are removed from every linear layer except the final decoder layer and from every LayerNorm, which spends more of the parameter budget on the weights themselves.
  • Unpadding. Padding tokens are removed and all sequences in a batch are concatenated into one, with FlashAttention's variable-length kernels keeping them separate (§12.4).
  • MLM only, no NSP, with a 30 percent masking rate. Wettig et al. (2023) had shown that 15 percent is not optimal, and that masking 40 percent beats 15 percent for BERT-large-sized models. More masked positions mean more predictions per sequence, at the price of less visible context.

Two sizes were released, base (149M parameters) and large (395M).

Encoder or decoder?

If decoders can do everything, why keep encoders? Because for a fixed budget, they are much cheaper at the jobs they're good at.

encoder (BERT family)decoder (GPT family)
attentionbidirectionalcausal
pretraining signalMLM on about 15 to 30% of tokensnext token on every token
natural outputsone vector per token, a pooled vector per texta distribution over the next token
strengthsclassification, tagging, span extraction, embeddings, reranking, at low latencygeneration, in-context learning, instruction following
typical size in production20M to 400M parameters1B to many hundreds of billions
can it generate textnot naturallyyes

Decoders can be turned into excellent embedders too (§7.3 shows how). But when you need to classify a billion documents a day, or embed every query to a search engine in a few milliseconds, a 100M-parameter encoder is usually the answer.

In PyTorch

Here is a compact BERT-style model with an MLM head whose decoder is tied to the input embeddings, trained for one step with the standard ignore_index=-100 convention. It uses nn.TransformerEncoder, which has no causal mask unless you pass one.

🔥 In PyTorchfor a machine with PyTorch

With Hugging Face, fine-tuning a pretrained encoder for classification takes a few lines. This needs a download, so it's shown and not run.

Read more

Watch next

🎬 BERT explained: Training, Inference, BERT vs GPT/LLamA, Fine tuning, [CLS] token · Umar JamilA careful whiteboard walk through masked language modeling, next sentence prediction, the CLS token, and fine-tuning, set side by side with decoder models.
🎬 BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding · Yannic KilcherA paper reading of the original BERT that explains why bidirectional pretraining needs masking and how the two pretraining tasks are built.
🎬 Encoder-Only Transformers (like BERT) for RAG, Clearly Explained!!! · StatQuest with Josh StarmerA gentle, step by step picture of how an encoder turns a sentence into context-aware embeddings, a good warm-up before the math.