深度学习进阶（二十）Transformer-XL

从Transformer的“短视”说起各位朋友，我们在之前的系列文章中，已经深入探讨了Transformer架构的方方面面。它凭借自注意力机制，成为了NLP领域的事实标准。但是，经典的Transformer有一个“天生”的缺陷——固定长度上下文。想象一下，你正在阅读一本小说。当你看完第100页时，前99页的内容其实已经“忘”得差不多了。Transformer就是这样，它在处理每个batch时，只能看到当前这个batch内的token，无法利用更早之前的信息。这导致它对长文本的理解非常有限，比如在预测一个段落结尾的词语时，它可能无法参考到文章开头埋下的伏笔。Transformer-XL（Extra Long）正是为了解决这个问题而生的。它的核心思想是：让模型在训练和推理时，能够跨越batch的边界，利用历史信息。这就像给模型装上了一块“外接硬盘”，让它能随时读取过去的记忆。### 核心机制一：Segment-Level Recurrence（片段级循环）这个机制是Transformer-XL的灵魂。在标准的Transformer中，每个segment（片段，即一个batch）是独立处理的。而在Transformer-XL中，我们引入了循环连接。具体做法是：在处理当前segment时，我们会把上一个segment的隐藏层状态（Hidden State）缓存下来，并作为计算当前segment注意力时的额外输入。这个缓存被称为Memory。让我们用代码来直观感受一下这个概念（伪代码，展示核心逻辑）：pythonimport torchimport torch.nn as nnclass TransformerXLLayer(nn.Module): def __init__(self, d_model, n_head): super().__init__() self.attn = nn.MultiheadAttention(d_model, n_head) self.linear = nn.Linear(d_model, d_model) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) def forward(self, x, memory): # x: 当前segment的输入, shape: [seq_len, batch, d_model] # memory: 上一个segment的隐藏状态, shape: [mem_len, batch, d_model] # 关键步骤：将memory和当前输入拼接起来，作为key和value key_value = torch.cat([memory, x], dim=0) # shape: [mem_len + seq_len, batch, d_model] # 自注意力：query来自当前x，key和value来自拼接后的序列 attn_output, _ = self.attn(x, key_value, key_value) # 残差连接 + 层归一化 out1 = self.norm1(x + attn_output) # 前馈网络 out2 = self.linear(out1) return self.norm2(out1 + out2)# 使用示例d_model = 128n_head = 4layer = TransformerXLLayer(d_model, n_head)# 假设当前segment长度5，memory长度10cur_seq = torch.randn(5, 2, d_model) # 当前segmentmemory = torch.randn(10, 2, d_model) # 历史memoryoutput = layer(cur_seq, memory)print(f"输出shape: {output.shape}") # 仍然是 [5, 2, 128]这个设计的巧妙之处在于：1. 梯度截断：虽然我们使用了memory，但梯度并不会通过memory反向传播到上一个segment。这保证了训练的稳定性，同时又能让模型“看到”历史信息。2. 推理加速：在推理阶段，我们可以逐步生成token，并一直缓存之前的隐藏状态，无需重新计算，效率极高。### 核心机制二：Relative Positional Encoding（相对位置编码）好了，有了循环机制，但还有一个问题需要解决。在标准Transformer中，位置编码是绝对的——每个位置有一个固定的向量。但在Transformer-XL中，由于我们拼接了memory和当前segment，如果仍然使用绝对位置编码，模型会混淆“当前位置”和“memory中的位置”。举个例子：假设memory长度是10，当前segment是第11-15个token。如果使用绝对位置编码，第11个token的位置编码是“11”，但模型并不知道这个“11”是相对于当前batch的起点还是整个序列的起点。解决方案：改用相对位置编码。我们不关心“token的绝对位置”，而关心“两个token之间的相对距离”。这样，无论memory有多长，模型都能正确理解注意力关系。具体实现上，Transformer-XL对注意力分数计算做了修改。原始的注意力分数是：score(q, k) = (q + u) · (k + v)其中u和v是绝对位置向量。而Transformer-XL将其改为：score(q_i, k_j) = (q_i + R_{i-j}) · (k_j) + (q_i + u) · (k_j)其中R_{i-j}是相对位置嵌入（一个可学习的参数，依赖于距离i-j），u和v现在变成了可学习的偏置向量。我们来看一下简化版的实现：pythonimport torchimport torch.nn as nnimport mathclass RelativePositionalAttention(nn.Module): def __init__(self, d_model, n_head, max_len=1000): super().__init__() self.d_model = d_model self.n_head = n_head self.head_dim = d_model // n_head # 可学习的相对位置嵌入表 self.pos_embedding = nn.Embedding(max_len, self.head_dim) # 可学习的偏置参数 self.u = nn.Parameter(torch.randn(n_head, self.head_dim)) self.v = nn.Parameter(torch.randn(n_head, self.head_dim)) def forward(self, q, k, v, distance): # q, k, v: [seq_len, batch, d_model] # distance: 相对距离矩阵, [seq_len, seq_len] batch_size, seq_len, _ = q.shape # 将q,k,v分割成多头 q = q.view(seq_len, batch_size, self.n_head, self.head_dim).transpose(0, 2) k = k.view(seq_len, batch_size, self.n_head, self.head_dim).transpose(0, 2) v = v.view(seq_len, batch_size, self.n_head, self.head_dim).transpose(0, 2) # 获取相对位置嵌入 rel_pos = self.pos_embedding(distance) # [seq_len, seq_len, head_dim] rel_pos = rel_pos.permute(2, 0, 1).unsqueeze(0) # [1, head_dim, seq_len, seq_len] # 计算注意力分数（简化版） # 第一部分：基于内容的注意力 content_score = torch.matmul(q, k.transpose(-2, -1)) # [batch, n_head, seq_len, seq_len] # 第二部分：基于相对位置的注意力 # 这里简化了计算，实际上需要对每个query进行不同的处理 # 我们使用一个简单的点积来演示 pos_score = torch.matmul(q, rel_pos) # [batch, n_head, seq_len, seq_len] # 加总两个分数 scores = (content_score + pos_score) / math.sqrt(self.head_dim) # 注意力权重 attn_weights = torch.softmax(scores, dim=-1) # 加权求和 output = torch.matmul(attn_weights, v) return output.transpose(0, 2).reshape(batch_size, seq_len, self.d_model)# 模拟使用d_model = 64n_head = 4attention = RelativePositionalAttention(d_model, n_head)q = torch.randn(10, 2, d_model)k = torch.randn(10, 2, d_model)v = torch.randn(10, 2, d_model)# 创建相对距离矩阵（这里简化为0-9）distance = torch.arange(10).unsqueeze(0).expand(10, 10) - torch.arange(10).unsqueeze(1)distance = torch.clamp(distance, 0, 9) # 限制范围output = attention(q, k, v, distance)print(f"输出shape: {output.shape}")### 训练与推理的完整流程在实际训练中，Transformer-XL是这样工作的：1. 准备数据：将长文本切成多个segment，每个segment长度固定（如128）。2. 前向传播：对于segment 2，我们使用segment 1的隐藏状态作为memory。3. 梯度更新：只计算当前segment的损失，梯度不回溯到memory。推理时，我们可以逐token生成。每当生成一个新token，我们把它加入当前segment，同时更新memory。### 性能提升与影响Transformer-XL在长文本任务上取得了巨大的成功，例如：- 语言建模：在WikiText-103等长文本数据集上刷新了SOTA。- 长距离依赖：能够捕捉超过1000个token的依赖关系。- 推理速度：比传统Transformer快很多，因为不需要重新计算历史状态。它的思想也直接启发了后续的XLNet、Compressive Transformer等模型，为处理超长序列提供了新的思路。### 总结Transformer-XL通过片段级循环机制和相对位置编码，成功突破了传统Transformer的固定长度限制。它让模型拥有了“记忆”，能够理解跨越很长的上下文的语义关联。虽然现在有了更新的架构（如Longformer、BigBird等），但Transformer-XL的核心思想——利用缓存和相对位置来扩展上下文——至今仍是处理长序列任务的重要基石。如果你正在处理类似文本摘要、长文档分类等问题，Transformer-XL绝对值得一试。