Transformer架构深度解析：为何必须使用多头注意力机制而非单头？

1次阅读

没有评论

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

在标准的单头注意力机制中，输入序列通过三个不同的线性变换得到查询 (Q)、键(K) 和值 (V) 矩阵。假设输入序列长度为 $n$，特征维度为 $d_{model}$，则：

$Q \in \mathbb{R}^{n \times d_k}$
$K \in \mathbb{R}^{n \times d_k}$
$V \in \mathbb{R}^{n \times d_v}$

其中 $d_k$ 通常等于 $d_v$，且 $d_k = d_{model}$ 在单头情况下。

注意力得分的计算公式为：
$$Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}})V$$

多头注意力将 $d_{model}$ 维的 Q、K、V 分别投影到 $h$ 个不同的子空间，每个子空间的维度为 $d_k = d_v = d_{model}/h$。最终将各头的输出拼接后再通过线性变换：

$$MultiHead(Q, K, V) = Concat(head_1,…,head_h)W^O$$
$$where\ head_i = Attention(QW_i^Q, KW_i^K, VW_i^V)$$

在长序列建模任务中，单头注意力的 Perplexity 指标明显劣于多头机制。实验表明，在 WikiText-103 数据集上，单头比 8 头模型的验证集困惑度高出 15-20%
单头注意力倾向于过度聚焦于局部区域，导致梯度消失。在图像描述生成任务中，单头模型的注意力分布熵值比多头低 30-40%

多头注意力的计算过程天然适合并行化，各注意力头可以独立计算。下图展示了维度变换流程：

[输入序列] -> [h 个线性投影] -> [h 个并行 attention 计算] -> [拼接] -> [输出投影]

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, h):
        super().__init__()
        assert d_model % h == 0, "d_model must be divisible by h"
        self.d_k = d_model // h
        self.h = h

        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, x, mask=None):
        batch_size = x.size(0)

        # 线性投影
        Q = self.W_q(x).view(batch_size, -1, self.h, self.d_k).transpose(1,2)
        K = self.W_k(x).view(batch_size, -1, self.h, self.d_k).transpose(1,2)
        V = self.W_v(x).view(batch_size, -1, self.h, self.d_k).transpose(1,2)

        # Scaled Dot-Product Attention
        scores = torch.matmul(Q, K.transpose(-2,-1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attn = torch.softmax(scores, dim=-1)

        # 拼接多头输出
        output = torch.matmul(attn, V)
        output = output.transpose(1,2).contiguous().view(batch_size, -1, self.h*self.d_k)

        return self.W_o(output)