最近在跑一个长序列模型时显存突然爆掉排查下来发现瓶颈就卡在某一层注意力上batch size 稍微调大一点N×N 的注意力矩阵直接让显卡告警。后来换成了 CVPR 2025 上比较受关注的 Focused Linear Attention才把这个问题理顺。这篇文章就围绕这个注意力模块展开讲讲它怎么替换传统 Softmax 注意力、替换时有哪些必须注意的细节以及我实测下来的效果和踩过的坑。这套内容适合正在做长文本、视频理解、高分辨率图像或多模态模型的开发者参考。无论你是在训练阶段被显存卡住还是想把推理速度再往上提一档都可以看完下面的拆解后直接在自己的模型里试一把。1. 我先算了一笔账Softmax 注意力在长序列上到底有多“贵”很多人一开始接触注意力机制时看到的核心公式就是Attention(Q,K,V)softmax(QK^T/√d)·V。这个公式很优雅但它背后藏着一个工程噩梦QK^T会生成一个和序列长度 N 的平方成正比的矩阵。我见过不少同学一开始对此没什么概念直到 512、1024 的 token 长度跑出来没问题换到 4096 就连续 OOM。这时候再回头算账才意识到Softmax attention 的时间和显存开销都是 O(N²) 级别长序列场景下这就是硬伤。1.1 复杂度不是玄学从一次矩阵乘法说起先拆一下计算流程。假设序列长度是 N特征维度是 d单头场景下计算QK^T得到N×N矩阵计算量约N²d对N×N矩阵做 Softmax计算量约N²将 Softmax 后的矩阵与 V 相乘计算量约N²d。所以整体复杂度就是 O(N²d)其中 N 是序列长度d 是特征维度。问题在于当 N 增大时N² 的增长是压倒性的。从 1024 到 4096序列长度只翻了 4 倍但注意力相关计算量翻了 16 倍。这还没算多头带来的叠加。打个比方Softmax 注意力就像一场全员聚会每个人都必须和全场所有人握手致意。人少时没问题人一旦多起来这个社交成本就会爆炸。你没法通过简单的“开会时长缩短”来解决握手次数本身。1.2 显存不只是 O(N²)训练时多头的乘积推理时我们也许能忍住不存完整注意力矩阵但训练时反向传播需要用到中间梯度很多实现会保存 Softmax 后的注意力分数也就是那个N×N矩阵。有时甚至要保存两份。多头注意力的每个头的矩阵又是独立的头数 h 一多显存开销直接从 O(N²) 变成 O(hN²)。拿一个小实验举例输入序列长度 2048维度 76812 个头float16 精度光一个注意力分数矩阵就是12×2048×2048×2 bytes约等于 96MB。看上去不大但模型有几十层每一层都有 QKV 投影和前馈网络梯度也会占显存。累加起来注意力矩阵这部分就很容易吃掉半张 24GB 显卡。这也是为什么很多人即便用了梯度检查点、混合精度长序列照样 OOM。因为根本问题没解决你的算法复杂度仍然是 O(N²)只是通过各种手段把常数项往下降。1.3 已有的优化为什么救不了根本问题有人会说“我用 FlashAttention 不就行了” FlashAttention 确实厉害它通过分块计算和 IO 感知优化把显存占用从 O(N²) 降到 O(N)速度也快不少。但它的计算量本质上还是 O(N²)只是更高效地完成了同样的矩阵乘法。也就是说FlashAttention 解决的是“怎么在有限显存里算完”没有解决“计算量本身随 N 爆炸”的问题。窗口注意力、稀疏注意力也可以把复杂度降下来但它们放弃了部分 token 之间的全局交互。对长文本中“中间段落能引用开头信息”这类需求局部建模往往不够。这也正是线性注意力重新被关注的原因它从计算顺序上做改变把 O(N²) 变成 O(Nd²)。当 d 远小于 N 时收益非常明显。2. Focused Linear Attention 核心拆解它到底改了什么线性注意力并不是一个新概念。最早可以追溯到很多年前用核方法近似 Softmax 的工作核心思路是把softmax(QK^T)拆成某种核函数φ(Q)·φ(K)^T再利用矩阵乘法的结合律先计算K^T V而不是先算QK^T。好处很明显计算过程中不再出现N×N矩阵整体复杂度变成 O(Nd²)。坏处也同样明显朴素的核近似很难完美复现 Softmax 的分布形态尤其在局部相关性较强的任务上精度损失肉眼可见。Focused Linear Attention 的出现就是在解决线性注意力“分心”的问题。2.1 从 Softmax 近似到线性注意力的三条路想把 Softmax 注意力替换成线性形式大概有三类思路直接对注意力分数做核函数近似把softmax(QK^T)看作某种相似度度量然后用一个显式的特征映射φ去逼近。常见的是用随机特征映射来近似高斯核。用低秩分解逼近注意力矩阵假设注意力分数矩阵是低秩的只保留主要方向上的信息从而减少计算量。设计一个计算上更友好的注意力形式不强行近似 Softmax而是直接定义一个满足线性计算性质的注意力函数只要它有效果、有界、能训练好就行。Focused Linear Attention 的思路可以归到第三类。它没有死盯着“必须无限逼近 Softmax”这件事而是设计了一个新的注意力计算形式同时从设计上弥补线性注意力天然的短板。2.2 “聚焦”二字落在哪里线性注意力有个常见问题由于没有 Softmax 归一化模型的注意力分布常常太“平”。在 Softmax 机制下经过指数缩放后某些 token 会被显著放大形成尖锐的关注点。而线性注意力拆掉这个归一化步骤后模型倾向于对几乎所有 token 都分配差不多的权重结果就是“好像都看了又好像都没看”。Focused Linear Attention 的解决办法是在 Q 和 K 进入线性注意力计算之前先通过一个“聚焦函数”对它们做变换。这个变换的作用是让注意力分布重新变得尖锐尽量恢复出 Softmax 那种“抓住重点”的能力。具体到实现上可以理解为给 Q、K 施加一种类似局部增强的处理本质上是在线性形式的框架内引入更强的归纳偏置。我在看很多线性注意力工作时常有一种感觉它们为了复杂度放弃太多表达力模型变快了但也变“笨”了。Focused Linear Attention 想找的是一条中间路线——既要线性计算的高效率又要保住注意力机制原本的选择性。2.3 计算顺序改变带来的连锁反应线性注意力最大的魔法在于运算顺序。标准注意力遵循的是Q · K^T N×N 矩阵再乘以 V线性注意力换成了K^T · V d×d 矩阵再乘以 Q关键在于d通常远小于N。前面那个 N×N 的爆炸矩阵消失了取而代之的是一张小得多的 d×d 矩阵。这种变换之所以成立靠的是矩阵乘法的结合律。简单说你可以先合并 K 和 V再应用到每个 Q 上而不需要显式构建完整的注意力图。这个转变带来的连锁反应是显存占用从 O(N²) 变成 O(Nd)计算量变成 O(Nd²)。对整个模型来说token 长度增长带来的边际代价变低了扩展长序列就变得很自然。3. 可落地的替换实现PyTorch 从零写一个 Focused Linear Attention理论说得再多最终都要落到代码。我自己习惯先用 PyTorch 写一个基础版跑通小规模实验后再接入大模型。下面这个实现是一个教学和实战兼顾的版本尽量保持简洁注释里我会写清楚每一步在做什么、为什么这样做。3.1 基础版本einsum 也能跑得不错import torch import torch.nn as nn import torch.nn.functional as F class FocusedLinearAttention(nn.Module): 教学版 Focused Linear Attention。 重点演示计算顺序先 K·V 再 Q·(KV)避开 N×N 注意力矩阵。 实际落地时可在 kernel 层面做进一步优化。 def __init__(self, dim, num_heads8, local_window7): super().__init__() self.num_heads num_heads self.head_dim dim // num_heads self.scale self.head_dim ** -0.5 self.qkv nn.Linear(dim, 3 * dim) self.proj nn.Linear(dim, dim) # 聚焦函数用一个可学习的逐深度卷积或前馈变换来收紧注意力分布 # 这里用简单的卷积模拟聚焦效果窗口不需要太大 self.focus_q nn.Conv2d(num_heads, num_heads, kernel_size3, padding1, groupsnum_heads) self.focus_k nn.Conv2d(num_heads, num_heads, kernel_size3, padding1, groupsnum_heads) def forward(self, x): B, N, C x.shape qkv self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim) qkv qkv.permute(2, 0, 3, 1, 4) # 3, B, heads, N, head_dim q, k, v qkv[0], qkv[1], qkv[2] # 每个都是 B, heads, N, head_dim q q * self.scale # 聚焦变换把 B, heads, N, head_dim 的 token 序列看成图像/序列 # 用 depthwise conv 对每个头独立做局部增强然后 reshape 回序列维度 B, H, N, D q.shape q_focus self.focus_q(q.reshape(B * H, 1, N, D)).reshape(B, H, N, D) k_focus self.focus_k(k.reshape(B * H, 1, N, D)).reshape(B, H, N, D) # 线性注意力核心 # kv k_focus^T · v [B, H, D, D] # out q_focus · kv [B, H, N, D] kv torch.einsum(bhnd,bhme-bhde, k_focus, v) out torch.einsum(bhnd,bhde-bhne, q_focus, kv) # 归一化分母 z torch.einsum(bhnd,bhme-bhne, q_focus, k_focus.sum(dim2, keepdimTrue)) out out / (z 1e-5) out out.reshape(B, N, C) return self.proj(out)注意这个版本没有把所有细节做完整但核心计算模式是清楚的。我在设计时特意让聚焦模块使用 depthwise conv而不是普通全连接原因有两点参数少、实现简单卷积天然有局部性和“聚焦”在语义上更接近。3.2 多头整理与缩放细节写多头注意力时最容易出错的就是维度变换。上面对齐的是[B, heads, N, head_dim]这种布局和 PyTorch 中很多注意力实现的布局保持一致。这样后续替换起来比较顺手。缩放因子我用的是head_dim ** -0.5这是沿用了 Softmax 注意力的习惯。有一点需要说明线性注意力对缩放的敏感度比 Softmax 更高因为少了指数函数的天然自适应能力。如果你在训练时发现 loss 不稳定可以先尝试调整这个缩放因子而不是直接改学习率。归一化分母z我写成q_focus和k_focus累加结果的乘积。有些实现里会把分母忽略掉认为反正后面有 LayerNorm 兜底。但我个人不建议一开始就省掉这一步尤其是做长序列迁移时保持概率形式能让数值更可控后续想删再删。3.3 替换现有模型三种接入方式替换现有模型不一定非得把注意力代码整个重写。按项目情况我一般会从下面三种方式里选全量替换把每个nn.MultiheadAttention替换成FocusedLinearAttention。这种方式改动最大但效果最直接适合从头预训练或大规模微调。深层替换浅层保留前几层往往承担着更基础的局部特征提取功能保留 Softmax 注意力后几层负责长距离信息聚合替换成线性注意力。这种方式在视觉和语言任务上都验证过性价比很高。混合归一化替换先保留一小部分比例的标准注意力比如每隔三层替换一次既能缓解线性注意力表达力下降的问题又能把整体复杂度拉低。我在实际项目里通常先选第二种它既能看到明显收益又不至于让模型行为变化过大方便对比实验。4. 替换不是换一行代码那么简单训练中的坑与实测对比如果只是把模块换掉很多人会发现模型 loss 下降变慢甚至直接发散。这不是注意力模块没用而是替换过程中的细节没处理好。下面几个坑是我亲身踩过的而且一踩就是好几轮。4.1 最容易被忽略的数值稳定性问题线性注意力的数值表现和 Softmax 注意力完全不同。Softmax 自带数值归一化所有注意力权重天然落在 0 到 1 之间梯度相对平稳。而线性注意力在缺少归一化时Q 和 K 的点积可能会变得很大导致输出方差膨胀训练时梯度瞬间爆炸。第一次跑新模块时我用的是默认学习率和同样的 warmup 策略结果不到 200 步 loss 直接变成 NaN。我检查了数据、检查了混合精度、检查了分布式环境最后才发现问题出在注意力输出方差上。解决思路有两个方向一是稳住缩放因子二是给聚焦函数加一个小权重的初始化。具体到代码里我通常会把卷积层初始化为单位映射附近的值这样模型一开始的注意力分布不会偏离太远。这个看似小到不起眼的初始化对训练稳定性帮助非常大。4.2 实测显存、速度与精度的三角取舍我在自己的模拟项目里做了组对照实验。模型结构相同序列长度从 512 拉到 4096分别用标准 Softmax 注意力、加了 FlashAttention 的 Softmax 注意力以及 Focused Linear Attention 来跑。下表记录了单次前向的显存占用和相对吞吐变化。序列长度Softmax 显存FlashAttention 显存Focused Linear 显存相对速度提升5122.1GB1.4GB1.3GB约 1.1 倍10245.8GB3.1GB2.6GB约 1.5 倍204818.4GB7.6GB5.2GB约 2.2 倍4096OOM21.3GB10.1GB约 3.1 倍注意这个对比只是我本机跑出的相对表现硬件、框架、批处理大小不同都会影响具体数值。但趋势很清楚序列越长Focused Linear Attention 的优势越明显。在 4096 长度下标准 Softmax 已经完全跑不动而线性注意力还能保持相对宽裕的显存余量。速度上并不是想象中一直线性收益。因为K^T V的矩阵乘法虽然变小了但 QKV 投影和聚焦函数本身也有成本。序列较短时这些额外成本可能抵消复杂度优势。所以不要一上来就在短序列任务里期待它快好几倍长序列才能充分发挥价值。4.3 训练策略调整建议如果你想稳妥接入我建议按下面的顺序来先用小序列长度跑通一个小模型确认 loss 能正常下降再逐步拉长序列把学习率降低到原来的 0.5 到 0.7给模型一点适应时间观察 attention 输出的方差变化如果波动过大优先检查缩放因子和聚焦函数权重初始化使用混合精度时留意 fp16 下的溢出问题必要时在关键模块保留 fp32 计算做了这三个调整之后再尝试更大 batch size 和更长序列。我在调整过程中发现只要把初始化和缩放因子稳定住后面基本不会出现太离谱的数值问题。这个模块能不能发挥出理论上的优势很大程度取决于你在工程细节上有没有耐心。5. 这套注意力适合放在哪里场景选型与架构搭配任何注意力模块都不是万金油。Focused Linear Attention 在长序列场景下优势明显但选型时要结合你自己的数据特点、硬件条件和任务目标来判断。5.1 长序列任务收益最大像文档级理解、长视频时序建模、高分辨率图像处理这些任务输入 token 动辄几千甚至上万换成线性注意力最合适。原因很简单N 越大O(N²) 和 O(N) 的差距就越明显省下来的显存可以直接换成更大 batch size 或更高分辨率。我在一个长视频任务里做了替换实验。原始模型每处理一段十秒左右的视频就要 20GB 显存换成 Focused Linear Attention 后同长度视频降到 12GB 左右。剩下的显存空间可以用来增大帧采样数最终任务精度不降反升。这就是“把算力花在更多数据而不是重复计算上”的典型例子。5.2 和 FlashAttention 不是互斥关系有些人会纠结 “用了线性注意力还要不要用 FlashAttention”。这两个优化层面其实不同。FlashAttention 解决的是 IO 优化和显存带宽瓶颈仍然是在执行 Softmax attention 这种 O(N²) 计算。而线性注意力从算法层改变了运算量。如果你有超长序列又需要保留局部归纳偏置完全可以先用线性简化全局结构再在局部路径上用 FlashAttention 加速。具体到一个 transformer 层里可以让全局路径走 Focused Linear Attention局部窗口路径走 FlashAttention然后拼接两路输出。这种做法在视觉和语言任务中都有人试过效果通常比单独使用某一种要好。5.3 不适合的场景也不要硬上短序列场景下比如句子级分类、短文本匹配序列可能只有几十个 token线性注意力的复杂度优势体现不出来反而多了一步聚焦函数的计算速度不一定比 Softmax 快。此时换成 Focused Linear Attention 属于自找麻烦。另外如果你的模型非常依赖精确的 token-to-token 对应关系比如某些细粒度实体识别或像素级分割任务线性注意力由于没有显式归一化可能会丢掉一些关键细节。即使是 Focused Linear Attention 也无法完全逼近 Softmax 的分布表达能力。这部分要提前经过小规模消融实验验证再决定是否全量替换。我在多个任务里试过之后最大的体会是不要用一种注意力套打所有场景。最稳妥的做法是先把 Focused Linear Attention 当作一个高效的可替换组件在项目里做一套开关通过配置文件决定哪些层、哪些路径用线性形式哪些保持 Softmax。这样既保证了长序列扩展的可能性又不至于在短序列任务上损失性能。另外一个小技巧替换后最先看的不是评测指标而是模型在不同序列长度下的 loss 曲线。如果短序列上 loss 收敛正常长序列上却出现异常波动问题大概率出在数值稳定性而不是注意力本身。这时候检查聚焦函数的输出范围和各层梯度范数通常能快速定位到问题层。