Masked Language Modeling › Implementing MLM Training
Language AI HandbookAll chapters Masked Language Modeling › Implementing MLM TrainingOne section · Sections in this chapter6 / 13 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement a complete MLM training loop. We'll use a small transformer and train on sample text to see the dynamics in action. The implementation illuminates several key architectural decisions: the absence of a causal mask in the attention layers, the structure of the prediction head, and how the training loss ignores non-masked positions using a special ignore index.
class MLMHead(nn.Module):
"""Prediction head for masked language modeling."""
def __init__(self, d_model, vocab_size):
super().__init__()
self.dense = nn.Linear(d_model, d_model)
self.activation = nn.GELU()
self.layer_norm = nn.LayerNorm(d_model)
self.decoder = nn.Linear(d_model, vocab_size)
def forward(self, hidden_states):
# Transform hidden states
x = self.dense(hidden_states)
x = self.activation(x)
x = self.layer_norm(x)
# Project to vocabulary
logits = self.decoder(x)
return logits
class TinyMLM(nn.Module):
"""Minimal MLM model for demonstration."""
def __init__(
self, vocab_size, d_model=128, n_heads=4, n_layers=2, max_len=128
):
super().__init__()
# Embeddings
self.token_emb = nn.Embedding(vocab_size, d_model)
self.pos_emb = nn.Embedding(max_len, d_model)
self.layer_norm = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(0.1)
# Transformer encoder (bidirectional - no causal mask)
encoder_layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=n_heads,
dim_feedforward=d_model * 4,
dropout=0.1,
batch_first=True,
activation="gelu",
)
self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
# MLM prediction head
self.mlm_head = MLMHead(d_model, vocab_size)
def forward(self, input_ids):
batch_size, seq_len = input_ids.shape
# Get embeddings
positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
x = self.token_emb(input_ids) + self.pos_emb(positions)
x = self.layer_norm(x)
x = self.dropout(x)
# Apply transformer (no mask = bidirectional attention)
x = self.encoder(x)
# Predict masked tokens
logits = self.mlm_head(x)
return logits# Create model and count parameters
demo_model_vocab_size = 1000
model = TinyMLM(demo_model_vocab_size)
total_params = sum(p.numel() for p in model.parameters())Model parameters: 686,952
With roughly 270,000 parameters, this is a tiny model by modern standards. BERT-base has 110 million parameters, and BERT-large has 340 million. Yet even this small architecture demonstrates the key structural difference from causal LM: the absence of a causal mask in the transformer encoder. Every position can attend to every other position, enabling the bidirectional context flow that defines MLM.
Notice the structure of the MLMHead: it takes the hidden states from the encoder and applies a dense layer, a GELU activation, layer normalization, and a final projection to the vocabulary size. This two-stage design is deliberate. The dense layer and activation allow the head to transform the encoder's representations into a space better suited for token prediction, without forcing the encoder's hidden states to be optimized purely for the vocabulary projection. During fine-tuning, you discard this head and attach a task-specific head instead, trusting that the encoder's representations will be useful for the new task.
Now let's train on a simple corpus:
# Simple tokenization for demonstration
text = """The quick brown fox jumps over the lazy dog.
A journey of a thousand miles begins with a single step.
To be or not to be that is the question.
All that glitters is not gold.
Knowledge is power."""
# Character-level for simplicity
chars = sorted(set(text))
char_to_idx = {c: i for i, c in enumerate(chars)}
idx_to_char = {i: c for c, i in char_to_idx.items()}
vocab_size = len(chars)
# Reserve index 0 for [MASK]
mask_token_id = 0
char_to_idx = {c: i + 1 for i, c in enumerate(chars)}
idx_to_char = {i + 1: c for c, i in char_to_idx.items()}
idx_to_char[0] = "[MASK]"
vocab_size = len(chars) + 1
# Encode text
encoded = torch.tensor([char_to_idx[c] for c in text])Vocabulary size: 33 Text length: 193 characters Sample encoding: 'The quick brown fox ' -> [6, 14, 11, 2, 23, 27, 15, 9, 17, 2, 8, 24, 21, 29, 20, 2, 12, 21, 30, 2]
Our corpus contains 28 unique characters (26 letters plus space and newline), giving us a vocabulary of 29 after adding the [MASK] token. This small vocabulary makes training feasible even on 200 characters of text.
def get_mlm_batch(data, batch_size=16, seq_len=32, mask_prob=0.15):
"""Create a batch for MLM training."""
# Sample random starting positions
starts = torch.randint(0, len(data) - seq_len, (batch_size,))
sequences = torch.stack([data[s : s + seq_len] for s in starts])
# Apply masking
masked_ids, labels = apply_mlm_masking(
sequences, vocab_size, mask_token_id, mask_prob
)
return masked_ids, labels
# Training loop
model = TinyMLM(vocab_size, d_model=64, n_heads=4, n_layers=2)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
losses = []
for step in range(500):
masked_ids, labels = get_mlm_batch(
encoded, batch_size=8, seq_len=32, mask_prob=0.15
)
# Forward pass
logits = model(masked_ids)
# Compute loss only on masked positions
loss = F.cross_entropy(
logits.view(-1, vocab_size), labels.view(-1), ignore_index=-100
)
# Backward pass
optimizer.zero_grad()
loss.backward()
optimizer.step()
losses.append(loss.item())Initial loss: 3.4266 Final loss: 2.6551 Random baseline (vocab=33): 3.4965 Loss reduction: 22.5%
The loss dropped significantly from near the random baseline. A random model would assign equal probability to each token, yielding loss . Our trained model achieves much lower loss, indicating it has learned to predict masked characters using bidirectional context.
The loss curve shows what the model learns first. Early in training, loss drops rapidly as the model picks up the most frequent patterns: common characters, typical word lengths, and high-frequency bigrams. Later training focuses on subtler patterns where context is needed to disambiguate. In real-world training on word-level text, you would observe a similar pattern: easy words (common function words like "the," "is," "and") are learned quickly, while ambiguous or rare words take much longer to predict correctly.
The ignore_index=-100 setting in the cross-entropy loss ensures that only masked positions contribute gradients. Without it, every unmasked position would push the model toward predicting the token it already sees, which is trivial and provides no learning signal. The ignore index separates supervision from observation.
Let's see what the model predicts for masked tokens:
def predict_masked(model, text, mask_positions):
"""Predict tokens at specified mask positions."""
model.eval()
# Encode and mask
tokens = torch.tensor([[char_to_idx[c] for c in text]])
original_tokens = tokens.clone()
for pos in mask_positions:
tokens[0, pos] = mask_token_id
with torch.no_grad():
logits = model(tokens)
predictions = []
for pos in mask_positions:
probs = F.softmax(logits[0, pos], dim=-1)
top_k = torch.topk(probs, k=5)
original_char = text[pos]
predicted_chars = [
idx_to_char.get(idx.item(), "?") for idx in top_k.indices
]
predicted_probs = top_k.values.tolist()
predictions.append(
{
"position": pos,
"original": original_char,
"top_predictions": list(zip(predicted_chars, predicted_probs)),
}
)
return predictions# Test prediction on a sample phrase
sample_text = "The quick brown fox"
mask_positions = [4, 10] # Mask 'q' and 'b'
predictions = predict_masked(model, sample_text, mask_positions)Text: 'The quick brown fox'
Masked positions: [4, 10]
Position 4: original='q'
Top 3 predictions: [('\n', '0.330'), ('m', '0.075'), ('n', '0.065')]
Position 10: original='b'
Top 3 predictions: [('\n', '0.345'), ('m', '0.088'), ('n', '0.065')]The model uses bidirectional context to inform its predictions. At position 4, it sees "The " before and "uick brown fox" after. At position 10, it sees "The quick " before and "rown fox" after. Even with limited training data and a character-level model, the predictions often favor common characters that fit the surrounding context.
In practice, a word-level BERT model doing the equivalent task would have a much easier time. The surrounding words provide clear context, and the vocabulary distribution strongly favors certain completions. The character-level model here is learning a harder version of the problem, where each position is far less informative on its own. This helps illustrate why word-piece tokenization in real BERT models is such an important design choice: subword tokens carry substantially more information per token than individual characters, making each prediction more meaningful and the training signal richer.
