Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Prefix Language Modeling › Generation with Prefix LM

Michael BrenndoerferUpdated July 14, 20255 min read

Language AI HandbookAll chapters

Part 22 of 67Pre-training Objectives

Chapter 159Prefix Language Modeling › Generation with Prefix LM

Reading section 8 of 16 · Generation with Prefix LM · 5 minSection 8 of 16 · 5 min

Section 8 of 16 · 5 min read

Prefix Language Modeling › Generation with Prefix LMOne section · Sections in this chapter8 / 16 ↓

Sections in this chapter16

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 generate text using our trained prefix LM. The key difference from causal generation is that we start with a bidirectional prefix:

In[19]:

Code

def generate_with_prefix(
    model, prefix_text, max_new_tokens=50, temperature=0.8
):
    """Generate continuation given a prefix."""
    model.eval()

    # Encode prefix
    prefix_tokens = [char_to_idx[c] for c in prefix_text]
    prefix_len = len(prefix_tokens)
    tokens = torch.tensor(prefix_tokens).unsqueeze(0)

    with torch.no_grad():
        for _ in range(max_new_tokens):
            # Get predictions (prefix_len stays fixed, sequence grows)
            logits = model(tokens, prefix_len)

            # Sample from last position
            next_logits = logits[0, -1, :] / temperature
            probs = F.softmax(next_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)

            # Append to sequence
            tokens = torch.cat([tokens, next_token.unsqueeze(0)], dim=1)

    # Decode
    generated = "".join([idx_to_char[t.item()] for t in tokens[0]])
    return generated, prefix_text

Out[20]:

Console

Prefix LM Generation Examples:
--------------------------------------------------
Prefix: 'The quick'
Full: The quick
A t atho be the thes question.
Alll gl 

Prefix: 'A journey of'
Full: A journey ofaa mila begins thogitho.
Ale ith nonours

Prefix: 'Knowledge is'
Full: Knowledge is p.
To be or nor no be nor nothat is le

The generation process treats the prefix specially: those tokens can attend to each other bidirectionally throughout generation. As new tokens are added, they see the full prefix context plus their left context, following the prefix LM pattern.

Visualizing Learned Attention PatternsLink Copied

To see the prefix LM attention pattern in action, we can extract and visualize the attention weights from our trained model. This reveals how the model distributes attention across prefix and target positions:

Out[21]:

Visualization

Heatmap showing attention weights with strong bidirectional pattern in prefix region and causal pattern in target region. — Attention weights from the trained prefix LM model (Layer 1, Head 1). The L-shaped pattern emerges clearly: prefix positions (0-9) attend broadly to each other, while target positions (10+) focus on the prefix and their left context. Brighter colors indicate stronger attention weights.

The attention heatmap reveals the prefix LM pattern in action. Notice how positions in the prefix region (left of the red dashed line) distribute attention broadly across all prefix positions, indicating bidirectional processing. Target positions (right of the line) show the characteristic causal pattern: each attends heavily to the prefix and to preceding target positions, but attention to future positions is blocked.

Out[22]:

Visualization

Bar chart showing attention distribution for a prefix position with relatively uniform weights across all positions. — Prefix position (position 5) distributes attention broadly across all positions it can see, with roughly uniform weights. This bidirectional access allows prefix tokens to incorporate context from the entire prefix.

Bar chart showing attention distribution for a target position with weights only on prefix and prior target positions. — Target position (position 15) concentrates attention on the prefix (positions 0-9) and accessible target positions (10-15), with zero weight on future positions (16+). The attention is blocked by the causal mask.

Comparing these attention distributions highlights the fundamental difference between prefix and target positions. The prefix position distributes attention across all accessible positions (itself and all others in the sequence), while the target position can only attend to the prefix and prior target positions, with future positions completely masked out.