这两年在大模型上做长文本几乎绕不开同一个坎注意力机制的计算复杂度是 O(n²)。我最初也觉得不就是矩阵乘吗直到亲手把序列长度从 4K 推到 32K、128K才发现全量注意力的账单会膨胀到让人怀疑人生。今天想把这几年看到的大模型注意力全景图整理一下重点聊聊业界常说的五条降本路线——稀疏注意力、线性注意力、低秩近似、KV Cache 压缩以及混合架构——它们到底都在解决什么问题又凭什么能省钱、省在哪、代价是什么。这篇内容适合正在做长文本模型训练、推理优化或者在自研模型结构里纠结“要不要换成线性注意力”的人。我尽量不用教科书写法全部按实际操作中会遇到的问题来讲也会给出可以照抄的选型思路和实现片段。1. 注意力机制的“贵”到底贵在哪1.1 一张注意力图的完整运算过程先回到最基础的东西。一个标准的缩放点积注意力输入是序列的隐藏状态 H形状是 [n, d]n 是序列长度d 是隐藏维度。通过三个权重矩阵分别得到 Q、K、V然后计算注意力分数矩阵 S QK^T / √d再做 softmax 得到权重矩阵 A最后输出 A V。这里最扎眼的就是 QK^T 这一步。它的形状是 [n, n]所以矩阵乘法的复杂度是 O(n² d)。当 n 从 4096 涨到 131072中间矩阵的元素数量会从约 1600 万个涨到约 170 亿个膨胀了一百倍还不止。FP16 精度下4K 序列的注意力分数矩阵才 32MB 左右挤一挤还在显存里放得下到 128K 序列时单这一个矩阵就要 33GB 以上一张主流显卡根本装不下。这个 n×n 矩阵就是整张“注意力图”。它表示每个 token 对其它所有 token 的关注强度。很多人第一次接触 O(n²) 时只会背“长序列成本高”这句话但真正训练时才会发现成本不只是矩阵乘法本身而是整条链路都被这张图卡住了。1.2 O(n²) 的三种账单算力、显存、带宽第一种账单是算力。训练阶段每个 token 都要和所有历史 token 做点积FLOPs 随序列长度平方增长。同样一个 batch把长度从 2K 提到 8K单次前向的注意力算力需求直接变成原来的 4 倍而不是大家直觉里的“长了 4 倍成本涨 4 倍”。第二种账单是显存。除了中间矩阵反向传播还要保存注意力权重用于梯度计算。全量注意力在训练时如果不开任何重计算优化会额外存一份 [n, n] 的注意力矩阵层数一多显存直接爆炸。这也是 FlashAttention 这类算子流行的原因——它把中间矩阵切块计算不显式存储整张图本质是在“用重新计算换显存”。第三种账单是推理带宽。自回归解码是一次生成一个 token计算量看起来不大但每生成一个新 token都要把之前所有 token 对应的 K 和 V 拿出来做一次 attention。这些 K、V 需要缓存在显存里也就是 KV Cache。序列越长KV Cache 越大显存带宽就成了瓶颈。我实测过一个 70 亿参数模型在 128K 上下文下单层单头的 KV 数据就可能到几十 MB几十层叠下来KV Cache 比模型权重还占显存。所以说O(n²) 不只是一个复杂的数学标记它在训练、长序列推理、并发服务三个环节分别变成了算力墙、显存墙和带宽墙。不同的人说“注意力太贵”指的其实是不同的墙。1.3 为什么不能直接砍复杂度既然全量注意力这么贵那直接简化它不就行了比如每个 token 只和最近的 token 做注意力这不就把复杂度降到 O(n) 了问题在于注意力机制的核心能力是“从任意位置回忆信息”。全量注意力就像开会时每个人都能随时向场上任意一个人提问而简化后的注意力相当于只允许你向旁边几个人提问。很多任务吃这个“全局回忆”能力。比如一段长文档里结论在第 10000 个 token而支撑它的关键证据在第 500 个 token局部注意力根本联系不起来。另外softmax 本身是一个全局归一化操作分母要对一整行做求和。哪怕这一行里大部分位置不关键你也不能随便把分母砍掉否则输出的概率分布不合法训练会不稳定。这也是很多“看起来能 O(n)”的方案真正实现时发现并不好用的原因——它们改变了 softmax 的归一化语义模型训练起来就不再是原来的优化问题了。所以所有降本路线本质上都在回答一个问题怎么在不动摇“注意力能选择性回忆”这个能力的前提下把计算复杂度从平方级降下来。接下来的五条路线就是在不同角度回答这个问题。2. 五条降本路线全景对比2.1 路线一稀疏注意力用先验换范围稀疏注意力是最直观、也最早被广泛使用的路线。它不做无差别的全量计算而是预设一个“哪些 token 之间需要交互”的规则只按照规则计算注意力矩阵中的部分元素。常见做法有几种滑动窗口每个 token 只关注前后 W 个 token复杂度变成 O(nW)。全局 token在一段序列里安排少数特殊 token它们可以关注整个序列序列里所有 token 也都能关注它们。随机稀疏随机选一部分 token 对做注意力提供一条“任意位置长距离连接”的通路。这种路线最大的优点是实现简单很多都能在现有 Transformer 上直接改 mask 就行。我自己试过在一个 1.8B 模型上把注意力改成滑动窗口加全局 token训练吞吐能提高 2 到 3 倍代价是部分长程推理任务的效果明显下降尤其是需要“精确定位某个遥远细节”的任务。稀疏注意力的缺点也显而易见——你提前假设了信息会在哪里出现。如果信息的分布和你的 mask 规则不匹配模型就只能靠层层传递来弥补深层堆叠加深之后一方面训练难度增加另一方面模型对远距离信号的感知会变得“模糊”。2.2 路线二线性注意力把 softmax 换成核函数线性注意力的思路比稀疏注意力更进一步。它不满足于“只看一部分”而是想彻底抛开 n×n 矩阵把注意力变成线性复杂度。关键点是softmax 可以写成 exp(QK^T) 的形式。如果把 exp 这个非线性核函数近似成两个特征映射的乘积也就是把 Q 和 K 先各自映射到一个高维空间那么注意力公式可以改写成$$(QK^T)V ≈ φ(Q)(φ(K)^T V)$$这里先算 φ(K)^T V得到一个相当于“全局记忆”的小矩阵再和 φ(Q) 相乘复杂度变成 O(n d²)不再依赖 n 的平方。只要 d 远小于 n这个复杂度基本就是线性的。这类方法的代表有随机特征近似也就是把 exp 核映射成有限维随机特征也有直接用简单核函数替换 softmax 的做法虽然和 softmax 语义不完全等价但实践中也能跑。我实际用下来线性注意力在超长序列64K 以上上的吞吐优势非常明显显存占用也稳定。但要注意它的核心问题是“表达能力不如 softmax”。softmax 能让某个 token 只关注极少数关键位置而很多线性注意力等价于对所有历史信息做了“平均式”汇总在需要强局部聚焦的任务上会弱一些。2.3 路线三低秩近似压缩注意力矩阵本身低秩近似走的是另一条路既然 n×n 矩阵太大那就假设它实际上是“低秩”的用一个小得多的矩阵去近似它。举一个简单例子不直接算 QK^T 得到 n×n而是先对 Q 做一次投影把它从 [n, d] 压成 [n, k]k 是一个固定常数或 log n 级别的大小。这样注意力矩阵变成 n×k后面再映射回 n。计算复杂度变成 O(n k d)当 k 远小于 n 时相当于 O(n)。这类方法在理论上很优雅但落地时有个尴尬的情况“注意力矩阵是不是低秩”这件事本身和任务强相关。如果输入序列里的信息高度冗余比如都是重复的固定格式文本低秩近似效果还行但如果是语义丰富、很难压缩的语料强行用低秩近似会丢掉很多细节。我的体验是低秩近似更适合做“压缩模块”而不是“完整替代”。比如在多层模型里只在靠中间的某一两层用低秩注意力做全局信息汇总前后还是用稀疏或全量注意力保存细节效果比整个模型全换低秩要稳。2.4 路线四KV Cache 压缩从工程上抠成本前面三条路线改的是模型结构第四条路线更偏推理工程——它不改变注意力的计算语义而是压缩 KV Cache 本身。因为长序列推理时真正的瓶颈往往不是计算矩阵乘而是每生成一个 token 就要把巨量 KV 数据从显存搬到计算单元。KV Cache 压缩的主流手法有量化。把 K、V 从 FP16 压成 INT8 或 INT4显存直接减半或减到 1/4。多查询共享。让多个 query 头共享同一个 key/value 头KV Cache 量可以除以 group 大小。淘汰策略。生成过程中只保留“历史中更可能被反复关注的 token”的 KV丢掉其它部分。经典的思路是看注意力分数的累积分布保留高分 token。层间分配。不同层的 KV Cache 重要程度不一样给重要层多留非重要层少留。这几种方案我都在线上服务里试过。量化最省事但要小心评估精度损失某些数据集上 INT4 KV 会让输出质量明显下滑多查询共享已经在不少开源模型里成了标配属于“免费的午餐”淘汰策略效果最好但实现复杂因为它需要实时维护哪些 KV 要丢、哪些要留还要防止后续 token 又需要已经被丢掉的 KV。2.5 路线五混合架构把不同复杂度组合起来第五条路线是近两年特别受关注的混合架构。它的核心信念是不需要让整个模型都只用一种注意力而是不同层、不同模块用不同复杂度。最常见的做法是把全量注意力层和线性复杂度层交替排列。比如浅层处理局部语法和词法信息用局部注意力就够深层需要长程推理放少量全量注意力层中间用线性层做信息压缩和传递。另一种做法是在一层内部混合比如一部分头做局部窗口注意力一部分头做全局 token 注意力。混合架构的好处是它把前几条路线的优势组合起来避免了“一刀切”的缺点。我见过的一个比较典型的 7B 级别模型前 1/3 层做滑动窗口注意力中间 1/3 层做窗口加全局稀疏注意力最后 1/3 层做全量注意力。和全量 Transformer 相比训练成本下降一半左右且大部分下游任务没有掉点长程检索任务还保留了接近全量的水平。难点在于混合架构的超参组合非常多每层放什么注意力、全局 token 放多少、窗口多大、全局注意力放在第几层这些都需要通过小规模实验来验证。配置不好会出现一种很难排查的现象模型某些层“吃不饱”某些层“过载”最终损失下降缓慢。2.6 五条路线到底在争什么看到这里你会发现五条路线并不是互斥的甚至经常组合使用。它们真正在争的是“在表达能力、计算开销、工程复杂度之间找到一个帕累托最优”。稀疏注意力争的是“用先验知识减少计算”线性注意力争的是“用核函数改写计算”低秩近似争的是“用矩阵结构压缩计算”KV Cache 压缩争的是“用工程手段绕开存储瓶颈”混合架构争的是“在不同位置各取所长”。没有哪条路线是绝对胜利的。决定胜负的是你的具体场景序列多长算力充裕还是显存紧张任务更吃长程推理还是局部建模下面我写一下实操落地时的选型思路和实现经验。路线核心思路训练侧效果推理侧效果主要坑稀疏注意力预设交互范围显存、算力下降明显吞吐提升明显全局信息易丢失线性注意力核函数改写复杂度长序列可训练显存稳定短任务掉点低秩近似压缩注意力矩阵维度理论复杂度低主要省显存信息冗余假设过强KV Cache 压缩压缩推理缓存不影响训练吞吐和显存双优化精度损失和实现复杂度混合架构不同层用不同方法综合性价比高视配置而定超参组合难调3. 实操落地如何选型与实现3.1 先问自己三个问题动手改结构之前先别急着追热点。我在选型前一定会先确认三件事第一现在最痛的是训练还是推理如果是训练长序列时 OOM首要目标应该是降低显存占用如果是线上推理并发上不去首要目标可能是 KV Cache 压缩和计算效率。不同的痛点最优路线完全不同。第二你需要多长的有效上下文如果你的场景里 4K 就够那完全没有必要去上线性注意力。反而是那些“名义上支持 128K实际超过 8K 就忘事”的模型才需要认真考虑结构层面的改造。第三任务里是否存在“长期依赖精确定位”的需求比如多轮对话里用户第一句提到的某个细节隔了一百轮还要能正确引用。这种任务对全量注意力的依赖非常强稀疏和线性方案都要小心。这三个问题想清楚方案基本就筛掉一大半了。拿我自己举例做一个“长文档问答 摘要”的私有化部署项目时第一版直接上全量注意力结果 32K 输入把一张 80GB 卡吃满并发只能开 1 路。后来确认任务属于“局部线索密集、偶尔需要跨文档引用”就改成了混合架构而不是纯线性方案。3.2 选型对照表我平时会给团队出一张非常务实的对照表适合直接抄作业场景画像推荐路线理由训练超长序列GPU 显存有限稀疏注意力 线性层交替显存成本可控训练稳定线上推理吞吐瓶颈上下文中等KV Cache 量化 多查询共享工程改动小立刻见效文档量极大长期依赖强允许离线混合架构全局局部在表达力和成本间平衡端侧部署内存严格受限低秩近似或线性注意力参数和缓存都更小通用基座模型不知道下游任务保守混合架构保持部分全量层保留通用能力风险最低注意表格只是起点不是终点。真正确定方案后我强烈建议先拿一个 100M 左右的小模型用两倍目标序列长度跑 5 到 10 轮小步数预热实验对比候选方案的 loss 曲线和下游任务基线再决定是否全量训练。这一步能省下大量后面返工的时间。3.3 一个稀疏注意力的实现片段如果只想快速验证稀疏注意力最直接的方式就是在原始全量注意力的 mask 上动手。下面是一个非常粗的伪代码示例用来生成滑动窗口加全局 token 的注意力 maskimport torch def build_sparse_mask(n, window256, global_every64): # n: 序列长度 mask torch.zeros(n, n, dtypetorch.bool) for i in range(n): # 滑动窗口只看前后 window 内的 token start max(0, i - window) end min(n, i window) mask[i, start:end] True # 全局 token每隔 global_every 个位置放一个 # 这些 token 能和所有位置交互 if i % global_every 0: mask[i, :] True mask[:, i] True return mask这里有几个细节值得说。第一mask 必须是 bool 或 float 类型不能直接用 0/1 整数否则在部分算子实现里会报错或静默出错。第二滑动窗口的尺寸不是越大越好我用 256 的窗口处理 32K 序列时模型还是能感知到“中间层堆叠带来的间接长程信息”但如果窗口小到 64很多文本生成任务会明显变差。第三全局 token 的位置要么均匀分布要么放到每段开头我试过随机分布效果不稳定。实际训练时这个 mask 要参与注意力分数的 softmax通常把 mask 为 False 的位置替换成负无穷而不是直接删掉那些位置。因为算子内部的并行计算还是按固定形状做的直接改形状反而会牺牲性能。3.4 一个线性注意力的粗实现线性注意力的实现核心在于“先算 K^T V再乘 Q”。下面是一个最小可用风格的示例重点展示结构没有加完整优化import torch import torch.nn.functional as F def linear_attention(q, k, v): # q, k, v: [batch, heads, seq_len, head_dim] # 简单核函数用 ReLU 映射做非线性 q F.relu(q) k F.relu(k) # 先算全局记忆k 转置后和 v 相乘 kv torch.einsum(bhnd,bhde-bhne, k.transpose(-1, -2), v) # 再和 q 相乘得到输出 out torch.einsum(bhnd,bhne-bhde, q, kv) # 归一化因子保证输出尺度稳定 z torch.einsum(bhnd,bhne-bhde, q, k.sum(dim-2, keepdimTrue)) return out / (z 1e-6)这段代码没有用随机特征近似 softmax而是用 ReLU 做了简单核替换。优点是实现极其简单训练长序列时不炸显存缺点是它和 softmax 的语义差别很大不能直接当作“换了一个更快的方法”就完事。用线性注意力时有三个坑。第一个是数值稳定性分母里的归一化因子如果接近 0输出会爆炸所以要加极小值保护。第二个是初始化要特别小心线性注意力的梯度分布和全量注意力不一样直接沿用原模型的初始化策略很容易在早期出现 loss 不下降。第三个是位置的相对信息会变弱因为注意力矩阵不再能显式地给不同距离的 token 分配不同的偏置如果模型本身没有 RoPE 这类位置编码效果会比预期差很多。3.5 混合架构的层分配经验混合架构是五条路线里上限最高的但也是最难调的。核心问题是哪几层用全量注意力哪几层用稀疏注意力哪几层用线性层我根据多次实验总结出一个相对可靠的初始经验浅层靠近输入更适合局部注意力。因为浅层主要负责抽取局部语法、实体和词汇信息还没到需要跨长距离整合的阶段中间层可以穿插少量全局注意力让信息在整个序列范围内流动深层靠近输出放少量全量注意力层给最后的高层语义整合提供充足表达能力。一个可以当作起点的模板是 20 层模型第 1-8 层滑动窗口注意力窗口 256。第 9-14 层窗口注意力 全局 token每 64 个 token 放一个。第 15-18 层全局注意力但间隔使用即每 2 层中只有 1 层做全量。第 19-20 层全量注意力。这个模板在短文本任务上几乎不掉点在长文档任务上比“全部稀疏”的模型好很多。当然不同模型的深度和隐藏维度不同模板要按比例缩放不要硬抄。还要注意混合架构各层的收敛速度可能不一致。全量注意力层通常收敛更快局部注意力层稍慢。实践里可以在训练中期对全量注意力层做较小的学习率或者给不同层设置不同的 dropout避免浅层过拟合、深层欠拟合。3.6 评估指标怎么定改完模型结构不能只看 loss 或者 BLEU 这类单一指标。我一般的做法是固定一组评测集至少包含三类任务短文本理解“只要局部信息就能回答”的任务用来确认局部建模能力没有退化。长程检索“答案在很远的某个位置出现”的任务用来确认长距离信息没有断裂。多跳推理“需要把两三个相隔很远的信息点串起来”的任务用来确认注意力语义里的“选择性回忆”仍然有效。同时还要量测训练吞吐、峰值显存、推理吞吐、首 token 延迟。毕竟降本的目标是把成本降下来不能只盯着指标好看却把速度拖垮。4. 常见问题与排查实录4.1 训练收敛了推理反而变慢有次我把一个模型从全量注意力改成稀疏注意力训练吞吐确实涨了但线上推理时发现延迟反而比原来高了一截。查了半天发现稀疏 mask 导致的“不规则内存访问”让 GPU 的并行效率下降如果没有对应的融合算子优化实际计算时间并不比全量少多少。排查思路是先做 profiling看算子是卡在矩阵乘法还是卡在 memory 拷贝。如果是内存访问不规则的问题优先检查算子库是否支持稀疏注意力模式或者在实现时把稀疏 mask 拆成块状结构保证访问的连续性。另一个办法是直接使用 FlashAttention 类的融合算子把 mask、softmax、乘 V 都合到一块处理。4.2 线性注意力在短任务上掉点线性注意力最典型的问题训练长序列时很爽但切到短文本任务分数反而比全量注意力模型低一大截。原因很简单短序列下 n 很小线性注意力省不了多少算力却因为表达能力的上限不如全量注意力精细。我的建议是不要对全部任务都使用同一套结构。如果模型要同时兼顾短文本和长文本可以让模型在短输入时走全量注意力分支长输入时走线性注意力分支或者直接在混合架构里保留充分的全局层。另一种做法是在线性注意力后面加一个轻量的局部注意力补充恢复一些局部精细建模能力。4.3 RoPE 与稀疏注意力打架旋转位置编码RoPE是目前很多模型的标准配置但它和稀疏注意力结合时有个容易被忽视的问题RoPE 的作用体现在两个 token 之间的相对位置差上而稀疏注意力因为 mask 只让部分 token 对有机会计算注意力这种“位置差”的表达就会在未被 mask 覆盖的路径上失效。结果就是模型对距离的感知不再连续偶尔会出现“换了表述方式就判若两人”的情况。解决办法是要么在训练阶段就保证 RoPE 的感知范围大于窗口大小要么改成可学习的相对位置偏置要么在窗口边界做重叠让相邻窗口之间有 token 可以互相看到减少距离信息的断层。我实测中窗口重叠 10% 到 20% 能明显改善这种现象。4.4 KV 量化后困惑度飙升KV Cache 压到 INT4 之后线上服务吞吐确实涨了但部分任务输出质量下滑困惑度上升明显。原因是 KV 量化对注意力分布的尾部位置比较敏感尤其是那些需要精确回忆的 token它的 K/V 一旦被量化损失影响检索精度就下降了。排查时可以先做逐层精度对比找出哪些层量化后损失最大。一般来说靠近输出的层对 KV 精度更敏感可以在这些层保留 FP16前面层用 INT8 或 INT4做一个“分层混合精度”。另外我建议量化时把异常值单独处理避免整体分布被少数大数值拉高误差。4.5 混合模型训练 loss 反复震荡混合架构训练时另一个常见现象是 loss 在中期反复震荡不下降。这通常是因为注意力层的种类的切换导致不同层的学习进度不匹配。比如局部注意力层已经收敛但一个全量注意力层还在大幅更新反馈给前面的局部层后局部层只能跟着反复调整。我处理这种问题时会做三件事先给全量注意力层单独设更低的峰值学习率再检查各层梯度的 L2 范数如果某一层的梯度明显高于其它层就加上梯度裁剪最后如果震荡还是严重可以适当提高全局 token 的比例让梯度信号更均匀地传递到所有层。4.6 速查表现象可能原因快速排查解决思路训练快但推理慢稀疏算子访问不连续profiling 看 memory 占比融合算子、块状 mask短任务掉点线性注意力表达能力不足对比短序列任务分数局部注意力补充、分支切换长程任务变差mask 规则和任务不匹配单测长程检索任务增大窗口、加全局 token位置感知异常RoPE 和稀疏 mask 冲突检查距离-分数相关性窗口重叠、改用相对偏置KV 量化后质量下降敏感层被过度压缩分层量化误差对比关键层保留高精度loss 震荡混合层学习速率不匹配分层梯度统计分层学习率、梯度裁剪显存偶尔峰值溢出长序列下中间矩阵分配看训练日志显存曲线FlashAttention、重计算最后再分享一点我的个人体会。做了这么多长文本优化最大的感受是模型结构没有免费的午餐所谓降本路线最后都是在跟“注意力要保留多少选择权”做交易。你在复杂度上省下来的每一分都要在表达能力或工程复杂度上还回去。真正可靠的方案不是选一条最酷的路线而是先测出自己项目的瓶颈到底在哪面墙再去五条路线里找最匹配的组合。如果你想动手我建议不要直接改大模型先用百万级参数的小模型跑通选型流程再把经验放大到大模型上这条路走得最稳。