Transformer 中的自注意力与多头机制

# Transformer 中的注意力机制：自注意力、交叉注意力与多头注意力

注意力机制让 AI 模型能够动态评估输入序列中各元素的重要性。对于每个 token，都会通过查询（Query）、键（Key）和值（Value）矩阵计算权重，从而生成上下文相关的表示。自注意力考察单一序列内部的联系，交叉注意力连接编码器与解码器，多头注意力则并行处理以捕捉数据不同维度。

自注意力原理

输入序列用矩阵 $X \in \mathbb{R}^{n \times d}$ 表示，其中 $n$ 为 token 数量，$d$ 为嵌入维度。从 $X$ 投影得到 $Q = XW^Q$、$K = XW^K$、$V = XW^V$，权重 $W^{Q,K,V} \in \mathbb{R}^{d \times d_k}$。

注意力分数矩阵 $A_{ij} = \frac{Q_i K_j^T}{\sqrt{d_k}}$ 衡量第 $j$ 个 token 对第 $i$ 个的关联度。用 $\sqrt{d_k}$ 缩放可稳定梯度，避免 softmax 中梯度消失。

行归一化矩阵 $N = \mathrm{softmax}(A)$ 形成概率分布。输出：$\mathrm{Attention}(Q, K, V) = N V$。

自注意力示例

考虑句子“Karina 要去商店”，分词为 ["karina", "是", "去", "商店"]，$n=5$。简化嵌入：

| Token | Embedding |

|--------|-----------|

| karina | [1, 0] |

| 是 | [0, 1] |

| 去 | [1, 1] |

| 该 | [1, 1] |

| 商店 | [0.5, 1] |

矩阵 $X_{5\times2} = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \\ 1 & 1 \\ 0.5 & 1 \end{bmatrix}$。为简化，$W^Q = W^K = W^V = I$，故 $Q=K=V=X$。

$QK^T = \begin{bmatrix} 1 & 0 & 1 & 1 & 0.5 \\ 0 & 1 & 1 & 1 & 1 \\ 1 & 1 & 2 & 2 & 1.5 \\ 1 & 1 & 2 & 2 & 1.5 \\ 0.5 & 1 & 1.5 & 1.5 & 1.25 \end{bmatrix}$，除以 $\sqrt{2} \approx 1.41$：

$A \approx \begin{bmatrix} 0.71 & 0 & 0.71 & 0.71 & 0.35 \\ 0 & 0.71 & 0.71 & 0.71 & 0.71 \\ 0.71 & 0.71 & 1.42 & 1.42 & 1.06 \\ 0.71 & 0.71 & 1.42 & 1.42 & 1.06 \\ 0.35 & 0.71 & 1.06 & 1.06 & 0.89 \end{bmatrix}$。

每行 softmax 后：$N \approx \begin{bmatrix} 0.25 & 0.15 & 0.25 & 0.25 & 0.10 \\ 0.17 & 0.23 & 0.23 & 0.23 & 0.23 \\ 0.15 & 0.15 & 0.30 & 0.30 & 0.10 \\ 0.15 & 0.15 & 0.30 & 0.30 & 0.10 \\ 0.14 & 0.23 & 0.28 & 0.28 & 0.07 \end{bmatrix}$。

输出 $N \times V = \begin{bmatrix} 0.73 & 0.68 \\ 0.57 & 0.87 \\ 0.68 & 0.81 \\ 0.68 & 0.81 \\ 0.63 & 0.84 \end{bmatrix}$。新嵌入融入了上下文：“是”的向量现在体现人类移动而非单纯时态。

import torch  
import torch.nn as nn  

class SelfAttention(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.embed_dim = embed_dim
        self.q_linear = nn.Linear(embed_dim, embed_dim)
        self.k_linear = nn.Linear(embed_dim, embed_dim)
        self.v_linear = nn.Linear(embed_dim, embed_dim)
        self.scale = embed_dim ** 0.5

    def forward(self, x):
        Q = self.q_linear(x)
        K = self.k_linear(x)
        V = self.v_linear(x)
        scores = torch.bmm(Q, K.transpose(1, 2)) / self.scale
        attn = torch.softmax(scores, dim=-1)
        out = torch.bmm(attn, V)
        return out

交叉注意力：连接序列

交叉注意力中，Query 来自解码器（$Q = X_{dec}W^Q$），Key 和 Value 来自编码器（$K = X_{enc}W^K$，$V = X_{enc}W^V$）。计算相同：$\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right) V$。

关键区别：

| 组件 | 自注意力 | 交叉注意力 |

|----------|-------------|-------------|

| Q | X | X_dec |

| K | X | X_enc |

| V | X | X_enc |

这让解码器聚焦编码器上下文的相关部分。

import torch  
import torch.nn as nn  

class CrossAttention(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.embed_dim = embed_dim
        self.q_linear = nn.Linear(embed_dim, embed_dim)
        self.k_linear = nn.Linear(embed_dim, embed_dim)
        self.v_linear = nn.Linear(embed_dim, embed_dim)
        self.scale = embed_dim ** 0.5

    def forward(self, x_dec, x_enc):
        Q = self.q_linear(x_dec)
        K = self.k_linear(x_enc)
        V = self.v_linear(x_enc)
        scores = torch.bmm(Q, K.transpose(1, 2)) / self.scale
        attn = torch.softmax(scores, dim=-1)
        out = torch.bmm(attn, V)
        return out

多头注意力：多视角并行

多头注意力运行 $h$ 个独立注意力计算，投影 $Q^{(i)} = X W_Q^{(i)}$、$K^{(i)} = X W_K^{(i)}$、$V^{(i)} = X W_V^{(i)}$。每个头：$\mathrm{head}_i = \mathrm{Attention}(Q^{(i)}, K^{(i)}, V^{(i)})$。

结果拼接并投影：$\mathrm{MultiHead}(X) = \mathrm{Concat}(head_1, \dots, head_h) W^O$。头维度 $d_k = d / h$，$d$ 为 embed_dim。

import torch
import torch.nn as nn

class MultiHeadSelfAttention(nn.Module):  
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        assert embed_dim % num_heads == 0
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.q = nn.Linear(embed_dim, embed_dim)
        self.k = nn.Linear(embed_dim, embed_dim)
        self.v = nn.Linear(embed_dim, embed_dim)
        self.out = nn.Linear(embed_dim, embed_dim)

    def forward(self, x):
        B, T, D = x.shape
        Q = self.q(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.k(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.v(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)
        attn = torch.softmax(scores, dim=-1)
        out = torch.matmul(attn, V).transpose(1, 2).contiguous().view(B, T, D)
        return self.out(out)

核心要点

• 自注意力为 token 添加上下文：每个输出向量依赖整个序列。
• $\sqrt{d_k}$ 缩放至关重要：稳定 softmax 和梯度。
• 交叉注意力用于 seq2seq：连接 Transformer 中的编码器与解码器。
• 多头捕捉多维度：语法、语义、长距离依赖。
• PyTorch 实现用 Linear 和 bmm：提升批处理效率。

— Editorial Team