💬 NLP

SPLA: Block Sparse Plus Linear Attention for Long Context Modeling

该论文介绍了 SPLA，一种结合了块状稀疏精确注意力与残差线性注意力模块的新型框架，通过准确选择相关块并压缩剩余数据且不产生 IO 开销，从而高效地对长上下文进行建模，并在长上下文基准测试中超越了密集注意力模型。

原作者： Bailin Wang, Dan Friedman, Tao Lei, Chong Wang

发布于 2026-02-02

📖 1 分钟阅读☕ 轻松阅读

原作者： Bailin Wang, Dan Friedman, Tao Lei, Chong Wang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，将大语言模型（LLM）比作一位才华横溢的图书管理员，正试图根据一座巨大的图书馆来编写一个故事。随着故事变得越来越长，管理员必须记住越来越多的过去细节。

问题：“沉重的背包”与“错误的猜测”
当故事变得非常长（数百万词）时，管理员面临两个大问题：

沉重的背包： 为了记住所有内容，管理员必须背着一个装满索引卡（即“KV cache”）的背包。随着故事的增长，背包变得越来越重，导致管理员移动得越来越慢，最终因为无法承受重量而陷入停滞。
错误的猜测： 为了节省空间，之前的一些方法尝试扔掉大部分索引卡，只保留那些管理员“认为”重要的卡片。但这些方法并不擅长猜测。它们经常扔掉关键页面（低“选择保真度”），更糟糕的是，它们会完全忽略图书馆的其余部分。这导致故事失去了主线剧情，因为那些不那么显眼但仍然重要的细节（即“长尾”部分）被完全删除了。

解决方案：SPLA（聪明的图书管理员）
该论文介绍了 SP LA（块稀疏加线性注意力，Block Sparse Plus Linear Attention），这是一种让管理员在不丧失理智且保持速度的情况下处理图书馆的新方法。它通过一个两部分组成的系统运作：

1. “智能搜索”（更好的选择）

SPLA 不仅仅是靠猜测哪些索引卡重要，而是使用一种数学技巧（称为“二阶泰勒展开”）来精确计算哪些文本块最重要。

类比： 想象管理员不再仅仅通过书名来判断一本书是否重要，而是快速检查这本书的“平均氛围”和“主题多样性”（均值与方差）。这能帮助他们比以前更准确地找到真正关键的页面，确保不会意外扔掉某个反转剧情的关键章节。

2. “魔法挤压”（残差线性注意力）

这是最核心的部分。在旧的方法中，如果一个文本块没有被选为“超级重要”，它就会被扔进垃圾桶。SPLA 说：“不准扔垃圾！”

类比： 想象管理员有一个特殊的“魔法海绵”。
- 对于最重要的页面（“峰值”），他们逐字逐句地阅读（精确注意力/Exact Attention）。
- 对于不太重要的页面（“长尾”），他们不是把它们扔掉，而是使用魔法海绵将所有信息挤压成一个微小、紧凑的摘要状态。
- 诀窍： 管理员实际上不需要再去重新阅读那些“被挤压”过的页面。他们通过从“整个图书馆的摘要”中减去“刚刚阅读的重要页面”来计算出这个摘要。这意味着他们无需在手中实际加载那些沉重且不重要的卡片，就能获得所有信息的益处。

为什么这很重要

不再“丢失剧情”： 通过保留不重要部分的“挤压”摘要，模型在故事变长时不会丢失上下文。它缩小了“快但笨”（稀疏模型）与“慢但聪明”（稠密模型）之间的差距。
速度： 因为管理员只在物理上加载最重要的卡片，所以他们可以运行得更快，尤其是在故事规模巨大的时候。
易于升级： 论文表明，你可以拿一个现有的、强大的图书管理员（预训练模型），并赋予他们这种新的“智能搜索 + 魔法海绵”系统，而无需从头开始重新训练。这就像是给一位资深图书管理员一套新工具，让他们在不改变思维方式的前提下变得更高效。

实验结果
作者在一个 140 亿参数的模型上进行了测试。他们发现 SPLA：

在通用知识和推理方面，表现得与那些缓慢、沉重的模型一样出色。
在极长任务（高达 256,000 词）上碾压了竞争对手，而在这些任务中，其他快速模型开始失效并出现错误。
保持了高水平的速度，证明了你无需在“快”与“聪明”之间做选择。

简而言之，SPLA 是一种让 AI 模型能够快速处理海量文本而不丢失细节的方法，其核心在于更聪明地决定哪些内容需要精读，以及如何对剩余内容进行总结。

技术摘要：SPLA —— 用于长上下文建模的块稀疏加线性注意力机制

1. 问题陈述

随着大语言模型（LLMs）扩展到处理数百万 token 的上下文窗口，解码阶段由于内存限制成为了关键瓶颈。全量注意力机制需要将日益庞大的键值（KV）缓存加载到内存中，这受限于内存带宽，并限制了推理吞吐量。

现有的块状稀疏注意力方法（如 NSA、InfLLM）试图通过仅加载最相关的 KV 缓存块来缓解这一问题。然而，本文指出了当前方法的两个根本局限性：

选择保真度低（Low Selection Fidelity）： 当前的块选择方法依赖于高度压缩的表示（例如 min/max 池化、简单的均值池化），这些表示缺乏与原始 token 级注意力目标在数学原理上的联系。这导致了高召回错误，即错失了具有语义关键性的块。
累积上下文损失（长尾发散，Long Tail Divergence）： 现有方法严格截断未被选中的块（将其贡献设为零）。随着上下文长度增加，注意力分布中的概率质量在“长尾”部分变得不可忽略。丢弃这些质量会导致稀疏注意力输出与稠密基准发生显著发散，从而降低生成质量，尤其是在长上下文场景下。

2. 方法论：稀疏加线性注意力（SPLA）

SPLA 通过将注意力机制分解为两个互补的部分来解决这些局限性：用于高相关性块的精确注意力（Exact Attention），以及用于剩余“长尾”部分的近似线性注意力（Approximate Linear Attention）。

2.1 通过二阶泰勒展开进行原则性块选择

为了提高选择保真度，SPLA 直接利用二阶泰勒展开从原始注意力目标中推导出块选择度量。

目标： 块相关性得分被定义为块内所有 token 的集成概率质量。
近似： 直接求和开销巨大。SPLA 近似计算了围绕块均值 ( $\bar{k}$ $\overset{ˉ}{k}$ ) 的指数注意力得分的期望。
- 该推导同时利用了一阶统计量（块均值）和二阶统计量（块协方差/方差）。
- 得到的得分公式为： $\exp(q^\top \bar{k}) \left(1 + \frac{1}{2}q^\top \text{Cov}(k)q\right)$ 。
实现： 为了避免 $O(d^2)$ 的存储开销，协方差矩阵被近似为对角矩阵。这使得利用块统计量的逐元素操作高效计算选择度量成为可能，在无需辅助训练的情况下，弥合了 token 级注意力与块级检索之间的差距。

2.2 用于“长尾”的残差线性注意力（RLA）

SPLA 并非丢弃未选中的块，而是通过一种专门的线性注意力机制将它们压缩为一个紧凑的循环状态。

递归性： RLA 状态专门累积来自未选中块的信息。
IO 高效的公式化： 朴素的实现可能需要重新扫描内存以获取未选中的块，这会违背稀疏化的初衷。SPLA 推导了一种优化的基于减法的公式化方法：
$o^{rla}_t = \bar{o}_t - \tilde{o}_t$
其中：
- $\bar{o}_t$ 是通过对所有块进行标准线性注意力计算得到的全局上下文估计。
- $\tilde{o}_t$ 是通过仅对选定块进行线性注意力计算得到的局部修正项。
效率： 由于选定的块已经加载到 SRAM 中进行精确稀疏注意力计算，因此 $\tilde{o}_t$ 可以在同一个融合算子（fused kernel）内完成计算。这确保了在稀疏解码过程中，未选中的块从未被显式访问高带宽内存（HBM），从而保持了稀疏解码的 IO 特性，同时保留了完整的全局上下文。

2.3 架构与集成

混合策略： 最终输出是稀疏注意力（精确）与 RLA 输出（近似）的门控加权和，并通过 RMS Norm 进行归一化。
参数效率： SPLA 旨在适配预训练的稠密模型。它在两个分支中使用了共享的查询（query）、键（key）和值（value）投影。唯一的额外参数是应用于 RLA 输出的 RMS Norm 的缩放参数。
特征映射： 使用了无参数的指数特征映射 ( $\phi(x) \sim \exp(x)$ )，确保不会引入新的投影权重。

3. 核心贡献

原则性块选择： 基于二阶泰勒展开推导出的块选择策略，利用块级均值和协方差统计量，实现了比启发式池化方法更高的选择保真度。
SPLA 框架： 一种全新的架构，将上下文严格划分为“精确”集和“近似”集。通过使用残差线性注意力处理未选中的块，它保留了完整的全局上下文，缩小了与稠密注意力的性能差距。
IO 高效实现： 一种优化的基于减法的 RLA 公式化方法，避免了对未选中块的显式内存访问，确保了与标准稀疏解码相当的推理效率。
适配性： 证明了稠密注意力模型可以通过持续预训练（continual pretraining）转化为 SPLA 架构，且仅需极少的额外参数。

4. 实验结果

作者通过在三个设置（通用知识、长上下文、推理）下进行持续预训练（CPT），对一个 14B 参数模型进行了评估。

通用知识： SPLA 在标准基准测试（ARC、MMLU、GSM8K 等）上保持了与稠密注意力模型竞争力的表现，通常能达到或略微超过稠密基准，并优于其他稀疏基准（NSA、InfLLM-v2）。
长上下文（RULER 基准）： SPLA 显著优于其他稀疏方法。
- 在 256k 上下文长度下，SPLA 达到了 72.3 的得分，相比之下，稠密基准为 69.3，NSA 为 32.5，InfLLM-v2 为 42.6。
- SPLA 在其他稀疏模型因长尾发散而导致性能严重下降的情况下，展现出了鲁棒性。
推理： SPLA 在复杂推理任务（AIME、HMMT、LiveCodeBench）上优于 NSA 和 InfLLM-v2，验证了其原则性选择和混合设计能够保留必要的细粒度信息。
消融实验：
- SPLA vs. SPA（不含 RLA 的稀疏加注意力）： SPLA 始终优于 SPA，证实了残差线性注意力模块在防止发散方面的关键作用。
- SPA vs. InfLLM-v2： SPA（使用基于泰勒的选择但没有 RLA）通常优于 InfLLM-v2，验证了二阶选择度量优于一阶近似。

5. 重要性与主张

论文声称，SPLA 通过解决效率与上下文保留之间的权衡，为下一代基础模型的服务提供了一条可扩展的路径。

弥合差距： SPLA 有效地缩小了稀疏注意力与稠密注意力之间的性能差距，在长上下文基准（高达 256k token）上甚至超越了稠密模型，同时保留了稀疏解码的效率。
适配策略： 该工作提出了一种极具前景的方向，即可以在昂贵的预训练阶段维持标准的注意力架构，然后在持续预训练期间将其无缝适配为高效的稀疏变体，且仅需极少的额外参数。
鲁棒性： 通过避免“硬截断”并利用原则性的选择度量，SPLA 缓解了随着上下文长度扩展而困扰现有稀疏方法的累积上下文损失问题。

作者指出，虽然由于块选择和在特定 14B 配置下的较小 GQA 组大小，目前的训练效率与高度优化的稠密算子相比略有开销，但该方法对于大规模训练是可行的，且预计随着更大组大小的使用，效率增益将会提升。