BERT Transformer 多头注意力机制：从原理到新手实践指南

1次阅读

没有评论

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

在自然语言处理领域，处理长序列数据一直是个难题。传统 RNN（循环神经网络）虽然能够处理序列数据，但它有个致命缺陷：

梯度消失问题：随着序列长度增加，RNN 难以捕捉远距离依赖关系
顺序计算限制：必须按时间步逐步计算，无法并行处理
信息瓶颈：最后时刻的隐藏状态需要记住整个序列信息

这些问题在 2017 年被 Transformer 架构彻底改变。Transformer 引入了自注意力机制，允许模型直接计算序列中任意两个位置的关系，无论它们相距多远。

自注意力机制的核心是三个矩阵：Query(Q)、Key(K)和 Value(V)。给定输入序列 $X \in \mathbb{R}^{n \times d_{model}}$，计算过程如下：

线性变换得到 Q /K/ V 矩阵：
$$
Q = XW^Q, \quad K = XW^K, \quad V = XW^V
$$
计算注意力分数：
$$
\text{Attention}(Q,K,V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})V
$$

其中 $d_k$ 是 key 的维度，缩放因子 $\frac{1}{\sqrt{d_k}}$ 防止点积过大导致梯度消失。

多头注意力就是将这个过程并行执行多次：
$$
\text{MultiHead}(Q,K,V) = \text{Concat}(head_1,…,head_h)W^O
$$
$$
head_i = \text{Attention}(QW_i^Q,KW_i^K,VW_i^V)
$$

下面我们实现一个完整的多头注意力模块，包含 mask 机制和梯度检查点优化：

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

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

        # 初始化参数
        self.d_model = d_model          # 模型维度
        self.n_heads = n_heads          # 注意力头数量
        self.d_head = d_model // n_heads # 每个头的维度

        # 定义线性变换层
        self.w_q = nn.Linear(d_model, d_model)  # Query 变换
        self.w_k = nn.Linear(d_model, d_model)  # Key 变换
        self.w_v = nn.Linear(d_model, d_model)  # Value 变换
        self.w_o = nn.Linear(d_model, d_model)  # 输出变换

    def forward(self, x, mask=None):
        """
        输入:
            x: 输入张量 [batch_size, seq_len, d_model]
            mask: 掩码张量 [batch_size, seq_len, seq_len]
        输出:
            output: 多头注意力输出 [batch_size, seq_len, d_model]
            attn_weights: 注意力权重 [batch_size, n_heads, seq_len, seq_len]
        """
        batch_size, seq_len, _ = x.size()

        # 1. 线性变换得到 Q /K/V
        q = self.w_q(x)  # [batch_size, seq_len, d_model]
        k = self.w_k(x)  
        v = self.w_v(x)

        # 2. 分割多头
        q = q.view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2)

        # 3. 计算缩放点积注意力
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) / (self.d_head ** 0.5)

        # 4. 应用 mask（如需要）if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))

        # 5. 计算注意力权重
        attn_weights = F.softmax(attn_scores, dim=-1)

        # 6. 计算输出
        output = torch.matmul(attn_weights, v)

        # 7. 合并多头
        output = output.transpose(1, 2).contiguous()
        output = output.view(batch_size, seq_len, self.d_model)

        # 8. 最终线性变换
        output = self.w_o(output)

        return output, attn_weights

在实际实现中，新手常会遇到以下问题：