Prefix Language Modeling › The Prefix LM Attention Pattern
Language AI HandbookAll chapters Prefix Language Modeling › The Prefix LM Attention PatternOne section · Sections in this chapter4 / 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
The mathematical objective we just defined has a direct physical manifestation: the attention mask. This mask controls which tokens can "see" which other tokens during the forward pass, and it must encode our two key requirements: bidirectional prefix and causal target.
Unlike the uniform lower-triangular mask of causal LM or the full-attention mask of masked LM, prefix LM uses a hybrid pattern that changes behavior partway through the sequence.
Anatomy of the Attention Mask
The mask divides naturally into three regions, each serving a distinct purpose:
- Prefix-to-prefix (upper-left quadrant): Full bidirectional attention. Every prefix token attends to every other prefix token, and can also attend to all target positions. This is what gives the prefix its representational richness, as each word can incorporate context from the entire sequence.
- Target-to-prefix (lower-left rectangle): Full attention. Every target token attends to all prefix tokens. This is how the conditioning flows, when generating "Le" in our translation example, the model can look back at any part of "Translate to French: The cat sat on the mat."
- Target-to-target (lower-right quadrant): Causal attention. Each target token attends only to itself and preceding target tokens. This preserves the autoregressive property needed for generation, you can't peek at future words you haven't generated yet.
Together, these regions create the characteristic L-shaped pattern visible in the mask visualization. The prefix forms a fully connected subgraph where information flows freely in all directions, while the target maintains causal structure that enables token-by-token generation.
Implementing the Mask
Translating this pattern into code requires building a matrix where each entry indicates whether position can attend to position :
def create_prefix_lm_mask(seq_len, prefix_len):
"""
Create attention mask for prefix language modeling.
Args:
seq_len: Total sequence length
prefix_len: Length of the prefix (bidirectional) region
Returns:
Mask where 0 indicates attend, -inf indicates block
"""
# Start with zeros (all attend)
mask = torch.zeros(seq_len, seq_len)
# For target positions (after prefix), apply causal masking
for i in range(prefix_len, seq_len):
for j in range(i + 1, seq_len):
mask[i, j] = float("-inf")
return maskPrefix LM mask (8 positions, prefix_len=4): (0 = attend, -inf = block) Pos 0: [ 0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 ] Pos 1: [ 0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 ] Pos 2: [ 0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 ] Pos 3: [ 0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 ] Pos 4: [ 0 , 0 , 0 , 0 , 0 , -inf, -inf, -inf] Pos 5: [ 0 , 0 , 0 , 0 , 0 , 0 , -inf, -inf] Pos 6: [ 0 , 0 , 0 , 0 , 0 , 0 , 0 , -inf] Pos 7: [ 0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 ]
The first four positions (prefix) have all zeros, meaning they attend to everything. Positions 4-7 (target) have zeros up to and including their own position, then negative infinity to block future tokens.