8头自注意力机制深度解析：从原理到高效实现

1次阅读

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

自注意力机制（Self-Attention）是 Transformer 架构的核心组件，它能够捕捉输入序列中各个位置之间的依赖关系。然而，传统的单头注意力机制在处理复杂语义关系时存在局限性，因此引入了多头注意力机制（Multi-Head Attention）。本文将深入解析 8 头自注意力机制的工作原理，对比不同头数对模型性能的影响，并提供高效的 PyTorch 实现方案。

Transformer 模型由 Vaswani 等人在 2017 年提出，其核心是自注意力机制。自注意力机制通过计算输入序列中每个位置与其他位置的权重，来捕获序列内部的依赖关系。具体计算过程如下：

输入表示 ：给定输入序列 $X \in \mathbb{R}^{n \times d_{model}}$，其中 $n$ 是序列长度，$d_{model}$ 是嵌入维度。
线性变换 ：通过三个可学习的权重矩阵 $W^Q, W^K, W^V$，将输入 $X$ 分别映射为查询（Query）、键（Key）和值（Value）矩阵：
$Q = X W^Q$
$K = X W^K$
$V = X W^V$
注意力分数 ：计算查询与键的点积，并通过缩放因子 $\sqrt{d_k}$ 进行归一化，以防止梯度消失或爆炸：
$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V$

多头注意力机制通过将输入序列映射到多个子空间，并行计算多个注意力头，从而捕捉不同子空间的特征。其优势主要体现在以下几个方面：

并行捕捉不同子空间特征 ：每个头可以关注输入序列的不同方面，例如局部依赖、全局依赖或特定语义关系。
头数选择对模型容量和计算开销的影响 ：增加头数可以提高模型的表达能力，但也会增加计算和内存开销。因此，头数的选择需要在模型性能和计算效率之间进行权衡。

以下是一个完整的 8 头自注意力机制的 PyTorch 实现代码，包含可配置的头数参数和高效的矩阵分块计算：

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, num_heads=8, dropout=0.1):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"

        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)

        self.dropout = nn.Dropout(dropout)

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)

        # Linear transformations and split into heads
        Q = self.W_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)

        # Scaled Dot-Product Attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32))

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

        attention = F.softmax(scores, dim=-1)
        attention = self.dropout(attention)

        # Apply attention to values and concatenate heads
        output = torch.matmul(attention, V)
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)

        # Final linear transformation
        output = self.W_o(output)
        return output

多头注意力机制的内存占用与头数成正比。假设输入序列长度为 $n$，嵌入维度为 $d_{model}$，头数为 $h$，则每个头的维度为 $d_k = d_{model} / h$。内存占用主要包括：