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

1次阅读

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

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

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

其中 $Q,K,V$ 分别表示查询、键和值矩阵。当模型在预测 $t$ 时刻的输出时，如果能够访问 $t+1$ 时刻的信息，会导致评估指标虚高，而实际部署时模型性能骤降。

RNN/LSTM/GRU：
优点：天然的单向结构避免信息泄漏
缺点：长距离依赖捕捉能力弱，并行计算效率低
标准 Transformer：
优点：强大的全局建模能力
缺点：需要额外约束才能保证因果性
Causal Transformer：
继承 Transformer 的并行计算优势
通过掩码实现严格时间因果性
实验显示在同等参数量下比 LSTM 提升 23% 的 perplexity

# 生成下三角布尔矩阵（True 表示需要屏蔽）def create_causal_mask(seq_len, device):
    return torch.triu(torch.ones(seq_len, seq_len, device=device), diagonal=1).bool()

class CausalAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_heads)

    def forward(self, x):
        seq_len = x.size(0)
        mask = create_causal_mask(seq_len, x.device)
        attn_output, _ = self.attn(x, x, x, attn_mask=mask)
        return attn_output

通过反向传播检查注意力权重梯度：

# 确保梯度不会流向未来位置
for t in range(1, seq_len):
    assert attn_weights[t, :t].requires_grad  # 过去位置应有梯度
    assert not attn_weights[t, t:].requires_grad  # 未来位置应无梯度

KV Cache：缓存先前时间步的键值对，避免重复计算

# 推理时维护 cache 的示例
cache = {'k': torch.empty(0), 'v': torch.empty(0)}
for step in range(max_len):
    k = compute_key(step_input)
    v = compute_value(step_input)
    cache['k'] = torch.cat([cache['k'], k], dim=1)
    cache['v'] = torch.cat([cache['v'], v], dim=1)