Causal Transformer 实战：解决长序列建模中的信息泄漏问题

1次阅读

没有评论

共计 2891 个字符，预计需要花费 8 分钟才能阅读完成。

传统 Transformer 的自注意力机制允许每个位置关注序列的所有位置（包括未来时间步），这在时间序列预测或文本生成任务中会导致信息泄漏。数学上，标准注意力分数计算为：

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

其中 $Q$, $K$, $V$ 分别表示查询、键和值矩阵。当模型在预测 $t$ 时刻的输出时，若能够访问 $t+1$ 时刻的键值信息，会导致训练和推理不一致，具体表现为：

训练阶段 ：模型利用未来信息 ” 偷看答案 ”，导致指标虚高
推理阶段 ：因无法获取真实未来信息，性能显著下降

LSTM/GRU
优势：天然的时间单向性，适合流式处理
劣势：
- 长程依赖建模能力弱（梯度消失问题）
- 难以并行化计算
Causal Transformer
优势：
- 通过掩码强制因果性，保留 Transformer 的并行计算能力
- 自注意力机制可显式建模任意距离的依赖关系
劣势：
- 需要显式处理 KV Cache（内存开销）
- 推理时仍需顺序解码

import torch

def generate_causal_mask(seq_len, device='cpu'):
    """生成下三角布尔矩阵（True 表示需要被掩盖）"""
    return torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool().to(device)

import math
import torch.nn as nn

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_k = d_model // n_heads
        self.n_heads = n_heads
        self.qkv_proj = nn.Linear(d_model, d_model * 3)
        self.out_proj = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        # 1. 投影得到 Q,K,V [batch, seq_len, d_model]
        qkv = self.qkv_proj(x)
        q, k, v = torch.chunk(qkv, 3, dim=-1)

        # 2. 分割多头 [batch, seq_len, n_heads, d_k]
        batch_size, seq_len, _ = q.shape
        q = q.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)

        # 3. 计算缩放点积注意力
        attn_scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k)

        # 4. 应用因果掩码（自动广播）if mask is not None:
            attn_scores = attn_scores.masked_fill(mask, float('-inf'))

        attn_weights = torch.softmax(attn_scores, dim=-1)
        output = attn_weights @ v

        # 5. 合并多头输出
        output = output.transpose(1, 2).contiguous()
        output = output.view(batch_size, seq_len, -1)
        return self.out_proj(output)

def test_gradient_flow():
    model = CausalSelfAttention(d_model=512, n_heads=8)
    x = torch.randn(2, 10, 512).requires_grad_(True)
    mask = generate_causal_mask(10)

    y = model(x, mask)
    loss = y.sum()
    loss.backward()

    # 检查梯度是否存在
    assert x.grad is not None  # 应通过
    print("梯度验证通过")

在自回归生成场景（如 GPT），可通过缓存历史时刻的 Key/Value 减少重复计算：

class GenerationCache:
    def __init__(self, max_length):
        self.k_cache = None
        self.v_cache = None
        self.max_len = max_length

    def update(self, new_k, new_v):
        if self.k_cache is None:
            self.k_cache = new_k
            self.v_cache = new_v
        else:
            self.k_cache = torch.cat([self.k_cache, new_k], dim=2)
            self.v_cache = torch.cat([self.v_cache, new_v], dim=2)

            # 截断保留最近 max_length 个 token
            if self.k_cache.size(2) > self.max_len:
                self.k_cache = self.k_cache[:, :, -self.max_len:]
                self.v_cache = self.v_cache[:, :, -self.max_len:]

对于超长序列（如 10k+ token），可采用：