Prefix-LM

Prefix-LM（Prefix Language Model）是一种混合 Attention 模式：prefix 内双向交互，后续生成保持 causal。

Attention 拓扑

mask_ar = [0, 0, 0, 1, 1, 1]
cumsum  = [0, 0, 0, 1, 2, 3]
          [prefix][causal]

Attention 矩阵：

          P1 P2 P3 A1 A2 A3
Query P1   1  1  1  0  0  0
      P2   1  1  1  0  0  0   ← Prefix 双向
      P3   1  1  1  0  0  0
      A1   1  1  1  1  0  0   ← 后续 causal
      A2   1  1  1  1  1  0
      A3   1  1  1  1  1  1

信息流：

┌────────────┐
│ P1 ↔ P2 ↔ P3│  双向交互
└────────────┘
      ↓
     A1         单向依赖
      ↓
     A2
      ↓
     A3

为什么需要 Prefix 双向？

Prefix 通常包含条件信息（prompt、image、context），它们的角色是：

"共同构成当前任务的 conditioning context"

而不是"正在逐个生成它们"。

因此希望 prefix 内充分交互：

image tokens ↔ text tokens ↔ state tokens

而后续生成的 action/text 保持 autoregressive。

与标准 Causal LM 的区别

	Causal LM	Prefix-LM
Attention	全局 causal（下三角）	Prefix 双向 + 后续 causal
信息流	每个 token 只能看过去	Prefix 内互相看到
应用	纯文本生成	条件生成（多模态、ICL）

实现

Prefix-LM 是 Block-Causal Attention 的特例：

# Prefix 在 block 0（内部双向）
# 后续每个 token 一个 block（causal）
mask_ar = [0, ..., 0, 1, 1, 1, ...]
          └prefix┘ └─causal─┘

通过 cumsum(mask_ar) 自动实现，无需分别处理。

多模态应用

在 PaliGemma / π0 这类模型中：

prefix = [image_tokens, text_tokens, state_tokens]
suffix = [action_tokens]

Prefix 内充分交互（图像-文本理解）
Action 逐步生成（保持 causal）

这不是"causal vs non-causal"的开关，而是定义了信息流拓扑图。

训练 vs 推理

训练

Prefix + 完整目标序列一次性 forward：

[P1 P2 P3 A1 A2 A3]

Mask 确保训练时 A2 看不到 A3。

推理

Prefix 一次性编码，后续 autoregressive 生成：

Encode: [P1 P2 P3]
Generate: A1 → A2 → A3

利用 KV Cache 避免重复计算 prefix。

Prefix-LM

Prefix-LM

Attention 拓扑

为什么需要 Prefix 双向？

与标准 Causal LM 的区别

实现

多模态应用

训练 vs 推理

训练

推理

相关概念