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

1次阅读

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

自注意力机制是 Transformer 架构的核心组件，它通过计算输入序列中每个位置与其他位置的关系权重，动态地聚合上下文信息。与传统 RNN 相比，自注意力机制具有两个显著优势：

能够直接建模任意距离的依赖关系，不受序列长度的限制
并行计算所有位置的关联，极大提升了训练效率

自注意力的核心公式为：

Attention(Q,K,V) = softmax(QK^T/√d_k)V

其中 Q(Query)、K(Key)、V(Value)都是输入序列的线性变换，d_k 是 Key 的维度。

尽管自注意力机制功能强大，但在实际应用中面临几个关键挑战：

计算复杂度 ：标准自注意力的 QK^T 矩阵乘法具有 O(n^2) 的时间和空间复杂度，对于长序列 (如 4096 tokens) 会变得极其昂贵
内存瓶颈：存储注意力矩阵需要大量显存，当 batch size 较大时容易导致 OOM
硬件利用率：原生实现可能无法充分利用现代 GPU 的并行计算能力

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

局部窗口注意力：限制每个 token 只关注其邻近的 w 个 token，复杂度降为 O(n×w)
轴向注意力：将 2D 注意力分解为行和列两个 1D 注意力操作
随机注意力：为每个 token 随机选择 k 个其他 token 计算注意力

分块计算：将大矩阵运算分解为小块，减少峰值内存使用
梯度检查点：在反向传播时重新计算部分中间结果，牺牲时间换空间
混合精度训练：使用 FP16 存储注意力矩阵，配合 Loss Scaling

以下是基于 PyTorch 的高效自注意力层实现：

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

class EfficientAttention(nn.Module):
    """内存优化的自注意力层"""
    def __init__(self, embed_dim, num_heads, chunk_size=64):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.chunk_size = chunk_size

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

    def forward(self, x):
        """
        输入: 
            x: [batch_size, seq_len, embed_dim]
        输出:
            [batch_size, seq_len, embed_dim]
        """
        batch_size, seq_len, _ = x.shape
        qkv = self.qkv_proj(x).chunk(3, dim=-1)
        q, k, v = [t.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) 
                  for t in qkv]

        # 分块计算注意力
        out = torch.zeros_like(q)
        for i in range(0, seq_len, self.chunk_size):
            end = min(i + self.chunk_size, seq_len)
            q_chunk = q[:, :, i:end, :]

            # 计算注意力分数
            attn = torch.matmul(q_chunk, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
            attn = F.softmax(attn, dim=-1)

            # 聚合 value
            out[:, :, i:end, :] = torch.matmul(attn, v)

        out = out.transpose(1, 2).reshape(batch_size, seq_len, -1)
        return self.out_proj(out)

我们在 NVIDIA V100 上测试了不同实现的性能(序列长度 1024)：