Transformer模型实战：自注意力与多头自注意力机制的高效实现与优化

1次阅读

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

Transformer 模型之所以在 NLP 领域大放异彩，关键在于其自注意力机制（Self-Attention）。这种机制能够动态地为输入序列中的每个位置分配不同的权重，从而捕捉长距离依赖关系。简单来说，自注意力机制让模型能够 ” 看到 ” 输入序列中哪些部分在当前任务中更为重要。

自注意力的计算过程可以分为三个主要步骤：

计算查询（Query）、键（Key）和值（Value）矩阵
计算注意力分数（通过 Q 和 K 的点积）
对注意力分数进行 softmax 归一化，然后与 V 矩阵相乘得到最终输出

多头自注意力（Multi-Head Attention）则是将这个过程并行执行多次，然后将结果拼接起来。这样做的好处是模型可以同时关注不同位置的不同特征表示。

虽然自注意力机制功能强大，但它也带来了两个主要挑战：

计算复杂度高 ：传统自注意力机制的计算复杂度为 O(n²)，其中 n 是序列长度。对于长序列（如文档级别的文本），这会带来显著的计算负担。
内存占用大 ：需要存储中间注意力矩阵，这在处理长序列时可能超出 GPU 内存限制。

针对上述问题，业界提出了多种优化方案：

稀疏注意力 ：只计算部分位置的注意力分数，如局部窗口注意力或基于内容的稀疏注意力。
内存高效实现 ：使用分块计算或重新计算技术减少内存使用。
低秩近似 ：用低秩矩阵近似注意力矩阵。
混合精度训练 ：在保证模型精度的同时减少内存占用。

以下是一个基本的自注意力实现（Python + PyTorch）：

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

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

        # 确保 embed_size 可以被 heads 整除
        assert (self.head_dim * heads == embed_size), "Embedding size needs to be divisible by 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):
        # 获取 batch size
        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 = torch.softmax(energy / (self.embed_size ** (1 / 2)), dim=3)

        # 应用注意力权重到 values
        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

不同的实现方式在时间和空间复杂度上存在显著差异：