Prefix Language Modeling › Implementing Prefix LM Training
Language AI HandbookAll chapters Prefix Language Modeling › Implementing Prefix LM TrainingOne section · Sections in this chapter6 / 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 the mathematical foundation and attention pattern in place, we can now build a complete prefix LM from scratch. This implementation brings together everything we've discussed: the hybrid attention mask, the target-only loss computation, and the conditioning structure that makes prefix LM unique.
We'll construct the model in three stages: first the attention mechanism that applies our L-shaped mask, then the full transformer architecture, and finally the loss function that focuses learning on target positions.
The Attention Mechanism
The core innovation happens in the attention layer. We need to apply the prefix LM mask during the attention computation. This keeps information flows according to our bidirectional-prefix, causal-target pattern:
class PrefixLMAttention(nn.Module):
"""Multi-head attention with prefix LM masking."""
def __init__(self, d_model, n_heads):
super().__init__()
self.d_model = d_model
self.n_heads = n_heads
self.head_dim = d_model // n_heads
self.q_proj = nn.Linear(d_model, d_model)
self.k_proj = nn.Linear(d_model, d_model)
self.v_proj = nn.Linear(d_model, d_model)
self.out_proj = nn.Linear(d_model, d_model)
def forward(self, x, prefix_len):
batch_size, seq_len, _ = x.shape
# Project to Q, K, V
q = self.q_proj(x).view(
batch_size, seq_len, self.n_heads, self.head_dim
)
k = self.k_proj(x).view(
batch_size, seq_len, self.n_heads, self.head_dim
)
v = self.v_proj(x).view(
batch_size, seq_len, self.n_heads, self.head_dim
)
# Transpose for attention: (batch, heads, seq, head_dim)
q = q.transpose(1, 2)
k = k.transpose(1, 2)
v = v.transpose(1, 2)
# Compute attention scores
scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim**0.5)
# Apply prefix LM mask
mask = create_prefix_lm_mask(seq_len, prefix_len).to(x.device)
scores = scores + mask
# Softmax and apply to values
attn_weights = F.softmax(scores, dim=-1)
output = torch.matmul(attn_weights, v)
# Reshape and project
output = (
output.transpose(1, 2)
.contiguous()
.view(batch_size, seq_len, self.d_model)
)
return self.out_proj(output)The key line is scores = scores + mask. By adding negative infinity to positions we want to block, the subsequent softmax converts those positions to zero probability. The mask is created dynamically based on prefix_len, allowing different sequences in a batch to have different split points.
The Complete Model
Now we wrap this attention mechanism in a full transformer architecture. The model includes embeddings, multiple transformer layers with our prefix LM attention, and an output projection to vocabulary logits:
class TinyPrefixLM(nn.Module):
"""Minimal prefix language 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)
# Attention layers
self.attention_layers = nn.ModuleList(
[PrefixLMAttention(d_model, n_heads) for _ in range(n_layers)]
)
# Feed-forward layers
self.ff_layers = nn.ModuleList(
[
nn.Sequential(
nn.Linear(d_model, d_model * 4),
nn.GELU(),
nn.Linear(d_model * 4, d_model),
nn.Dropout(0.1),
)
for _ in range(n_layers)
]
)
# Layer norms
self.attn_norms = nn.ModuleList(
[nn.LayerNorm(d_model) for _ in range(n_layers)]
)
self.ff_norms = nn.ModuleList(
[nn.LayerNorm(d_model) for _ in range(n_layers)]
)
# Output projection
self.output = nn.Linear(d_model, vocab_size)
def forward(self, x, prefix_len):
batch_size, seq_len = x.shape
# Embeddings
positions = torch.arange(seq_len, device=x.device).unsqueeze(0)
h = self.token_emb(x) + self.pos_emb(positions)
h = self.layer_norm(h)
h = self.dropout(h)
# Transformer layers with prefix LM attention
for attn, ff, attn_norm, ff_norm in zip(
self.attention_layers,
self.ff_layers,
self.attn_norms,
self.ff_norms,
):
# Self-attention with residual
h = h + attn(attn_norm(h), prefix_len)
# Feed-forward with residual
h = h + ff(ff_norm(h))
# Project to vocabulary
logits = self.output(h)
return logits
def compute_prefix_lm_loss(logits, targets, prefix_len):
"""Compute loss only on target positions (after prefix)."""
# Only compute loss on positions after the prefix
target_logits = logits[:, prefix_len:, :]
target_labels = targets[:, prefix_len:]
# Flatten and compute cross-entropy
vocab_size = logits.size(-1)
loss = F.cross_entropy(
target_logits.reshape(-1, vocab_size),
target_labels.reshape(-1),
ignore_index=-100,
)
return lossNotice how compute_prefix_lm_loss implements the target-only loss from our mathematical formulation. By slicing logits[:, prefix_len:, :], we extract only the predictions for target positions. The model produces logits for all positions (including the prefix), but we discard the prefix predictions when computing the loss, focusing the learning signal entirely on the generation task.
Model parameters: 670,184 Architecture: - Embedding dim: 128 - Attention heads: 4 - Layers: 2 - Vocabulary: 1000
The model applies prefix LM attention in each layer, passing the prefix length to determine where the bidirectional/causal boundary lies. This is more flexible than hardcoding the mask, as different examples in a batch could have different prefix lengths.