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

XLNet排列语言模型的训练复现双向上下文的捕获方式与BERT的差异XLNetYang et al., NeurIPS 2019通过排列语言模型Permutation Language Modeling, PLM捕获双向上下文在多个NLP基准上超越了BERT。其核心创新在于不是像BERT那样通过[MASK] token破坏输入来引入双向性而是在自回归框架下通过对输入序列的分解顺序进行排列来实现对上下文的双向感知。本文从数学原理和代码实现两个层面复现PLM的训练过程重点分析排列机制如何通过双流自注意力Two-Stream Self-Attention实现看到上下文但不知道自己的位置这一关键特性。一、BERT的Masking与XLNet的排列两种双向性方案BERT通过Masked Language ModelingMLM来学习双向表示随机选择15%的token用[MASK]替换模型基于未Mask的上下文预测被Mask的token。这一方法有两个被广泛讨论的局限性预训练-微调不一致预训练时模型看到[MASK] token微调时没有[MASK]造成输入分布偏移。独立性假设BERT假设被Mask的token之间相互独立。对于New York is a [MASK]和[MASK] York is a city两个MaskBERT独立预测每个[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), devicex.device, dtypex.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, p0.1, trainingself.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, p0.1, trainingself.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基准上的对比任务BERT-baseXLNet-base差异MNLI-m84.686.82.2QNLI90.591.10.6SST-293.594.20.7RACE64.366.11.8SQuAD v2.076.380.23.9XLNet在需要长距离推理的任务RACE阅读理解、SQuAD v2.0上优势最为明显这与排列语言模型能够无死角的双向上下文的理论优势一致。但也需要指出XLNet的训练成本约为BERT的1.5倍因为排列采样的额外开销和双流注意力导致的更高计算量。五、总结XLNet通过排列语言模型以一种比BERT更为自然的方式捕获了双向上下文——它避免了[MASK] token引入的预训练-微调不一致并通过排列机制消除了独立性假设。双流自注意力是PLM能够工作的关键工程创新其中内容流编码完整的上下文查询流编码知道位置但不知道内容的预测状态。在长距离推理任务上XLNet相较BERT的优势体现了排列语言模型在上下文捕获完整性方面的理论优势。XLNet的主要代价在于训练效率——排列采样和双流注意力带来了约50%的额外计算开销。