Transformer 为什么使用多头注意力，而不是单头？

在现代深度学习中，Transformer 模型的多头注意力（Multi-Head Attention）机制被证明是自然语言处理和其他领域中极其强大的工具。一个常见的问题是：为什么 Transformer 使用多头注意力，而不是简单地使用一个头的注意力？ 本文将从公式推导和代码实现的角度进行详细且专业的讲解。

1. 多头注意力的直觉

在单头注意力中，我们可以通过自注意力机制计算每个输入元素之间的关系，从而捕捉到全局的依赖关系。但是，单头注意力存在局限性，它仅能够在单个表示空间中进行线性变换，从而可能无法充分表示输入数据的不同特征关系。

多头注意力通过在不同的子空间中进行注意力计算，从而获得多个表示。这样可以捕捉到输入的不同方面的特征，使得模型的表示能力大大增强。换句话说，多头注意力提供了一种并行计算的方式，通过在多个空间中独立捕捉不同的上下文信息，从而使得 Transformer 的表达能力更加丰富。

2. 多头注意力的公式推导

首先，我们回顾一下单头注意力的计算过程。假设输入矩阵为 $X \in \mathbb{R}^{n \times d}$ ，其中 $n$ 是序列长度， $d$ 是特征维度。注意力计算的步骤如下：

通过线性变换得到 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}$ 是可学习的参数矩阵。
计算注意力权重：
$A = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)$
计算输出：
$H = AV$

在多头注意力中，我们将输入 $X$ 变换为多个不同的 Query, Key, Value，分别进行注意力计算。假设有 $h$ 个头，每个头的维度为 $d_k$ ，则计算过程如下：

对每个头 $i$ ，分别计算 Query, Key, Value：
$Q_i = XW_i^Q, \quad K_i = XW_i^K, \quad V_i = XW_i^V$
对每个头 $i$ ，计算注意力权重和输出：
$A_i = \text{softmax}\left(\frac{Q_iK_i^T}{\sqrt{d_k}}\right), \quad H_i = A_iV_i$
将所有头的输出拼接在一起，并通过线性变换得到最终的输出：
$H_{\text{multi}} = [H_1, H_2, \ldots, H_h]W^O$
其中， $W^O \in \mathbb{R}^{hd_k \times d}$ 是可学习的参数矩阵。

通过上述步骤，多头注意力能够在多个子空间中独立捕捉输入的不同特征关系，使得模型能够更好地理解复杂的上下文依赖。

3. 代码实现

以下是多头注意力的代码实现，基于 PyTorch：

python
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_k = d_model // num_heads
        self.num_heads = 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, x):
        batch_size = x.size(0)
        
        # 线性变换并分割成多头
        q = self.w_q(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        k = self.w_k(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        v = self.w_v(x).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))
        attn = F.softmax(scores, dim=-1)
        
        # 计算注意力输出
        h = torch.matmul(attn, v)
        h = h.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k)
        
        # 最后线性变换
        output = self.w_o(h)
        return output

4. 多头注意力的优势

捕捉不同特征：多头注意力可以在不同的子空间中进行注意力计算，从而能够捕捉输入序列的不同特征关系。
提升模型的稳定性：单头注意力可能会陷入某些特定的注意力模式，而多头注意力通过多个头的并行计算，可以提高模型的稳定性，减少对某些特征的过度依赖。
增强表达能力：通过在多个子空间中计算注意力，多头注意力能够增强模型的表达能力，使得模型在面对复杂任务时具有更好的表现。

5. 总结

多头注意力机制是 Transformer 中的关键组件，其通过在多个子空间中并行地计算注意力，提升了模型的特征提取能力和表达能力。这使得 Transformer 在自然语言处理等任务中取得了非常优异的效果。

使用多头注意力的主要原因在于它能够捕捉输入数据的不同方面，从而提供更加全面的特征表示。如果仅使用一个头，模型将难以充分表示复杂的上下文依赖关系，进而影响整体性能。通过本文的公式推导和代码实现，希望能够帮助读者更好地理解多头注意力的工作原理及其优势。