深入解析attention自注意力机制：从数学原理到PyTorch实现

1次阅读

没有评论

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

自注意力机制的核心是计算查询（Query）、键（Key）和值（Value）矩阵之间的关系。给定输入序列 $X \in \mathbb{R}^{n \times d_{model}}$，我们通过三个不同的线性变换得到 Q、K、V 矩阵：

$$
Q = XW^Q, \quad K = XW^K, \quad V = XW^V
$$

其中 $W^Q, W^K \in \mathbb{R}^{d_{model} \times d_k}$，$W^V \in \mathbb{R}^{d_{model} \times d_v}$ 是可学习参数。注意力权重通过缩放点积计算：

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

梯度消失问题：
RNN 在长序列上存在梯度消失 / 爆炸问题，反向传播时梯度需要连乘多个时间步
自注意力通过直接连接所有位置解决了长距离依赖问题
计算效率：
RNN 的 $O(n)$ 序列计算无法并行
自注意力的矩阵运算可完全并行，时间复杂度 $O(n^2d)$
尽管理论复杂度更高，但实际训练速度更快

import torch
import torch.nn as nn
from einops import rearrange, einsum

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, n_heads=8):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_k = d_model // n_heads
        self.n_heads = n_heads

        self.Wq = nn.Linear(d_model, d_model)
        self.Wk = nn.Linear(d_model, d_model)
        self.Wv = nn.Linear(d_model, d_model)
        self.out = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        # 1. 线性变换并分头
        q = rearrange(self.Wq(x), "b n (h d) -> b h n d", h=self.n_heads)
        k = rearrange(self.Wk(x), "b n (h d) -> b h n d", h=self.n_heads)
        v = rearrange(self.Wv(x), "b n (h d) -> b h n d", h=self.n_heads)

        # 2. 计算缩放点积注意力
        scores = einsum(q, k, "b h i d, b h j d -> b h i j") / (self.d_k ** 0.5)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        attn = torch.softmax(scores, dim=-1)

        # 3. 聚合 value 并合并头部
        out = einsum(attn, v, "b h i j, b h j d -> b h i d")
        out = rearrange(out, "b h n d -> b n (h d)")

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