多头注意力（Multi‑Head Attention）

1. 多头注意力（Multi‑Head Attention）原理

设输入序列表示为矩阵 $X\in {\mathbb{R}}^{B\times L\times d_{\text{model}}}$，其中

• $B$：批大小（batch size），
• $L$：序列长度（sequence length），
• $d_{\text{model}}$：模型隐层维度（model dimension）。

多头注意力基于对缩放点乘注意力的并行化扩展，引入了 $h$ 个注意力头（heads），每个头在不同子空间中学习不同的表示。

1.1 线性映射与切分

我们首先为每个头定义三组可学习权重：
$$W_i^Q\in {\mathbb{R}}^{d_{\text{model}}\times d_k},W_i^K\in {\mathbb{R}}^{d_{\text{model}}\times d_k},W_i^V\in {\mathbb{R}}^{d_{\text{model}}\times d_v},i=1,\dots ,h$$
其中

• $h$：头数（number of heads），
• $d_k$：每个头中 Query/Key 的维度（key/query dimension），
• $d_v$：每个头中 Value 的维度（value dimension），
• 通常 $d_{\text{model}}=h\times d_k$ 且取 $d_v=d_k$。

对输入 $X$ 进行投影，得到第 $i$ 个头的查询、键、值：
$$Q_i=X{W}_i^Q,K_i=X{W}_i^K,V_i=X{W}_i^V$$
其中

• $Q_i\in {\mathbb{R}}^{B\times L\times d_k}$，
• $K_i\in {\mathbb{R}}^{B\times L\times d_k}$，
• $V_i\in {\mathbb{R}}^{B\times L\times d_v}$。

1.2 缩放点乘注意力

对第 $i$ 个头，分别对所有位置做点积注意力：

1. 打分矩阵
   $${\mathrm{Score}}_i=Q_i{K}_i^{\top }\in {\mathbb{R}}^{B\times L\times L}$$
2. 缩放
   $${\widetilde{\mathrm{Score}}}_i=\frac{{\mathrm{Score}}_i}{\sqrt{d_k}}$$
3. Softmax 归一化
   $$A_i=\mathrm{softmax}\left({\widetilde{\mathrm{Score}}}_i\right)\in {\mathbb{R}}^{B\times L\times L}$$
4. 加权求和
   $${\mathrm{head}}_i=A_i{V}_i\in {\mathbb{R}}^{B\times L\times d_v}$$

1.3 拼接与线性变换

将所有头的输出在最后一维拼接，再做一次线性投影：
$$\mathrm{Concat}=[{\mathrm{head}}_1,\dots ,{\mathrm{head}}_h]\in {\mathbb{R}}^{B\times L\times (h{d}_v)}$$
定义输出权重矩阵
$$W^O\in {\mathbb{R}}^{(h{d}_v)\times d_{\text{model}}}$$
最终输出：
$$\mathrm{MultiHead}(X)=\mathrm{Concat}{W}^O\in {\mathbb{R}}^{B\times L\times d_{\text{model}}}$$

---

2. PyTorch 代码实现

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, h: int):
        """
        d_model: 模型维度 d_model
        h: 注意力头数 h
        """
        super().__init__()
        assert d_model % h == 0, "d_model 必须能被 h 整除"
        self.d_model = d_model      # d_model
        self.h = h                  # h
        self.d_k = d_model // h     # d_k = d_model / h
        self.d_v = self.d_k         # d_v 通常等于 d_k

        # 投影矩阵 W_i^Q, W_i^K, W_i^V，实际上合并为一个大矩阵后在 forward 再切分
        self.w_q = nn.Linear(d_model, d_model)  # 同时生成 h 个 Q 投影
        self.w_k = nn.Linear(d_model, d_model)  # 同时生成 h 个 K 投影
        self.w_v = nn.Linear(d_model, d_model)  # 同时生成 h 个 V 投影

        # 输出线性变换 W^O
        self.w_o = nn.Linear(d_model, d_model)

    def forward(self, X: torch.Tensor, mask: torch.Tensor = None):
        """
        X: 输入张量，形状 (B, L, d_model)
        mask: 可选掩码，形状 (B, 1, L, L) 或 (B, L, L)
        """
        B, L, _ = X.size()

        # 1. 线性映射到 Q, K, V，然后切分 h 头
        #    先得到 (B, L, h*d_k)，再 view/transpose 为 (B, h, L, d_k)
        Q = self.w_q(X).view(B, L, self.h, self.d_k).transpose(1, 2)
        K = self.w_k(X).view(B, L, self.h, self.d_k).transpose(1, 2)
        V = self.w_v(X).view(B, L, self.h, self.d_k).transpose(1, 2)
        # 此时 Q, K, V 形状均为 (B, h, L, d_k)

        # 2. 计算点积注意力
        #    scores = Q @ K^T  -> (B, h, L, L)
        scores = torch.matmul(Q, K.transpose(-2, -1))  
        #    缩放：除以 sqrt(d_k)
        scores = scores / math.sqrt(self.d_k)
        #    可选掩码：将被屏蔽位置设为 -inf 
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        #    Softmax 归一化 -> (B, h, L, L)
        A = F.softmax(scores, dim=-1)

        #    加权求和 -> head_i 形状 (B, h, L, d_k)
        heads = torch.matmul(A, V)

        # 3. 拼接 h 个头：transpose 回 (B, L, h, d_k) 再 reshape
        concat = heads.transpose(1, 2).contiguous().view(B, L, self.h * self.d_k)
        #    concat 形状 (B, L, h*d_k) == (B, L, d_model)

        # 4. 最后一次线性变换 W^O
        output = self.w_o(concat)  # -> (B, L, d_model)
        return output, A  # 返回输出及注意力权重 A