Prefix Language Modeling › UniLM-Style Unified Training

Michael BrenndoerferUpdated July 14, 20254 min read
Language AI HandbookAll chapters
Part 22 of 67Pre-training Objectives
Chapter 159Prefix Language Modeling › UniLM-Style Unified Training
Reading section 9 of 16 · UniLM-Style Unified Training · 4 minSection 9 of 16 · 4 min

Book contents

Language AI Handbook67 parts · 513 chaptersView the whole book →
Active chapter 159 of 513

Section 9 of 16 · 4 min read

Prefix Language Modeling › UniLM-Style Unified TrainingOne section · Sections in this chapter9 / 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

A powerful extension of prefix LM is training a single model on multiple objectives simultaneously. UniLM (Unified Language Model) pioneered this approach, training on three objectives in each batch:

  1. Bidirectional LM (like MLM): Full bidirectional attention, predict masked tokens
  2. Unidirectional LM (like CLM): Causal attention, predict next token
  3. Prefix LM: Bidirectional prefix, causal target
Out[23]:
Visualization
Full attention mask for bidirectional LM showing all positions attend to all.
Bidirectional uses full attention for understanding tasks where every position attends to every other position.
Lower triangular attention mask for unidirectional causal LM.
Unidirectional uses causal attention for generation where each position only attends to earlier positions.
L-shaped attention mask for prefix LM with bidirectional prefix and causal target.
Prefix LM combines both patterns for conditioned generation, with bidirectional prefix and causal target.

The unified training produces a model capable of all three modes. At inference time, you select the appropriate attention pattern for your task. This flexibility makes UniLM-style models particularly versatile.

In[24]:
Code
def create_unified_mask(seq_len, mode, prefix_len=None):
    """
    Create attention mask for different LM modes.

    Args:
        seq_len: Sequence length
        mode: 'bidirectional', 'unidirectional', or 'prefix'
        prefix_len: Required for 'prefix' mode

    Returns:
        Attention mask (0 = attend, -inf = block)
    """
    mask = torch.zeros(seq_len, seq_len)

    if mode == "bidirectional":
        # All attend to all
        pass

    elif mode == "unidirectional":
        # Causal masking
        for i in range(seq_len):
            for j in range(i + 1, seq_len):
                mask[i, j] = float("-inf")

    elif mode == "prefix":
        # Prefix is bidirectional, target is causal
        if prefix_len is None:
            raise ValueError("prefix_len required for prefix mode")
        for i in range(prefix_len, seq_len):
            for j in range(i + 1, seq_len):
                mask[i, j] = float("-inf")

    return mask
Out[25]:
Console
Unified attention masks for sequence length 6:

BIDIRECTIONAL:
  Attend positions: 36, Block positions: 0
UNIDIRECTIONAL:
  Attend positions: 21, Block positions: 15
PREFIX:
  Attend positions: 33, Block positions: 3

Bidirectional mode allows all 36 position pairs to attend to each other. Unidirectional mode blocks 15 positions (the upper triangle), allowing only 21 attend pairs. Prefix LM falls between these extremes: with a prefix of 3 tokens, it allows more attention than causal but less than bidirectional, as the last 3 positions still use causal masking.