深入解析attention反向传播机制：从数学原理到PyTorch实现

1次阅读

没有评论

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

在 Transformer 模型中，attention 机制是核心组件，但其反向传播过程存在显著的内存和计算瓶颈。传统实现中，attention 矩阵的计算复杂度为 O(N^2)，其中 N 是序列长度。这意味着：

内存消耗随序列长度平方级增长，尤其在长序列场景（如文档处理或语音识别）下，显存不足成为常见问题。
反向传播需要存储中间变量（如 attention 权重和 softmax 结果），进一步加剧内存压力。

举个例子，当序列长度达到 1024 时，单精度浮点数的 attention 矩阵将占用 4MB 显存（1024×1024×4 字节），而反向传播时可能需要存储多个这样的中间结果。

缩放点积注意力（scaled dot-product attention）的前向计算可表示为：

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

其中 $Q$, $K$, $V$ 分别对应 query、key、value 矩阵，$d_k$ 是 key 的维度。

假设损失函数 $L$ 对 attention 输出 $O$ 的梯度为 $\frac{\partial L}{\partial O}$，则各参数的梯度计算如下：

value 矩阵梯度 ：
$$\frac{\partial L}{\partial V} = P^T \frac{\partial L}{\partial O}$$
其中 $P = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)$
query 矩阵梯度 ：
$$\frac{\partial L}{\partial Q} = \frac{1}{\sqrt{d_k}} \left(\frac{\partial L}{\partial P} \circ (P – P \odot P) \right) K$$
这里 $\circ$ 表示逐元素乘法，$\odot$ 表示外积。
key 矩阵梯度 ：
$$\frac{\partial L}{\partial K} = \frac{1}{\sqrt{d_k}} \left(\frac{\partial L}{\partial P} \circ (P – P \odot P) \right)^T Q$$

关键点在于 softmax 梯度的计算：$\frac{\partial P}{\partial z} = P \circ (I – P^T)$，其中 $z = QK^T/\sqrt{d_k}$。

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q: torch.Tensor,  # [batch, heads, seq_len, dim]
    K: torch.Tensor,
    V: torch.Tensor,
    mask: torch.Tensor = None
) -> torch.Tensor:
    """基础 attention 实现，支持 mask 处理"""
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)

    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)

    attn_weights = F.softmax(scores, dim=-1)
    return torch.matmul(attn_weights, V)

from torch.utils.checkpoint import checkpoint

def memory_efficient_attention(Q, K, V, mask=None):
    """使用梯度检查点减少内存占用"""
    def _inner_forward(Q, K, V):
        return scaled_dot_product_attention(Q, K, V, mask)

    return checkpoint(_inner_forward, Q, K, V)

def test_gradients():
    batch, heads, seq_len, dim = 2, 4, 64, 32
    Q = torch.randn(batch, heads, seq_len, dim, requires_grad=True)
    K = torch.randn_like(Q, requires_grad=True)
    V = torch.randn_like(Q, requires_grad=True)

    # 前向计算
    output = scaled_dot_product_attention(Q, K, V)
    dummy_loss = output.sum()

    # 反向传播
    dummy_loss.backward()

    # 检查梯度是否存在
    assert Q.grad is not None
    assert K.grad is not None
    assert V.grad is not None
    print("梯度检查通过")

通过实验对比不同实现的性能表现（测试环境：NVIDIA V100, 序列长度 =512）：