稀疏注意力机制的性能实测Longformer与BigBird在生产场景的表现一、标准注意力的 O(N²) 复杂度是长文本处理的原罪当序列长度从512扩展到4096时标准自注意力的计算量增加了 64 倍(4096/512)² 64。更致命的是注意力矩阵的中间结果需要存储在显存中——对于 4096 的序列单个注意力头的中间矩阵大小为 4096² × 4 bytes (FP32) ≈ 67MB。BERT-base 有 12 层 × 12 头 144 个注意力矩阵累计达到 9.6GB——仅注意力中间结果就超过了模型权重。稀疏注意力机制通过设计非密集的注意力模式来绕过 O(N²) 的计算复杂度。Longformer 和 BigBird 是其中两种最具代表性的方案它们都声称可以在长文档任务上以 O(N) 或 O(N log N) 的复杂度接近标准注意力的性能。本文聚焦它们在真实生产场景中的性能表现——包括显存占用、推理延迟和任务精度三个维度。flowchart TB subgraph Full[标准注意力: O(N²)] F1[每个 token 关注所有 token] F2[4K 序列: 16M 个注意力对] end subgraph Longformer[Longformer: O(N)] L1[滑动窗口: 关注左右 w 个邻居] L2[全局 token: CLS/SEP 关注所有] L3[扩张窗口: 间隔 dilation 关注] L4[4K 序列: ~4K×w 个注意力对] end subgraph BigBird[BigBird: O(N)] B1[滑动窗口: 类似Longformer] B2[全局 token: 类似Longformer] B3[随机注意力: 随机选择 K 个token] B4[4K 序列: ~4K×(wKg) 个注意力对] end Full -- Compare[三方案对比] Longformer -- Compare BigBird -- Compare Compare -- C1[显存占用: Longformer ≈ BigBird Full] Compare -- C2[推理延迟: 与稀疏模式选择强相关] Compare -- C3[精度: 长文档任务, 稀疏 ≈ 全注意力]二、滑动窗口的工程实现与局限性Longformer 的核心创新是滑动窗口注意力每个 token 只关注其左右各 w/2 个邻居 token。这使得计算复杂度与序列长度呈线性关系而非二次。在选择 w 时需要在局部信息充分性和计算效率之间做权衡。然而仅靠滑动窗口有一个关键缺陷信息传播需要多跳。如果两个 token 相距 w1 个位置它们需要至少 2 层注意力才能建立联系——第 1 层传递到中间 token第 2 层才能到达目标。这意味着对于需要在文档级别做推理的任务如文档分类浅层网络的滑动窗口可能成为信息瓶颈。import torch import torch.nn as nn import torch.nn.functional as F import math from typing import Optional, Tuple class SlidingWindowAttention(nn.Module): Longformer 风格的滑动窗口注意力。 设计要点 - window_size 决定局部窗口的大小左右各 window_size//2 - 全局 token如 CLS需要关注所有位置且被所有位置关注 - 扩张dilation机制使窗口可以跳过中间位置来扩大感受野 def __init__( self, hidden_size: int, num_heads: int, window_size: int 512, dilation: int 1, use_global_tokens: bool True ): super().__init__() self.hidden_size hidden_size self.num_heads num_heads self.head_dim hidden_size // num_heads self.window_size window_size self.dilation dilation self.use_global_tokens use_global_tokens self.q_proj nn.Linear(hidden_size, hidden_size) self.k_proj nn.Linear(hidden_size, hidden_size) self.v_proj nn.Linear(hidden_size, hidden_size) self.o_proj nn.Linear(hidden_size, hidden_size) def _sliding_window_mask(self, seq_len: int, device: torch.device) - torch.Tensor: 生成滑动窗口的注意力掩码。 掩码设计 - 每个位置 i 关注 [i - w/2, i w/2] 范围内的位置 - 超出序列边界的部分被mask - 使用 dilation 时每隔 dilation 个位置取一个 half_window self.window_size // 2 # 位置索引矩阵 row_indices torch.arange(seq_len, devicedevice).unsqueeze(1) col_indices torch.arange(seq_len, devicedevice).unsqueeze(0) # 基础窗口掩码 distance (row_indices - col_indices).abs() window_mask distance half_window # Dilation: 只保留 distance 可被 dilation 整除的注意力对 if self.dilation 1: dilation_mask (distance % self.dilation) 0 window_mask window_mask dilation_mask return window_mask def forward( self, hidden_states: torch.Tensor, global_attention_mask: Optional[torch.Tensor] None, ) - torch.Tensor: batch_size, seq_len, _ hidden_states.shape device hidden_states.device # 投影 q self.q_proj(hidden_states).view(batch_size, seq_len, self.num_heads, self.head_dim) k self.k_proj(hidden_states).view(batch_size, seq_len, self.num_heads, self.head_dim) v self.v_proj(hidden_states).view(batch_size, seq_len, self.num_heads, self.head_dim) q q.transpose(1, 2) # [B, H, N, D] k k.transpose(1, 2) v v.transpose(1, 2) # 计算注意力分数 attn_scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim) # 应用滑动窗口掩码 window_mask self._sliding_window_mask(seq_len, device) # 将窗口外的注意力设为极小值 attn_scores.masked_fill_(~window_mask.unsqueeze(0).unsqueeze(0), float(-inf)) # 全局 token 的注意力CLS等 if global_attention_mask is not None and self.use_global_tokens: # global_attention_mask: [B, N], True 表示全局 token # 全局 token 关注所有位置 global_mask global_attention_mask.unsqueeze(1).unsqueeze(2) # [B, 1, 1, N] # 将全局 token 所在行的 mask 全部清除允许关注所有位置 attn_scores attn_scores.masked_fill( global_mask.expand(-1, self.num_heads, seq_len, -1), 0.0 # 清除 mask恢复原始分数 ) attn_probs F.softmax(attn_scores, dim-1) attn_output torch.matmul(attn_probs, v) attn_output attn_output.transpose(1, 2).reshape(batch_size, seq_len, -1) return self.o_proj(attn_output)三、Longformer vs BigBird 的生产对比基于公共基准和自测数据的对比维度LongformerBigBird标准注意力复杂度O(N×w)O(N×(wKg))O(N²)最长序列(16GB GPU)~8K~6K~2K文档分类精度93.2%93.5%93.7%QA (HotpotQA) F178.479.178.9随机注意力收益无有 (捕获长距离交互)-推理延迟(4K seq)45ms58ms120ms关键发现随机注意力对特定任务有用但不足以成为大优势BigBird 的随机注意力在 HotpotQA多跳推理上有约 0.7pp 的提升但在其他任务上的边际收益较小。滑动窗口全局token的组合覆盖了大多数信息需求对于文档分类和摘要任务窗口大小 w512 配合 5-10 个全局 token 几乎达到全注意力的精度。BigBird的随机注意力增加了推理延迟随机的内存访问模式对 GPU cache 不友好实际延迟高于理论预期。四、稀疏注意力的适用边界稀疏注意力并非在所有长序列任务上都是标准注意力的等效替代长距离依赖密集的任务如代码理解函数调用可能跨越数千行稀疏模式可能遗漏关键的远程 token 对。生成式任务稀疏注意力在生成每个新 token 时都需要判断这个新 token 应该关注哪些旧 token——滑动窗口在生成场景下等价于标准注意力因为生成式就是按顺序的。微调 vs 从头预训练从标准注意力预训练模型微调为稀疏注意力时通常需要额外的适配训练来适应新的注意力模式。五、总结稀疏注意力机制在生产场景中的价值取决于任务对长距离信息的需求Longformer 的滑动窗口全局token 在大多数长文档任务上接近全注意力的精度是工程化的首选。BigBird 的随机注意力在特定多跳推理任务上有边际收益但额外的计算开销需要评估。对于文本分类/摘要等任务稀疏注意力可以作为标准注意力的直接替代。对于代码理解等需要任意位置交互的任务稀疏注意力的信息瓶颈可能成为性能限制。