Prefix Language Modeling › Training on Prefix-Target Pairs
Language AI HandbookAll chapters Prefix Language Modeling › Training on Prefix-Target PairsOne section · Sections in this chapter7 / 16 ↓
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
With our model architecture complete, we can now see prefix LM in action. We'll train on a simple task: given the first part of a sentence (the prefix), generate a plausible continuation (the target). Character-level tokenization keeps things interpretable, letting us watch the model learn character-by-character patterns.
Preparing the Data
First, we create a small corpus and tokenize it at the character level:
# Training corpus
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 and power is knowledge."""
# Character-level tokenization
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)
# Encode text
encoded = torch.tensor([char_to_idx[c] for c in text])Vocabulary size: 32 characters Text length: 216 characters Sample: 'The quick brown fox jumps over the lazy dog. A jou...'
This tiny corpus is far smaller than what you'd use for real pretraining, but it's enough to demonstrate the mechanics. The vocabulary of unique characters becomes our token space.
The Training Loop with Random Splits
A key design choice in prefix LM training is how to select the split point. We could fix it (always split at position 10), but that would limit generalization. Instead, we randomize the prefix length within a reasonable range. This teaches the model to handle varying amounts of context:
def get_prefix_lm_batch(
data, batch_size=16, seq_len=48, min_prefix_frac=0.3, max_prefix_frac=0.7
):
"""
Create a batch for prefix LM training.
Each sequence is split at a random point into prefix and target.
The prefix fraction is randomly sampled between min and max.
"""
# 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])
# Random prefix length for this batch (same for all sequences in batch)
prefix_frac = (
torch.rand(1).item() * (max_prefix_frac - min_prefix_frac)
+ min_prefix_frac
)
prefix_len = int(seq_len * prefix_frac)
# Input is the full sequence, target is shifted by 1
input_ids = sequences[:, :-1]
labels = sequences[:, 1:]
# Adjust prefix_len for the shifted sequences
prefix_len = max(1, prefix_len - 1)
return input_ids, labels, prefix_len
# Training loop
model = TinyPrefixLM(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):
input_ids, labels, prefix_len = get_prefix_lm_batch(
encoded, batch_size=8, seq_len=48
)
# Forward pass
logits = model(input_ids, prefix_len)
# Compute loss only on target (non-prefix) positions
loss = compute_prefix_lm_loss(logits, labels, prefix_len)
# Backward pass
optimizer.zero_grad()
loss.backward()
optimizer.step()
losses.append(loss.item())Initial loss: 3.8022 Final loss: 0.3725 Random baseline: 3.4657 Loss reduction: 90.2%
The loss dropped substantially from the random baseline, indicating the model has learned meaningful patterns. A random model would assign equal probability to each token, yielding loss near the random baseline. The significant reduction shows the model is successfully using the bidirectional prefix context to improve target predictions.
The model learns to generate continuations conditioned on the prefix. Unlike standard causal LM where loss is computed everywhere, prefix LM focuses the learning signal on the generation task. The prefix is rich context, processed bidirectionally, that informs the target generation.
The per-position loss analysis reveals an important pattern: the first few target positions often achieve lower loss than later positions. This makes intuitive sense, the first target token has direct access to the full bidirectional prefix, while later tokens must rely increasingly on autoregressively generated context, which may be less informative than the original prefix.