深入理解Attention反向传播：从数学原理到PyTorch实现

1次阅读

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

标准的 Scaled Dot-Product Attention 计算流程如下：

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

其中：
– $Q \in \mathbb{R}^{n\times d_k}$ 是查询矩阵
– $K \in \mathbb{R}^{m\times d_k}$ 是键矩阵
– $V \in \mathbb{R}^{m\times d_v}$ 是值矩阵
– $d_k$ 是键向量的维度

设损失函数为 $L$，我们需要计算 $\frac{\partial L}{\partial Q}$, $\frac{\partial L}{\partial K}$ 和 $\frac{\partial L}{\partial V}$。

计算过程可以分解为：

计算 $\frac{\partial L}{\partial V}$
计算 $\frac{\partial L}{\partial \text{softmax}(S)}$，其中 $S = \frac{QK^T}{\sqrt{d_k}}$
计算 $\frac{\partial L}{\partial S}$
计算 $\frac{\partial L}{\partial Q}$ 和 $\frac{\partial L}{\partial K}$

$$
\frac{\partial L}{\partial V} = \text{softmax}(S)^T \frac{\partial L}{\partial \text{Output}}
$$

设 $P = \text{softmax}(S)$，则：

$$
\frac{\partial L}{\partial S} = P \circ (\frac{\partial L}{\partial P} – \sum_{j}(\frac{\partial L}{\partial P_{ij}} \circ P_{ij})))
$$

其中 $\circ$ 表示逐元素乘法。

$$
\frac{\partial L}{\partial Q} = \frac{1}{\sqrt{d_k}} \frac{\partial L}{\partial S} K
$$

$$
\frac{\partial L}{\partial K} = \frac{1}{\sqrt{d_k}} Q^T \frac{\partial L}{\partial S}
$$

import torch
import torch.nn as nn
import torch.nn.functional as F

class CustomAttention(nn.Module):
    def __init__(self, d_k):
        super().__init__()
        self.d_k = d_k

    def forward(self, Q, K, V):
        """
        Q: (batch_size, n, d_k)
        K: (batch_size, m, d_k)
        V: (batch_size, m, d_v)
        """
        # 计算缩放点积注意力
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
        attn_weights = F.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, V)

        # 保存中间结果用于反向传播
        self.save_for_backward(Q, K, V, attn_weights, scores)

        return output

    def backward(self, grad_output):
        """grad_output: (batch_size, n, d_v)"""
        Q, K, V, attn_weights, scores = self.saved_tensors

        # 计算 V 的梯度
        grad_V = torch.matmul(attn_weights.transpose(-2, -1), grad_output)

        # 计算 softmax 部分的梯度
        grad_attn = torch.matmul(grad_output, V.transpose(-2, -1))
        grad_scores = attn_weights * (grad_attn - torch.sum(attn_weights * grad_attn, dim=-1, keepdim=True))

        # 计算 Q 和 K 的梯度
        grad_Q = torch.matmul(grad_scores, K) / (self.d_k ** 0.5)
        grad_K = torch.matmul(Q.transpose(-2, -1), grad_scores) / (self.d_k ** 0.5)

        return grad_Q, grad_K, grad_V

# 创建测试数据
d_k = 64
d_v = 128
batch_size = 2
n = 10
m = 15

Q = torch.randn(batch_size, n, d_k, requires_grad=True)
K = torch.randn(batch_size, m, d_k, requires_grad=True)
V = torch.randn(batch_size, m, d_v, requires_grad=True)

# 前向传播
attention = CustomAttention(d_k)
output = attention(Q, K, V)

# 随机生成梯度
fake_loss = output.sum()
fake_loss.backward()

# 使用 hook 验证梯度
Q.register_hook(lambda grad: print(f"Q grad norm: {grad.norm()}"))
K.register_hook(lambda grad: print(f"K grad norm: {grad.norm()}"))
V.register_hook(lambda grad: print(f"V grad norm: {grad.norm()}"))

当 $d_k$ 较大时，点积结果可能过大导致 softmax 饱和
解决方法：确保使用缩放因子 $\sqrt{d_k}$
初始化时应注意 Q、K 矩阵的方差

使用 fp16 时 softmax 容易溢出
建议在计算 softmax 前转换为 fp32
梯度缩放策略需要调整

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter()

def log_gradients(name, param, step):
    if param.grad is not None:
        writer.add_scalar(f"gradients/{name}", param.grad.norm(), step)

# 在训练循环中调用
for name, param in model.named_parameters():
    log_gradients(name, param, global_step)

当 Key 和 Query 维度不同时，梯度计算会如何变化？
如何修改反向传播实现来支持因果掩码？
多头注意力机制的反向传播有何不同？

通过本文的推导和实现，我们详细分析了 Attention 机制的反向传播过程。理解这些基础原理对于调试 Transformer 模型和实现自定义 Attention 变体非常重要。建议读者尝试在 PyTorch 中实现更复杂的 Attention 变体，如局部 Attention 或稀疏 Attention，以加深理解。

正文完