多头注意力机制(Multi-head Attention)在Transformer架构中的关键作用与实现解析

1次阅读

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

多头注意力机制 (Multi-head Attention) 最早由 Vaswani 等人在 2017 年的经典论文《Attention is All You Need》中提出，是 Transformer 架构的核心组件。在传统的 RNN 和 LSTM 模型中，处理长序列依赖关系存在梯度消失和顺序计算的瓶颈。多头注意力的出现，使得模型可以并行处理序列中任意位置的关系，极大提升了处理长文本、时间序列等数据的效率。

自注意力机制的核心思想是通过计算序列中每个元素与其他元素的关联程度，来动态调整每个元素的表示。具体来说，对于输入序列中的每个元素，模型会生成三个向量：

Query(Q)：用于 ” 询问 ” 与其他元素的相关性
Key(K)：用于 ” 回答 ” 与其他 Query 的匹配程度
Value(V)：携带实际的信息内容

多头注意力的创新之处在于将注意力机制并行化：

将原始的 Q、K、V 矩阵通过不同的线性变换投影到多个子空间
在每个子空间中独立计算注意力
将多个头的输出拼接后做最后一次线性变换

这种设计允许模型：

同时关注不同位置的不同关系模式
学习更丰富的表示能力
提高模型的并行计算效率

单头注意力只能学习一种关系模式
不同头可以关注不同方面的信息（如语法、语义、远距离依赖等）
类似于 CNN 中使用多个滤波器提取不同特征

以下是一个完整的 PyTorch 多头注意力实现：

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

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        assert embed_dim % num_heads == 0, "Embedding dimension must be divisible by number of heads"

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        # 线性变换层
        self.qkv_proj = nn.Linear(embed_dim, 3*embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x, mask=None):
        batch_size, seq_len, embed_dim = x.size()

        # 生成 Q,K,V [batch_size, seq_len, 3*embed_dim]
        qkv = self.qkv_proj(x)

        # 分割为多头 [batch_size, num_heads, seq_len, head_dim]
        qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]

        # 计算缩放点积注意力 [batch_size, num_heads, seq_len, seq_len]
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)

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

        # Softmax 归一化
        attn_weights = F.softmax(attn_scores, dim=-1)

        # 加权求和 [batch_size, num_heads, seq_len, head_dim]
        output = torch.matmul(attn_weights, v)

        # 拼接多头输出 [batch_size, seq_len, embed_dim]
        output = output.permute(0, 2, 1, 3).contiguous()
        output = output.reshape(batch_size, seq_len, embed_dim)

        # 最终线性变换
        output = self.out_proj(output)

        return output

多头注意力的计算复杂度主要由矩阵乘法决定：