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

1次阅读

没有评论

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

Attention 机制的核心思想是让模型在处理序列数据时，能够动态地关注与当前任务最相关的部分。简单来说，它通过计算权重来决定不同位置输入的重要性。这种机制广泛应用于自然语言处理、计算机视觉等领域。

Query, Key, Value：这是 Attention 的三个基本要素。Query 代表当前需要关注的内容，Key 是待匹配的内容，Value 则是实际被加权求和的内容
权重计算 ：通过 Query 和 Key 的相似度计算得到权重，通常使用点积或缩放点积
加权求和 ：用计算出的权重对 Value 进行加权求和，得到最终的 Attention 输出

在反向传播过程中，Attention 的参数主要通过梯度下降法进行更新。具体来说，包括以下几个步骤：

计算损失函数对 Attention 输出的梯度
这个梯度会沿着计算图反向传播到权重计算部分
再进一步传播到 Query、Key、Value 的投影矩阵
最终根据链式法则更新这些投影矩阵的参数

数学上，假设我们有缩放点积 Attention，其计算过程可以表示为：

Attention(Q,K,V) = softmax(QK^T/√d_k)V

其中，Q、K、V 都是通过可学习的投影矩阵得到的。在反向传播时，我们需要计算损失函数对这些投影矩阵的梯度。

Attention 层的参数更新会通过两种方式影响模型其他部分：

直接影响 ：Attention 的输出会作为后续层的输入，因此其参数更新会改变后续层接收到的梯度
间接影响 ：由于 Attention 权重是动态计算的，这些权重的变化会影响模型关注的重点区域，从而改变整个模型的行为

特别是在多层 Transformer 结构中，这种影响会通过残差连接和层归一化进一步放大。

下面是一个简单的 Self-Attention 层的 PyTorch 实现，展示了反向传播的过程：

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)

        # 计算 Q,K,V
        values = self.values(values)
        keys = self.keys(keys)
        queries = self.queries(queries)

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

        # 计算 Attention 权重
        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)

        # 最后的线性变换
        out = self.fc_out(out)
        return out

在这个实现中，PyTorch 的自动微分机制会自动处理反向传播时的梯度计算。我们可以清楚地看到，梯度会沿着计算图从输出反向传播到各个投影矩阵。

在深层 Attention 网络中，梯度可能会变得非常小或非常大。解决方法包括：