XLNet排列语言模型的训练复现：双向上下文的捕获方式与BERT的差异

XLNet排列语言模型的训练复现：双向上下文的捕获方式与BERT的差异

XLNet（Yang et al., NeurIPS 2019）通过排列语言模型（Permutation Language Modeling, PLM）捕获双向上下文，在多个NLP基准上超越了BERT。其核心创新在于：不是像BERT那样通过[MASK] token破坏输入来引入双向性，而是在自回归框架下通过对输入序列的分解顺序进行排列来实现对上下文的双向感知。本文从数学原理和代码实现两个层面复现PLM的训练过程，重点分析排列机制如何通过双流自注意力（Two-Stream Self-Attention）实现"看到上下文但不知道自己的位置"这一关键特性。

一、BERT的Masking与XLNet的排列：两种双向性方案

BERT通过Masked Language Modeling（MLM）来学习双向表示：随机选择15%的token用[MASK]替换，模型基于未Mask的上下文预测被Mask的token。这一方法有两个被广泛讨论的局限性：

预训练-微调不一致：预训练时模型看到[MASK] token，微调时没有[MASK]，造成输入分布偏移。

独立性假设：BERT假设被Mask的token之间相互独立。对于"New York is a [MASK]"和"[MASK] York is a city"两个Mask，BERT独立预测每个[MASK]，但在排列语言模型中，预测"New"时可以知道"York"已经被预测。

XLNet的排列语言模型避免了这两个问题：它不引入[MASK] token，而是在所有可能的排列顺序上训练自回归模型。例如对于序列(x1, x2, x3, x4)，排列语言模型可能按顺序x3→x2→x4→x1逐token预测，预测x1时可以看到x2、x3、x4（因为它们在该排列中位于x1之前），从而以自回归方式捕获双向上下文。

二、双流自注意力的核心设计

排列语言模型面临一个看似矛盾的需求：在预测token $x_{z_t}$时，模型需要使用它之前所有token ${x_{z_1}, \\ldots, x_{z_{t-1}}}$的内容（content）以及当前token $x_{z_t}$的位置（position），但不能使用当前token的内容（因为那正是要预测的目标）。

双流自注意力通过维护两组独立的隐藏状态来解决这一问题：

内容流（Content Stream） $h_\\theta(\\mathbf{x}{\\mathbf{z}{\\leq t}})$：包含到当前token为止的所有token的内容信息。用于编码上下文，和标准Transformer的隐藏状态一样。

查询流（Query Stream） $g_\\theta(\\mathbf{x}{\\mathbf{z}{<t}}, z_t)$：仅包含所有前序token的内容信息 + 当前token的位置信息，但不包含当前token的内容。仅用于预测当前token。

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class TwoStreamSelfAttention(nn.Module):
"""
XLNet 双流自注意力的核心实现。
维护内容流（Content Stream）和查询流（Query Stream）两组状态。
"""

def __init__(
self,
hidden_dim: int = 768,
num_heads: int = 12,
dropout: float = 0.1,
):
super().__init__()
self.hidden_dim = hidden_dim
self.num_heads = num_heads
self.head_dim = hidden_dim // num_heads

# QKV 投影矩阵（内容流和查询流共享）
self.q_proj = nn.Linear(hidden_dim, hidden_dim)
self.k_proj = nn.Linear(hidden_dim, hidden_dim)
self.v_proj = nn.Linear(hidden_dim, hidden_dim)
self.o_proj = nn.Linear(hidden_dim, hidden_dim)

# 查询流的额外 Q 投影（使用位置嵌入而非内容嵌入）
self.q_proj_query_stream = nn.Linear(hidden_dim, hidden_dim)

def _rel_shift(self, x: torch.Tensor) -> torch.Tensor:
"""
相对位置编码的 shift 操作（XLNet 特有）。
将相对位置矩阵向左上方平移一位。
"""
# x shape: (batch, num_heads, seq_len, seq_len)
zero_pad = torch.zeros(
(x.shape[0], x.shape[1], x.shape[2], 1),
device=x.device, dtype=x.dtype
)
x_padded = torch.cat([zero_pad, x], dim=-1)
x_padded = x_padded.view(
x.shape[0], x.shape[1], x.shape[3] + 1, x.shape[2]
)
x = x_padded[:, :, 1:].view_as(x)
return x

def forward(
self,
content_h: torch.Tensor, # 内容流: (B, S, H)
query_h: torch.Tensor, # 查询流: (B, S, H)
attention_mask: torch.Tensor, # 排列注意力掩码
pos_emb: torch.Tensor, # 相对位置编码
) -> tuple:
"""
双流自注意力前向传播。

关键差异：
– 内容流：Q 来自 content_h（包含自己的内容）
– 查询流：Q 来自 query_h（不包含自己的内容，仅位置信息）
– 两个流共享 K 和 V
"""
B, S, H = content_h.shape

# === 内容流的注意力 ===
# Q 来自内容流（可以 "看到自己"）
Q_content = self._reshape_to_heads(
self.q_proj(content_h)
) # (B, num_heads, S, head_dim)

# K, V 始终来自内容流
K = self._reshape_to_heads(self.k_proj(content_h))
V = self._reshape_to_heads(self.v_proj(content_h))

# 计算内容流注意力
attn_scores_content = torch.matmul(
Q_content, K.transpose(-2, -1)
) / math.sqrt(self.head_dim) # (B, num_heads, S, S)

# 加上相对位置偏置
attn_scores_content = attn_scores_content + pos_emb

# 应用排列注意力掩码（确保自回归性）
attn_scores_content = attn_scores_content + attention_mask

attn_probs_content = F.softmax(attn_scores_content, dim=-1)
attn_probs_content = F.dropout(
attn_probs_content, p=0.1, training=self.training
)
content_output = torch.matmul(
attn_probs_content, V
) # (B, num_heads, S, head_dim)

# === 查询流的注意力 ===
# Q 来自查询流（不能 "看到自己" 的内容）
Q_query = self._reshape_to_heads(
self.q_proj_query_stream(query_h)
) # (B, num_heads, S, head_dim)

# 使用和内容流相同的 K, V
attn_scores_query = torch.matmul(
Q_query, K.transpose(-2, -1)
) / math.sqrt(self.head_dim)

attn_scores_query = attn_scores_query + pos_emb
attn_scores_query = attn_scores_query + attention_mask

attn_probs_query = F.softmax(attn_scores_query, dim=-1)
attn_probs_query = F.dropout(
attn_probs_query, p=0.1, training=self.training
)
query_output = torch.matmul(
attn_probs_query, V
) # (B, num_heads, S, head_dim)

# 合并多头并投影
content_output = self._reshape_from_heads(content_output)
query_output = self._reshape_from_heads(query_output)

content_output = self.o_proj(content_output)
query_output = self.o_proj(query_output)

return content_output, query_output

def _reshape_to_heads(self, x: torch.Tensor) -> torch.Tensor:
"""将 (B, S, H) 重塑为 (B, num_heads, S, head_dim)。"""
B, S, H = x.shape
return x.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)

def _reshape_from_heads(self, x: torch.Tensor) -> torch.Tensor:
"""将 (B, num_heads, S, head_dim) 重塑为 (B, S, H)。"""
return x.transpose(1, 2).contiguous().view(
x.shape[0], -1, self.hidden_dim
)

三、排列机制的工程实现

XLNet在工程层面并没有真正将输入序列按所有排列重新排序再输入模型——那样计算量不可接受。实际的实现方式是：保持输入序列的原始顺序不变，通过精心设计的注意力掩码（Attention Mask）来模拟不同的排列顺序。

对于排列 $\\mathbf{z} = [z_1, \\ldots, z_T]$，注意力掩码 $M_{ij}$ 定义为：

$$M_{ij} = \\begin{cases} 0 & \\text{if } \\text{pos}^{-1}{\\mathbf{z}}(j) \\leq \\text{pos}^{-1}{\\mathbf{z}}(i) \\text{ (可以关注)} \\ -\\infty & \\text{otherwise (不能关注)} \\end{cases}$$

其中 $\\text{pos}^{-1}_{\\mathbf{z}}(i)$ 是token $i$ 在排列 $\\mathbf{z}$ 中的位置。这只需要在每个训练步采样一个新的排列并构建对应的掩码，无需重组输入数据。

四、与BERT的对比实验

在相同的训练资源（8×V100，训练1M步）和类似的模型规模（~110M参数）下，XLNet-base与BERT-base在GLUE基准上的对比：

XLNet在需要长距离推理的任务（RACE阅读理解、SQuAD v2.0）上优势最为明显，这与排列语言模型能够"无死角的双向上下文"的理论优势一致。但也需要指出，XLNet的训练成本约为BERT的1.5倍（因为排列采样的额外开销和双流注意力导致的更高计算量）。

五、总结

XLNet通过排列语言模型以一种比BERT更为自然的方式捕获了双向上下文——它避免了[MASK] token引入的预训练-微调不一致，并通过排列机制消除了独立性假设。双流自注意力是PLM能够工作的关键工程创新，其中内容流编码完整的上下文，查询流编码"知道位置但不知道内容"的预测状态。在长距离推理任务上，XLNet相较BERT的优势体现了排列语言模型在上下文捕获完整性方面的理论优势。XLNet的主要代价在于训练效率——排列采样和双流注意力带来了约50%的额外计算开销。