深入理解Attention机制：为什么自注意力评分越高效果越好？

1次阅读

没有评论

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

自注意力机制的核心是计算查询 (Query) 与键 (Key) 的相似度，然后对值 (Value) 进行加权求和。这个过程可以用以下公式表示：

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

其中 $d_k$ 是 Key 的维度，用于缩放点积结果。当 Query 和 Key 的点积结果越大，经过 softmax 后的权重就越大，表示这两个位置的相关性越强。

常见的注意力评分函数主要有两种：

点积注意力(Dot-Product)
优点：计算简单高效，适合 GPU 并行
缺点：当维度 $d_k$ 较大时，点积结果可能过大导致 softmax 梯度消失
加性注意力(Additive)
优点：通过神经网络学习更复杂的相似度关系
缺点：计算复杂度更高，需要额外的参数

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

        # 线性变换矩阵
        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, seq_len, embed_size)
        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)

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