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.
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 from attending to anything after it. Remove the mask and the task becomes trivial. In the very first layer, position can attend straight to token 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.
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 drawn from a vocabulary of size . We pick a random set of positions , corrupt the input at those positions to get , run a bidirectional Transformer on , and ask it to recover the original tokens,
Here is the final hidden state at position (after a small dense layer with GELU and LayerNorm), is the input embedding table reused as the output layer (weight tying, as in §1.5 and §6.2), and 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.
- The loss covers only . Positions that weren't selected contribute nothing. In BERT, 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).
- Each masked token is predicted independently given . 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.
- The prediction head is cheap. The heavy lifting is in the encoder. After pretraining, the head is thrown away, and is the contextual embedding you keep (exactly the "contextual token embeddings" of §1.8 and §3.4).
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.
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 percent of all tokens, which BERT's authors found didn't hurt.
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] | same | random | MNLI (fine-tune) | NER (fine-tune) | NER (frozen features) |
|---|---|---|---|---|---|
| 80% | 10% | 10% | 84.2 | 95.4 | 94.9 |
| 100% | 0% | 0% | 84.3 | 94.9 | 94.0 |
| 80% | 0% | 20% | 84.1 | 95.2 | 94.6 |
| 80% | 20% | 0% | 84.4 | 95.2 | 94.7 |
| 0% | 20% | 80% | 83.7 | 94.8 | 94.6 |
| 0% | 0% | 100% | 83.6 | 94.9 | 94.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 tokens on average. Of those, about become [MASK], become random tokens, and stay as they were. So about 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.
MLM masking with the 80-10-10 rule
Implement mlm_mask(ids, special, vocab_size, mask_id, first_regular, rng, p=0.15).
idsis an int array of shape(B, T)andspeciala 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 fromrange(first_regular, vocab_size)with probability 0.1, and the original token with probability 0.1. - Return
(inputs, labels).inputsis the corrupted copy (do not modifyids), andlabelsholds the original id at selected positions and-100elsewhere.
No Python loops over tokens. This is a real interview question, and the tests check the statistics on a large batch.
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 be the final hidden states, with row written , and let be the row for the [CLS] token at position 1.
Sequence classification. With classes, add and predict . 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 through a tanh "pooler" layer trained on NSP.)
Token classification. For named entities or part-of-speech tags, put a classifier on every . 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 and an end vector . The probability that the answer starts at token is a softmax over passage positions,
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 that maximizes subject to 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 is not decoration. Picking the start and end by separate argmaxes is a classic bug.
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, scores , while scores , so the answer is tokens 2 to 4.
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.
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 , , or 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 layers, width , 12 heads, and a feed-forward width of . BERT-large has , , 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).
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, M parameters, the same table from §1.2. Each layer has about M parameters, matching the 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.
| change | BERT | RoBERTa | why it helps |
|---|---|---|---|
| masking | static, fixed in preprocessing | dynamic, new mask every time | more varied targets for the same data |
| NSP | yes, segment pairs | removed, full contiguous sentences | NSP was too easy, long inputs teach long-range context |
| batch | 256 sequences | 8K sequences | better perplexity and end-task accuracy at equal compute, and easier to parallelize |
| data | 16 GB | 160 GB (adds CC-News, OpenWebText, Stories) | more and more diverse text |
| tokenizer | WordPiece, 30K | byte-level BPE, 50K | no unknown tokens (§2.3) |
| training | 1M steps | up to 500K steps at 8K batch | far 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) | |
|---|---|---|
| attention | bidirectional | causal |
| pretraining signal | MLM on about 15 to 30% of tokens | next token on every token |
| natural outputs | one vector per token, a pooled vector per text | a distribution over the next token |
| strengths | classification, tagging, span extraction, embeddings, reranking, at low latency | generation, in-context learning, instruction following |
| typical size in production | 20M to 400M parameters | 1B to many hundreds of billions |
| can it generate text | not naturally | yes |
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.
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
- Devlin et al. (2019). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. Read §3.1, Appendix A.1 on masking, and Appendix C.2 for the ablation table quoted above.
- Liu et al. (2019). RoBERTa: A Robustly Optimized BERT Pretraining Approach. §4 is a model of careful ablation, covering masking, NSP, batch size, and tokenization.
- Lan et al. (2020). ALBERT: A Lite BERT for Self-supervised Learning of Language Representations. The argument for why NSP is easy, and sentence order prediction.
- Clark et al. (2020). ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators. Replaced token detection and the sample-efficiency argument.
- Wettig et al. (2023). Should You Mask 15% in Masked Language Modeling? Why higher masking rates can be better.
- Warner et al. (2024). Smarter, Better, Faster, Longer: A Modern Bidirectional Encoder for Fast, Memory Efficient, and Long Context Finetuning and Inference (ModernBERT). The modern encoder recipe, with ablations in the appendix.
- Jay Alammar. The Illustrated BERT, ELMo, and co. (How NLP Cracked Transfer Learning). A gentle visual companion to this section.