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

1次阅读

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

自注意力机制是 Transformer 架构的核心组件，它让模型能够动态地关注输入序列中不同位置的信息。下面我们从原理到实现逐步拆解这个看似复杂的机制。

想象你在图书馆找资料：

Query：你的问题（” 我想找关于深度学习的书 ”）
Key：书籍的索引标签（书名 / 分类）
Value：书籍的实际内容

数学上这三个矩阵的关系是：

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

其中：
– Q ∈ R^{n×d_k}：查询矩阵
– K ∈ R^{m×d_k}：键矩阵
– V ∈ R^{m×d_v}：值矩阵
– d_k：键向量的维度

相似度计算：QK^T 得到注意力分数矩阵
缩放处理：除以√d_k 防止梯度消失
归一化：softmax 得到注意力权重
加权求和：权重与 V 相乘得到输出

具体推导示例：

\begin{aligned}
&\text{设} Q = \begin{bmatrix}1\\2\end{bmatrix}, K = \begin{bmatrix}0.5\\1.5\end{bmatrix}, V = \begin{bmatrix}10\\20\end{bmatrix}, d_k=1 \\
&\text{步骤 1：} QK^T = \begin{bmatrix}0.5&1.5\\1&3\end{bmatrix} \\
&\text{步骤 2：} \frac{QK^T}{\sqrt{1}} = \begin{bmatrix}0.5&1.5\\1&3\end{bmatrix} \\
&\text{步骤 3：} \text{softmax}(\begin{bmatrix}0.5&1.5\\1&3\end{bmatrix}) ≈ \begin{bmatrix}0.27&0.73\\0.12&0.88\end{bmatrix} \\
&\text{步骤 4：} \begin{bmatrix}0.27&0.73\\0.12&0.88\end{bmatrix} \begin{bmatrix}10\\20\end{bmatrix} = \begin{bmatrix}17.3\\18.8\end{bmatrix}
\end{aligned}

就像有多组专家同时从不同角度分析问题：

每组注意力头学习不同的关注模式
并行计算提高效率
最终拼接所有头的输出获得更丰富的表示

import torch
import torch.nn as nn
import torch.nn.functional as F

class SingleHeadAttention(nn.Module):
    def __init__(self, d_model, d_k):
        super().__init__()
        self.W_q = nn.Linear(d_model, d_k)  # [d_model, d_k]
        self.W_k = nn.Linear(d_model, d_k)  # [d_model, d_k]
        self.W_v = nn.Linear(d_model, d_k)  # [d_model, d_k]

    def forward(self, x, mask=None):
        # x: [batch, seq_len, d_model]
        Q = self.W_q(x)  # [batch, seq_len, d_k]
        K = self.W_k(x)  # [batch, seq_len, d_k]
        V = self.W_v(x)  # [batch, seq_len, d_k]

        scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(K.size(-1)))
        # scores: [batch, seq_len, seq_len]

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

        attn_weights = F.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, V)  # [batch, seq_len, d_k]

        return output

在解码器中需要防止看到未来信息：

# 创建下三角掩码矩阵
seq_len = 10
mask = torch.tril(torch.ones(seq_len, seq_len))  # [seq_len, seq_len]

# 应用示例
attn = SingleHeadAttention(d_model=512, d_k=64)
output = attn(x, mask=mask)  # x: [batch, 10, 512]

当 d_k 较大时，QK^T 的值可能非常大
导致 softmax 进入梯度饱和区
缩放因子√d_k 保持梯度稳定

对于长序列（如 >512 tokens）：

分块计算：将序列拆分为多个块处理
稀疏注意力：只计算局部位置的注意力
内存高效实现：使用 FlashAttention 等优化

调节 softmax 温度系数观察注意力分布变化：

def visualize_attention(temperature):
    # 随机生成注意力分数
    scores = torch.randn(5, 5)  # [seq_len, seq_len]

    # 应用不同温度系数的 softmax
    attn = F.softmax(scores/temperature, dim=-1)

    # 可视化
    plt.matshow(attn.detach().numpy())
    plt.title(f"Temperature={temperature}")
    plt.colorbar()
    plt.show()

# 实验对比
temperatures = [0.1, 1.0, 10.0]
for temp in temperatures:
    visualize_attention(temp)

预期现象：
– 低温（0.1）：注意力更集中（接近 one-hot）
– 常温（1.0）：合理分布
– 高温（10.0）：接近均匀分布

通过这个实验可以直观理解温度系数如何影响模型的关注模式。

掌握自注意力机制需要注意：