自注意力机制在长序列建模中的性能优化实战

1次阅读

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

传统自注意力机制的计算复杂度为 $O(n^2)$，其中 n 是序列长度。当处理长序列时（如 2048 tokens 以上），这会带来两个主要问题：

显存占用爆炸 ：存储 $n \times n$ 的注意力矩阵需要巨大内存。例如 n =4096 时，单精度浮点矩阵占用：$4096^2 \times 4 \text{bytes} \approx 67\text{MB}$，而实际模型通常有多个注意力头和多层，显存需求会成倍增加。
计算量剧增 ：在 A100 GPU 上，标准 self-attention 处理 4096 长度序列的 FLOPs 约为 $3 \times 4096^2 \times d_{model}$，当 $d_{model}=1024$ 时高达 51GFLOPs。

方案	计算复杂度	适用场景	PPL 损失	训练速度提升
稀疏注意力	O(n√n)	局部相关序列	<5%	3-5x
LSH 注意力 (Reformer)	O(n log n)	全局依赖序列	2-8%	4-8x
分块计算 (Longformer)	O(n)	超长文档	3-10%	5-10x

import torch
from torch.nn import functional as F

class LSHAttention(torch.nn.Module):
    def __init__(self, num_hashes=4, bucket_size=64):
        super().__init__()
        self.num_hashes = num_hashes
        self.bucket_size = bucket_size

    def hash_vectors(self, query, num_buckets):
        # 使用随机旋转投影进行哈希
        rotation_shape = (query.size(-1), self.num_hashes, num_buckets // 2)
        rotations = torch.randn(rotation_shape, device=query.device)

        # 计算哈希签名 [batch, seq_len, num_hashes]
        rotated_vecs = torch.einsum('btf,fhb->bth', query, rotations)
        rotated_vecs = torch.cat([rotated_vecs, -rotated_vecs], dim=-1)
        hashes = torch.argmax(rotated_vecs, dim=-1)
        return hashes

    def forward(self, query, key, value):
        batch_size, seq_len, dim = query.shape
        num_buckets = seq_len // self.bucket_size

        # 梯度检查点节省显存
        def compute_attention(query, key, value, masks):
            return F.scaled_dot_product_attention(query, key, value, attn_mask=masks)

        hashes = self.hash_vectors(query, num_buckets)  # [batch, seq_len, num_hashes]
        masks = (hashes.unsqueeze(-1) == hashes.unsqueeze(-2))  # 同桶标记

        return torch.utils.checkpoint.checkpoint(compute_attention, query, key, value, masks)

在 WikiText-103 上的测试数据（序列长度 4096）：