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

1次阅读

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

多头注意力机制（Multi-Head Attention）是 Transformer 架构中的核心组件，广泛应用于自然语言处理、计算机视觉等领域。本文将从数学原理出发，详细解析多头注意力的实现细节，并提供优化后的 PyTorch 代码实现。

注意力机制的核心思想是通过计算查询（Query）、键（Key）和值（Value）之间的关系，动态地为每个查询分配不同的权重。公式如下：

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V

其中，Q、K、V分别是查询、键和值矩阵，d_k是键的维度。

多头注意力通过将 Q、K、V 分别投影到 h 个不同的子空间（称为“头”），并行计算注意力，然后将结果拼接起来。这样可以捕捉输入数据的不同特征。公式如下：

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) W^O

其中，每个头的计算方式为：

head_i = Attention(QW_i^Q, KW_i^K, VW_i^V)

W_i^Q、W_i^K、W_i^V和 W^O 是可学习的参数矩阵。

单头注意力：计算简单，内存占用低，但捕捉特征能力有限。
多头注意力：通过并行计算多个注意力头，可以捕捉输入数据的不同特征，但计算复杂度和内存占用较高。

以下是多头注意力的完整 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__()
        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):
        # Q, K, V 的形状: (batch_size, seq_len, d_model)
        batch_size = Q.size(0)

        # 线性变换并分头
        Q = self.W_q(Q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)  # (batch_size, num_heads, seq_len, d_k)
        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))  # (batch_size, num_heads, seq_len, seq_len)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

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

        # 加权求和
        output = torch.matmul(attn_weights, V)  # (batch_size, num_heads, seq_len, d_k)

        # 拼接多头结果
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)  # (batch_size, seq_len, d_model)

        # 线性变换
        output = self.W_o(output)

        return output, attn_weights

多头注意力的计算复杂度为 O(n^2 * d)，其中n 是序列长度，d是模型维度。内存占用主要来自注意力分数矩阵，大小为(batch_size, num_heads, seq_len, seq_len)。