自注意力机制入门指南：从数学原理到PyTorch实现

1次阅读

没有评论

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

在理解自注意力机制之前，我们需要先掌握向量点积（dot product）的几何意义。给定两个向量 $\mathbf{a}$ 和 $\mathbf{b}$，它们的点积定义为：

$$\mathbf{a} \cdot \mathbf{b} = \sum_{i=1}^n a_i b_i = |\mathbf{a}| |\mathbf{b}| \cos \theta$$

其中 $\theta$ 是两个向量的夹角。这个公式告诉我们：

点积可以衡量两个向量的相似度：当向量方向相同时点积最大，垂直时为 0，相反时为负值
在自注意力中，点积被用来计算查询向量（Query）和键向量（Key）的关联程度

输入表示 ：假设我们有输入序列的矩阵表示 $X \in \mathbb{R}^{n \times d}$，其中 $n$ 是序列长度，$d$ 是特征维度
线性变换 ：通过可学习的权重矩阵 $W^Q, W^K, W^V$ 计算 Query、Key、Value 矩阵：
$$Q = X W^Q, \quad K = X W^K, \quad V = X W^V$$
注意力分数 ：计算所有 Query-Key 对的点积并缩放（除以 $\sqrt{d_k}$，$d_k$ 是 Key 的维度）：
$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$
Softmax 归一化 ：对每行进行 softmax 操作，使得注意力权重和为 1，保证数值稳定性

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

        # 确保 embed_size 可以被 heads 整除
        assert self.head_dim * heads == embed_size, "Embed size needs to be divisible by 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 = 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])  # 高效矩阵乘法

        # 处理 mask（用于 decoder 的因果注意力）if mask is not None:
            energy = energy.masked_fill(mask == 0, float("-1e20"))

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

        # 应用注意力权重到 values
        out = torch.einsum("nhql,nlhd->nqhd", [attention, values])
        out = out.reshape(N, query_len, self.heads * self.head_dim)  # 合并多头输出

        return self.fc_out(out), attention

并行计算 ：矩阵乘法在现代 GPU 上可以高度并行化，这也是自注意力比 RNN 更快的原因之一
内存管理 ：注意力矩阵大小是 $O(n^2)$，长序列会消耗大量内存，建议使用内存高效的注意力实现

梯度消失 ：softmax 可能导致梯度消失，解决方案：
使用更稳定的 softmax 实现（如 PyTorch 的 log_softmax）
添加残差连接

注意力可视化 ：

import matplotlib.pyplot as plt

def plot_attention(attention, sentence):
    fig = plt.figure(figsize=(12, 12))
    ax = fig.add_subplot(111)
    cax = ax.matshow(attention.numpy(), cmap='bone')
    fig.colorbar(cax)

    # 设置 x 轴和 y 轴标签
    ax.set_xticks(range(len(sentence)))
    ax.set_yticks(range(len(sentence)))

    plt.show()