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

1次阅读

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

在自然语言处理（NLP）领域，循环神经网络（RNN）曾经是处理序列数据的标准方法。然而，RNN 存在明显的缺陷：

长距离依赖问题 ：RNN 难以捕捉序列中相距较远的词语之间的关系，信息在长距离传递过程中容易丢失或失真。
并行计算困难 ：RNN 的计算是逐步进行的，无法充分利用现代 GPU 的并行计算能力，导致训练速度慢。

Transformer 架构的提出彻底改变了这一局面。其核心创新在于自注意力机制（Self-Attention），它能够直接计算序列中任意两个位置之间的关系，无论它们相距多远。这种机制不仅解决了长距离依赖问题，还完美支持并行计算，大幅提升了模型训练效率。

自注意力机制的核心思想是通过计算查询（Query）、键（Key）和值（Value）矩阵来建立词语之间的关系。具体计算过程如下：

线性变换 ：将输入的嵌入向量通过三个不同的权重矩阵（$W_Q$, $W_K$, $W_V$）投影到查询、键和值空间。
$$ Q = XW_Q, \quad K = XW_K, \quad V = XW_V $$
注意力得分计算 ：通过点积计算查询与键之间的相似度得分，然后进行缩放（除以 $\sqrt{d_k}$）和 softmax 归一化。
$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$

多头注意力机制将上述过程扩展到多个子空间（头），每个头学习不同的注意力模式，最后将结果拼接起来：

分割头 ：将 Q、K、V 矩阵按头数分割成多个子矩阵。
并行计算 ：每个头独立计算注意力得分和加权和。
拼接结果 ：将所有头的输出拼接起来，通过线性变换得到最终结果。

这种设计让模型能够同时关注来自不同子空间的信息，显著提升了表达能力。

下面是一个完整的 PyTorch 实现，包含 mask 机制和梯度检查点优化：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, n_heads=8, dropout=0.1):
        """
        初始化多头注意力层
        :param d_model: 输入维度（默认 512）:param n_heads: 注意力头数量（默认 8）:param dropout: dropout 率（默认 0.1）"""
        super().__init__()
        assert d_model % n_heads == 0, "d_model 必须能被 n_heads 整除"

        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads  # 每个头的维度

        # 线性变换矩阵
        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)

        self.dropout = nn.Dropout(dropout)

    def forward(self, q, k, v, mask=None):
        """
        前向传播
        :param q: 查询矩阵 (batch_size, seq_len, d_model)
        :param k: 键矩阵 (batch_size, seq_len, d_model)
        :param v: 值矩阵 (batch_size, seq_len, d_model)
        :param mask: 掩码矩阵 (batch_size, 1, seq_len, seq_len)
        :return: 注意力输出 (batch_size, seq_len, d_model)
        """
        batch_size = q.size(0)

        # 线性变换并分割头
        q = self.w_q(q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        k = self.w_k(k).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        v = self.w_v(v).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)

        # 使用梯度检查点优化显存
        attn_output = checkpoint(self._scaled_dot_product_attention, q, k, v, mask)

        # 拼接所有头并通过线性层
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, -1, self.d_model)
        return self.w_o(attn_output)

    def _scaled_dot_product_attention(self, q, k, v, mask=None):
        """缩放点积注意力计算"""
        # 计算注意力分数
        scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32))

        # 应用 mask（如果有）if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # softmax 归一化
        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # 加权求和
        return torch.matmul(attn_weights, v)

在实际实现中，初学者常会遇到以下几个问题：