深入解析BERT多头注意力机制：从数学原理到PyTorch实现

1次阅读

没有评论

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

在自然语言处理领域，传统的 RNN（循环神经网络）在处理长序列时存在明显的局限性。最突出的问题是：

梯度消失 / 爆炸 ：随着序列长度的增加，RNN 难以有效捕捉远距离依赖关系
顺序计算 ：无法并行处理序列，训练速度受限于序列长度
信息瓶颈 ：最后一个隐状态需要压缩整个序列信息

而自注意力机制通过计算词与词之间的关联度，实现了：

直接建模任意距离的依赖关系
完全并行的序列计算
动态权重分配（不同位置关注不同重要性的上下文）

标准的缩放点积注意力公式为：

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

其中：
– $Q \in \mathbb{R}^{n\times d_k}$（查询矩阵）
– $K \in \mathbb{R}^{m\times d_k}$（键矩阵）
– $V \in \mathbb{R}^{m\times d_v}$（值矩阵）
– $\sqrt{d_k}$ 缩放因子防止内积过大导致 softmax 饱和

将 Q、K、V 通过不同的线性投影拆分成 $h$ 个头：

$$
\begin{aligned}
head_i &= Attention(QW_i^Q, KW_i^K, VW_i^V) \
MultiHead(Q,K,V) &= Concat(head_1,…,head_h)W^O
\end{aligned}
$$

每个头的维度通常为 $d_{model}/h$，这样拼接后能保持总维度不变。

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, num_heads=8):
        super().__init__()
        assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除"

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_head = d_model // num_heads

        # 定义 QKV 的线性变换层
        self.wq = nn.Linear(d_model, d_model)  # [d_model, d_model]
        self.wk = nn.Linear(d_model, d_model)
        self.wv = nn.Linear(d_model, d_model)

        # 最终输出的线性层
        self.wo = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        """
        输入 x: [batch_size, seq_len, d_model]
        输出: [batch_size, seq_len, d_model]
        """
        batch_size, seq_len, _ = x.shape

        # 1. 计算 QKV [batch_size, seq_len, d_model]
        Q = self.wq(x)
        K = self.wk(x)
        V = self.wv(x)

        # 2. 拆分为多头 [batch_size, num_heads, seq_len, d_head]
        Q = Q.view(batch_size, -1, self.num_heads, self.d_head).transpose(1, 2)
        K = K.view(batch_size, -1, self.num_heads, self.d_head).transpose(1, 2)
        V = V.view(batch_size, -1, self.num_heads, self.d_head).transpose(1, 2)

        # 3. 计算缩放点积注意力 [batch_size, num_heads, seq_len, seq_len]
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_head))

        # 应用 mask（如需要）if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, -1e9)

        # softmax 归一化
        attn_weights = F.softmax(attn_scores, dim=-1)

        # 4. 计算注意力输出 [batch_size, num_heads, seq_len, d_head]
        attn_output = torch.matmul(attn_weights, V)

        # 5. 合并多头 [batch_size, seq_len, d_model]
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, self.d_model)

        # 6. 最终线性变换
        output = self.wo(attn_output)
        return output

头数 $h$ 与显存占用的关系近似为：

$$
显存 \approx batch_size \times seq_len^2 \times h \times \frac{d_{model}}{h}
$$

实践中建议：