Lesson 12

Encoder-only models and masked-token objectives

Why hide tokens?

BERT learns context by predicting selected tokens that have been hidden.

1. Original text: the cat went home
2. Hide a token: the cat [MASK] home
3. Read both sides: use the cat on the left and home on the right.
4. Predict the hidden token: compare the prediction with went.

Reading both sides is called bidirectional context. This helps with tasks such as classifying text, searching, and finding names.

BERT and GPT ask different questions

• BERT: “What belongs in this gap?” Both sides are visible; the hidden original token is not.
• GPT: “What comes next?” Future tokens are hidden.

In this small BERT-style exercise, calculate loss only at the selected masked positions.

Important terms

Follow the shapes

1. Token IDs: (B,T) — B sequences, T positions.
2. Token + position embeddings: (B,T,D) — D features per position.
3. Encoder output: (B,T,D) — same shape, now with context.
4. Prediction scores: (B,T,V) — V vocabulary choices per position.
5. Loss: compare scores and targets only at masked positions.

Worked example: tiny masked encoder

This small CPU model teaches the data flow; it is not a useful pretrained BERT. Token ID 5 is reserved for [MASK]. Unlike GPT, the transformer receives no causal mask.

import torch
from torch import nn

class TinyMaskedModel(nn.Module):
    def __init__(self, vocabulary_size, maximum_length, model_size=12):
        super().__init__()
        self.token_embedding = nn.Embedding(vocabulary_size, model_size)
        self.position_embedding = nn.Embedding(maximum_length, model_size)
        layer = nn.TransformerEncoderLayer(
            d_model=model_size,
            nhead=3,
            dim_feedforward=24,
            dropout=0.0,
            batch_first=True,
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=2)
        self.language_head = nn.Linear(model_size, vocabulary_size)

    def forward(self, token_ids):
        length = token_ids.shape[1]
        positions = torch.arange(length, device=token_ids.device)
        x = self.token_embedding(token_ids)
        x = x + self.position_embedding(positions)
        encoded = self.encoder(x)
        return self.language_head(encoded)

torch.manual_seed(12)
model = TinyMaskedModel(vocabulary_size=6, maximum_length=5)
mask_id = 5
inputs = torch.tensor([[0, 1, mask_id, 3, 4]], dtype=torch.long)
logits = model(inputs)

print(tuple(inputs.shape))
print(tuple(logits.shape))

# Expected:
# (1, 5)
# (1, 5, 6)

The representation at position two can collect information from positions on both sides. The language head produces predictions everywhere, but we select the masked location for this simple example.

Loss only where the answer was hidden

Count before calculating: if a sequence has six positions and two are masked, the masked-language loss should average two token losses. Averaging all six positions gives four easy or irrelevant terms extra influence.

original_token_id = torch.tensor([2])
masked_logits = logits[:, 2, :]
loss = nn.functional.cross_entropy(masked_logits, original_token_id)
loss.backward()

print(tuple(masked_logits.shape))  # Expected: (1, 6)
print(loss.ndim)                   # Expected: 0

A full data pipeline chooses several mask positions in each sequence and records their original IDs. Unmasked or padding positions are commonly assigned an ignore value in the target tensor so they do not affect cross-entropy. If the original token remains visible at the masked position, the task leaks its own answer.

A suspicious symptom of leakage is near-perfect training accuracy almost immediately. Inspect the actual model input at each scored position. The target token must be replaced or otherwise hidden there.

From pretraining to classification

A pretrained encoder already produces context-aware vectors. For text classification, a small output layer can map one pooled sequence vector to class logits. A common design uses a special first token whose encoded vector represents the sequence. Another simple choice is a masked mean over real token positions.

encoded_features = torch.randn(4, 12)  # four pooled examples
classification_head = nn.Linear(12, 3)
class_logits = classification_head(encoded_features)
labels = torch.tensor([0, 2, 1, 0])
classification_loss = nn.functional.cross_entropy(class_logits, labels)

print(tuple(class_logits.shape))  # Expected: (4, 3)
print(classification_loss.ndim)   # Expected: 0

In a real project, use a trusted pretrained checkpoint and its matching tokenizer rather than training a full BERT from random values. Begin with a small baseline, track validation metrics, and consider freezing the encoder before fine-tuning all layers.

Try it

A batch has eight sequences of length 20, model width 48, and vocabulary size 1000. State the encoder-output and masked-language-logit shapes. If there are 12 masked positions in total, what shape should the selected logits have before cross-entropy?

Recap

BERT is an encoder-only transformer trained to reconstruct hidden tokens using context on both sides. Its input and encoder shapes are (B,T) and (B,T,D); a language head produces (B,T,V). Compute pretraining loss only at chosen masked positions. For practical classification, start from a pretrained encoder and add a small task head.

References: PyTorch TransformerEncoder and encoder layer documentation.