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

1次阅读

没有评论

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

在传统的单头注意力机制中，模型通过单一的注意力头来计算输入序列中各个位置之间的关系。这种设计存在两个主要问题：

语义单一性 ：单头注意力只能捕捉到一种类型的语义关系，无法同时关注不同方面的特征。例如，在处理自然语言时，我们可能需要同时关注语法结构、指代关系和情感倾向等多个维度的信息。
信息瓶颈 ：对于长序列建模，单头注意力的计算能力有限，容易导致信息丢失或过拟合。特别是在处理复杂任务时，单一的注意力头难以充分捕获序列中的多样性和复杂性。

单头注意力的计算可以表示为：

$$
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V
$$

其中，$Q$、$K$、$V$ 分别表示查询（Query）、键（Key）和值（Value）矩阵，$d_k$ 是键的维度。

多头注意力通过将输入线性投影到多个子空间，并行计算多个注意力头，最后将结果拼接起来：

$$
\text{MultiHead}(Q, K, V) = \text{concat}(\text{head}_1, \text{head}_2, \ldots, \text{head}_h)W^O
$$

每个头的计算方式为：

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

其中，$W_i^Q$、$W_i^K$、$W_i^V$ 是投影矩阵，$W^O$ 是输出投影矩阵。

并行化计算 ：多头注意力可以并行计算多个头的注意力权重，充分利用现代 GPU 的并行计算能力。
子空间语义分化 ：不同的头可以关注不同的语义特征。例如，一个头可能关注语法结构，另一个头关注指代关系，第三个头关注情感倾向。这种分工合作使得模型能够更全面地理解输入序列。

以下是一个完整的 PyTorch 实现，包含维度切分和合并操作，以及使用爱因斯坦求和约定（einsum）优化矩阵运算：

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, einsum

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        assert (self.head_dim * num_heads == embed_dim), "Embedding dimension must be divisible by number of heads"

        self.qkv_proj = nn.Linear(embed_dim, 3 * embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)

    def split_heads(self, x):
        batch_size, seq_len, _ = x.shape
        return rearrange(x, "b s (h d) -> b h s d", h=self.num_heads)

    def concat_heads(self, x):
        return rearrange(x, "b h s d -> b s (h d)")

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape

        # Project Q, K, V
        qkv = self.qkv_proj(x)
        q, k, v = torch.chunk(qkv, 3, dim=-1)

        # Split into multiple heads
        q = self.split_heads(q)
        k = self.split_heads(k)
        v = self.split_heads(v)

        # Scaled dot-product attention
        scores = einsum(q, k, "b h i d, b h j d -> b h i j") / (self.head_dim ** 0.5)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))

        attn_weights = F.softmax(scores, dim=-1)
        output = einsum(attn_weights, v, "b h i j, b h j d -> b h i d")

        # Concatenate heads
        output = self.concat_heads(output)

        # Final linear projection
        output = self.out_proj(output)

        return output, attn_weights

为了监控 GPU 内存占用，可以在训练循环中添加以下代码：

def train_step(model, batch):
    inputs, labels = batch
    inputs = inputs.to(device)
    labels = labels.to(device)

    # Clear previous gradients
    optimizer.zero_grad()

    # Forward pass
    outputs, attn_weights = model(inputs)

    # Compute loss
    loss = criterion(outputs, labels)

    # Backward pass
    loss.backward()

    # Update weights
    optimizer.step()

    # Print GPU memory usage
    print(f"GPU memory allocated: {torch.cuda.memory_allocated() / 1024 ** 2:.2f} MB")

    return loss.item()

我们在 IMDb 电影评论数据集上进行了实验，比较单头和多头注意力在文本分类任务上的表现。实验结果如下：