为什么在多头注意力中需要对每个 Head 进行降维

在 Transformer 模型中，多头注意力机制（Multi-Head Attention）是一个非常重要的组成部分。它通过并行地计算多个注意力头（Attention Head）来增强模型的表示能力。然而，为了控制计算复杂度和内存使用量，通常对每个注意力头进行降维。本文将详细分析这种设计背后的原因，并通过公式和代码展示多头注意力的实现过程。

多头注意力机制概述

多头注意力机制的核心思想是将输入序列拆分成多个独立的子空间，分别计算注意力，然后将各个头的结果拼接起来，以获得更丰富的上下文信息。假设原始输入的维度是 $d_{\text{model}}$ ，多头注意力机制会将其分成 $h$ 个子空间，每个子空间的维度为 $d_k = \frac{d_{\text{model}}}{h}$ ，其中 $h$ 表示头的数量。

多头注意力的公式如下：

\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W^O

其中每个注意力头 $i$ 的计算方式为：

\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)

其中， $W_i^Q$ 、 $W_i^K$ 和 $W_i^V$ 分别是每个头对应的降维矩阵。

为什么需要对每个 Head 进行降维

1. 控制计算复杂度

在计算注意力得分时，涉及矩阵乘法运算。假设原始输入维度是 $d_{\text{model}}$ ，如果不进行降维，则每个头都需要使用 $d_{\text{model}}$ 维度的全矩阵计算。这样做的计算量和内存开销都非常大，尤其是当 $d_{\text{model}}$ 较大且 $h$ 较多的情况下。通过将每个头的维度降低为 $d_k = \frac{d_{\text{model}}}{h}$ ，计算复杂度会显著降低：

\text{计算复杂度} \approx O\left(h \cdot d_k^2\right) = O\left(h \cdot \left(\frac{d_{\text{model}}}{h}\right)^2\right) = O\left(\frac{d_{\text{model}}^2}{h}\right)

从公式中可以看出，降维后的计算复杂度与 $d_{\text{model}}^2$ 成正比，但与 $h$ 成反比，从而有效降低了计算复杂度。

2. 保持多头注意力的维度一致性

在多头注意力中，各个头的输出最终会进行拼接和映射。如果每个头不进行降维，则拼接后的结果维度将是 $h \times d_{\text{model}}$ 。这样会导致拼接后的结果维度增长，影响后续的网络结构。而将每个头的维度降至 $\frac{d_{\text{model}}}{h}$ 后，拼接后的维度仍为 $d_{\text{model}}$ ，从而与输入维度保持一致，方便与后续层相连接。

3. 提供不同的子空间表示

降维后，每个注意力头都可以在不同的子空间中独立计算不同的注意力权重。这样一来，多头注意力可以在多个子空间中捕捉到不同的特征和关系，从而提高模型的表现力和泛化能力。

代码实现

以下代码展示了多头注意力机制中的降维过程，以及如何将各个头的结果拼接在一起。

python
import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model 必须是 num_heads 的整数倍"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        # 定义 Q, K, V 的线性变换矩阵
        self.q_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        
        # 定义输出的线性变换矩阵
        self.out_linear = nn.Linear(d_model, d_model)

    def forward(self, Q, K, V):
        # 获取批次大小
        batch_size = Q.size(0)

        # 线性变换并分割为多个头
        Q = self.q_linear(Q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.k_linear(K).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.v_linear(V).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)

        # 计算注意力得分并应用到 V
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32))
        attn_probs = torch.nn.functional.softmax(attn_scores, dim=-1)
        attn_output = torch.matmul(attn_probs, V)

        # 拼接所有头的结果
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)

        # 通过输出线性变换
        output = self.out_linear(attn_output)
        return output

# 初始化参数
d_model = 64
num_heads = 8
Q = torch.rand(32, 10, d_model)
K = torch.rand(32, 10, d_model)
V = torch.rand(32, 10, d_model)

# 实例化多头注意力模块
multi_head_attention = MultiHeadAttention(d_model, num_heads)
output = multi_head_attention(Q, K, V)
print(output.shape)  # 输出形状 (32, 10, 64)

代码详解

定义维度：通过 d_model 和 num_heads 计算每个头的降维维度 d_k。
线性变换和分割：对输入进行线性变换，并将结果 reshape 为 (batch_size, num_heads, seq_len, d_k)。
计算多头注意力：在每个降维的子空间内计算注意力，并将结果拼接成 (batch_size, seq_len, d_model) 维度。
输出线性变换：通过线性层进行转换，以便于和后续网络层衔接。

总结

在多头注意力机制中，对每个头进行降维能够有效地控制计算复杂度、保持输出维度一致性，并且允许模型在不同的子空间中学习到更丰富的特征。通过上述分析和代码示例，可以更好地理解降维在多头注意力中的重要作用，并在实际应用中高效地实现这一过程。

目录