从零理解AI自注意力机制：原理剖析与PyTorch实战

1次阅读

没有评论

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

在自然语言处理中，循环神经网络（RNN）和长短期记忆网络（LSTM）曾经是处理序列数据的标配。但随着序列长度的增加，它们暴露出两个致命缺陷：

长期依赖问题：相距较远的 token 难以建立直接联系，信息需要通过多个时间步逐步传递，容易丢失或失真
并行计算困难：必须按时间步顺序计算，无法充分利用 GPU 的并行能力

自注意力机制（Self-Attention）的提出完美解决了这些问题。它让序列中的每个元素都能直接 ” 看到 ” 其他所有元素，并通过可学习的权重决定关注哪些部分。这种特性使其成为 Transformer 架构的核心组件。

自注意力机制的核心是计算查询（Query）、键（Key）和值（Value）三个矩阵之间的关系。我们用 LaTeX 展示完整计算流程：

线性变换：将输入序列 $X \in \mathbb{R}^{n \times d_{model}}$ 转换为 Q /K/V
$$Q = XW^Q, \quad K = XW^K, \quad V = XW^V$$
注意力得分：计算 query 和 key 的点积并缩放
$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

其中 $\sqrt{d_k}$ 的缩放是为了防止点积结果过大导致 softmax 梯度消失。

下面我们实现一个带 mask 支持的自注意力层（建议使用 Python 3.8+ 和 PyTorch 1.10+）：

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange

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

        assert (self.head_dim * heads == embed_size), "Embedding size needs to be divisible by 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):
        # 获取 batch size
        N = query.shape[0]

        # 分割嵌入维度到多个头
        values = rearrange(values, "n (h d) -> h n d", h=self.heads)
        keys = rearrange(keys, "n (h d) -> h n d", h=self.heads)
        queries = rearrange(query, "n (h d) -> h n d", h=self.heads)

        # 计算注意力得分
        energy = torch.einsum("hqd,hkd->hqk", [queries, keys])

        # 缩放得分
        energy = energy / (self.embed_size ** (1/2))

        # 应用 mask（如处理 padding）if mask is not None:
            energy = energy.masked_fill(mask == 0, float("-1e20"))

        # 计算注意力权重
        attention = torch.softmax(energy, dim=2)

        # 应用注意力权重到 values
        out = torch.einsum("hql,hld->hqd", [attention, values])

        # 合并多头输出
        out = rearrange(out, "h n d -> n (h d)")
        out = self.fc_out(out)

        return out

关键点说明：

使用 einops 库的 rearrange 函数简化维度操作
masked_fill处理 padding 位置（用极大负值使 softmax 后权重接近 0）
通过 torch.einsum 实现高效的矩阵运算

理解模型关注哪些 token 非常重要。我们可以用 matplotlib 绘制注意力热力图：

import matplotlib.pyplot as plt
import seaborn as sns

# 假设我们有一个注意力矩阵 attention [n_heads, seq_len, seq_len]
def plot_attention(attention, sentence):
    plt.figure(figsize=(10, 8))
    ax = sns.heatmap(attention[0].cpu().detach().numpy(), 
                    xticklabels=sentence.split(),
                    yticklabels=sentence.split(),
                    cmap="YlGnBu",
                    annot=True)
    ax.set_title("Attention Weights")
    plt.show()

# 示例用法
sentence = "the cat sat on the mat"
plot_attention(attention_matrix, sentence)

在实践中容易遇到以下问题：