Transformer架构中的2.2.2多头注意力机制：原理剖析与高效实现

1次阅读

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

注意力机制在自然语言处理（NLP）中扮演着至关重要的角色，它能够帮助模型在处理序列数据时动态地关注输入的不同部分。传统的注意力机制在处理长序列时存在效率低下的问题，而多头注意力机制通过将注意力分散到多个“头”上，显著提升了模型的表达能力和计算效率。

多头注意力的核心思想是将输入数据分成多个子空间，每个子空间独立计算注意力权重，最后将所有子空间的结果合并。这种设计不仅增强了模型的并行处理能力，还能捕获输入数据中的多种依赖关系。

尽管多头注意力机制在理论上非常高效，但在实际实现中常常面临以下挑战：

计算复杂度高：多头注意力的计算复杂度与序列长度的平方成正比，长序列的处理会显著增加计算负担。
内存占用大：每个注意力头需要存储中间结果，导致内存占用急剧上升，尤其是在大规模模型训练中。
并行化难度大：传统的实现方式难以充分利用 GPU 的并行计算能力，导致训练和推理速度受限。

多头注意力机制的核心是计算查询（Q）、键（K）和值（V）矩阵的注意力分数。具体步骤如下：

将输入数据线性映射到多个子空间，生成多组 Q、K、V 矩阵。
对每组 Q 和 K 矩阵计算点积注意力分数，然后通过 Softmax 函数归一化。
将归一化的注意力分数与 V 矩阵相乘，得到每个头的输出。
将所有头的输出拼接在一起，并通过线性变换合并为最终输出。

以下是一个基于 PyTorch 的优化实现代码：

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

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

        # 线性映射层
        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)

    def forward(self, values, keys, query, mask=None):
        N = query.shape[0]
        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]

        # 分割输入到多个头
        values = values.reshape(N, value_len, self.heads, self.head_dim)
        keys = keys.reshape(N, key_len, self.heads, self.head_dim)
        queries = query.reshape(N, query_len, self.heads, self.head_dim)

        # 计算注意力分数
        energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])
        if mask is not None:
            energy = energy.masked_fill(mask == 0, float("-1e20"))
        attention = F.softmax(energy / (self.embed_size ** (1/2)), dim=3)

        # 计算输出
        out = torch.einsum("nhql,nlhd->nqhd", [attention, values])
        out = out.reshape(N, query_len, self.heads * self.head_dim)
        out = self.fc_out(out)
        return out