深入解析Attention机制在反向传播中的参数更新原理及其影响

1次阅读

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

Attention 机制已成为现代深度学习模型的核心组件，尤其在自然语言处理和计算机视觉领域表现突出。其核心思想是通过动态计算权重来捕捉输入序列中不同部分的重要性，从而实现对关键信息的聚焦。在 Transformer 架构中，Attention 机制通过 Query、Key 和 Value 三个矩阵的交互来实现这一目标。

给定输入序列 $X$，Attention 机制的计算过程可表示为：

计算 Query、Key 和 Value 矩阵：
$$Q = XW_Q, K = XW_K, V = XW_V$$
计算注意力分数：
$$A = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)$$
计算输出：
$$\text{Attention}(Q,K,V) = AV$$

反向传播时，我们需要计算损失函数 $L$ 对各个参数的梯度：

对 Value 矩阵 $V$ 的梯度：
$$\frac{\partial L}{\partial V} = A^T \frac{\partial L}{\partial O}$$
对注意力权重 $A$ 的梯度：
$$\frac{\partial L}{\partial A} = \frac{\partial L}{\partial O} V^T$$
对 Query 和 Key 矩阵的梯度涉及更复杂的链式法则计算，需要考虑 softmax 和缩放操作的导数。

Attention 参数的更新会通过计算图影响模型其他部分的参数：

由于 Attention 层的输出会被传递到后续层，其参数更新会直接影响后续层的梯度计算。
在多头 Attention 中，不同头的参数更新会相互影响，需要谨慎调整学习率。
在深层 Transformer 中，Attention 参数的更新会通过残差连接影响整个模型的训练动态。

以下是 PyTorch 实现的自定义 Attention 层：

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

class SelfAttention(nn.Module):
    def __init__(self, embed_size, heads):
        super(SelfAttention, self).__init__()
        self.embed_size = embed_size
        self.heads = heads
        self.head_dim = embed_size // heads

        # 初始化 Q,K,V 矩阵
        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)

    def forward(self, values, keys, query, mask):
        N = query.shape[0]
        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]

        # 分割输入到多个头
        values = values.reshape(N, value_len, self.heads, self.head_dim)
        keys = keys.reshape(N, key_len, self.heads, self.head_dim)
        queries = query.reshape(N, query_len, self.heads, self.head_dim)

        # 计算注意力分数
        energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])
        if mask is not None:
            energy = energy.masked_fill(mask == 0, float("-1e20"))

        attention = torch.softmax(energy / (self.embed_size ** (1/2)), dim=3)

        # 计算输出
        out = torch.einsum("nhql,nlhd->nqhd", [attention, values])
        out = out.reshape(N, query_len, self.heads * self.head_dim)

        return self.fc_out(out)