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

1次阅读

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

自注意力机制（Self-Attention）是 Transformer 架构的核心组件，也是 BERT 等预训练模型成功的关键。与传统的 RNN 和 CNN 相比，自注意力机制能够直接建模输入序列中任意两个位置之间的关系，而不受限于局部窗口或顺序依赖。这种全局建模能力使得模型能够更好地捕捉长距离依赖关系，从而在自然语言处理（NLP）任务中取得了突破性进展。

自注意力机制的核心思想是通过计算输入序列中每个位置与其他位置的注意力权重，动态地聚合上下文信息。具体来说，给定输入序列 (X \in \mathbb{R}^{n \times d})，其中 (n) 是序列长度，(d) 是隐藏层维度，自注意力机制的计算过程如下：

首先，通过线性变换将输入 (X) 映射到查询（Query）、键（Key）和值（Value）三个空间：
[
Q = XW_Q, \quad K = XW_K, \quad V = XW_V
]
其中，(W_Q, W_K, W_V \in \mathbb{R}^{d \times d_k}) 是可学习的参数矩阵，(d_k) 是每个头的维度。
计算注意力分数：
[
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V
]
这里，(\frac{QK^T}{\sqrt{d_k}}) 是缩放点积注意力，除以 (\sqrt{d_k}) 是为了防止点积值过大导致梯度消失。
多头注意力（Multi-Head Attention）将上述过程重复 (h) 次，每次使用不同的参数矩阵，最后将结果拼接起来：
[
\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)W_O
]
其中，(\text{head}_i = \text{Attention}(QW_Q^i, KW_K^i, VW_V^i))，(W_O \in \mathbb{R}^{hd_v \times d}) 是输出投影矩阵。

以下是一个使用 PyTorch 实现多头注意力的代码示例，代码中包含了详细的注释：

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_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)

    def forward(self, Q, K, V, mask=None):
        batch_size = Q.size(0)

        # 线性变换并分头
        Q = self.W_Q(Q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_K(K).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_V(V).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)

        # 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32))
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attention = F.softmax(scores, dim=-1)

        # 加权求和并拼接
        context = torch.matmul(attention, V).transpose(1, 2).contiguous()
        context = context.view(batch_size, -1, self.d_model)
        output = self.W_O(context)

        return output, attention

在实际应用中，多头注意力机制的计算复杂度和内存占用是主要瓶颈。以下是几种常见的优化策略：