S2O：基于在线置换的稀疏注意力早停
机器学习
2026-05-06 v2 人工智能
摘要
注意力随序列长度呈二次增长，根本限制了长上下文推理。现有的块级稀疏化可以降低延迟，但粗粒度块会施加内在稀疏性上限，使得即使经过仔细设计的系统也难以进一步提升。我们提出 S2O，通过在线置换实现稀疏注意力的早停。灵感来自内存系统中的虚拟地址到物理地址映射，S2O 重新审视并分解 FlashAttention 的执行，使推理能够加载原顺序中非连续的标记。受注意力热图中细粒度结构的启发，我们将显式置换转化为一种在线、索引导导的离散加载策略；在极轻的预处理和索引置换开销下，它将重要性集中于小集合的高优先级块上。基于这种重要性导导的在线置换加载，S2O 进一步引入一种早停规则：计算从高到低重要性进行；一旦当前块得分低于阈值，S2O 即提前终止并跳过剩余的低贡献块，从而在受控误差预算下提高有效稀疏性并降低计算量。结果显示，S2O 在 128K 上下文下，Llama-3.1-8B 的单运算 MSE 提升最高可达 3.82 倍，在匹配稀疏性下的前置计算密度降低 3.31 倍；同时保持端到端准确率，实现 7.51 倍注意力和 3.81 倍端到端加速。
关键词
引用
@article{arxiv.2602.22575,
title = {S2O: Early Stopping for Sparse Attention via Online Permutation},
author = {Yu Zhang and Songwei Liu and Chenqian Yan and Sheng Lin and Beichen Ning and Fangmin Chen and Xing Wang},
journal= {arXiv preprint arXiv:2602.22575},
year = {2026}
}