Block Sparse Attention with Log-Linear Complexity
Bohao Tang, Zhen Qin, Yuqi Pan, Zheng Li, Pengfei Liu
cs.LG
2026-09-25
交大与ByteDance Seed的PISA用粗到细金字塔和LogSumExp打分选关键块,选择复杂度从平方降到O(N log N)。2.67B上常识题与BSA持平,检索更好,256K选块比BSA快约10倍。
长上下文的二次注意力太贵。块稀疏注意力把 key/value 切成固定块,每条 query 只看 Top-K 个块,真正算注意力可以线性。但选哪些块,传统做法要对每个 query 打一遍所有块的分,序列长 N、块大小 C,选择仍是 O(N²/C)。选块阶段把线性注意力的红利吃掉了。
NSA、MoBA、HiLS 这类可训练稀疏方法,prefill 仍是平方级。训练无关的 HiP 能做到 O(N log N),但不进训练。PISA(Pyramid Sparse Attention)要的是可训练、选择也是 log-linear 的块稀疏注意力。作者来自上海交大、上海创智学院和 ByteDance Seed。
先把 key 做成金字塔。最细一层是大小 C 的叶子块,往上每层把相邻块均值池化,分支因子叶子层为 C、更粗层为 2,一共 O(log N) 层。选择从最粗层开始:该层候选很少,对每个候选用 LogSumExp 打分,留下 Top-K,再展开成下一层的子块。没选中的整枝丢掉。每层候选数量有上界(约 gK),所以单条 query 的路由是 O(log N),整段序列 O(N log N)。decode 复杂度相应是 O(log N)。
为什么用 LSE 而不是块均值点积:LSE 近似块里最热的那些 key,对检索这种峰值信号更友好。消融里 PISA-1 只用均值、PISA-2 用均值加方差,完整 LSE 在三个尺度上训练损失和 containment 都更好。叶子块还会强制保留当前块、前一块和序列首块,避免因果局部窗口被金字塔剪掉。
实现上,训练/prefill 用两段 Triton kernel:中间层路由和叶子 LSE 打分融合,不物化 dense 分值矩阵;decode 用单段 kernel 减少 launch。同一 GQA 组共享选中的 KV 块。Qtile=4 时最多 4 条 query 复用同一片叶子 key。
从零预训练 100B token、上下文 4K,再 10B token 续训到 16K。稀疏方法块大小 C=64,预训练 K=8、续训 K=32。对照是 Full Attention、BSA(NSA 的选中支路加均值摘要)、NSA、HiLS。规模 418M / 1.47B / 2.67B。
2.67B 预训练后:
| 方法 | 训练损失 | 常识 8 项平均 | 检索 6 项平均 |
| Full Attention | 2.2209 | 58.23 | 53.01 |
| BSA | 2.2184 | 58.61 | 51.13 |
| PISA | 2.2151 | 58.10 | 52.14 |
常识题上稀疏方法彼此咬得很紧,PISA 并不领先。检索平均是稀疏里最好,仍低于满注意力。418M 和 1.47B 是同一趋势:PISA containment 41.77 / 49.83,BSA 是 38.18 / 48.54。RULER 续训后 2.67B 的四类 needle、1K 到 16K 平均:满注意力 66.24,PISA 62.80,PISA-2 62.92,BSA 只有 54.99。选块诊断固定满注意力的 Q/K、K=8:PISA Recall@8 为 90.95,BSA 85.91,捕获注意力质量比 99.46 vs 98.42。
延迟才是卖点。选块 kernel、C=64、K=8:4K 到 16K 上 BSA 更快;64K / 128K / 256K 上 PISA 相对 BSA 分别快 2.86×、5.31×、9.95×。短序列别用这个金字塔。
块稀疏要进训练,选块复杂度必须一起降,否则 prefill 仍被平方项卡住。PISA 证明层次路由加 LSE 可以在不牺牲太多语言建模的前提下做到 O(N log N),并且检索比均值打分更稳。Triton 实现是可跟的工程资产。
对已经在 4K–8K 推理的人,这篇几乎没有立即收益。对要做 64K 以上 prefill 的稀疏注意力,选块阶段值得换成金字塔。它是 NSA/BSA 路线的复杂度补丁,不是新的注意力原语。
作者承认实验只覆盖到 2.67B 和 100B+10B token,规模一变,相对收益可能翻脸。RULER 上 PISA-2 平均 62.92,略高于完整 PISA 的 62.80,「LSE 一定更好」在长上下文检索上并不干净。续训后 containment 平均 PISA 还低于 BSA(附录 2.67B:66.40 vs 68.15),主文「检索更好」主要指 4K 预训练表。延迟数字排除了选中块上的注意力计算,端到端能省多少要另测。没有和 HiP 这类训练无关 O(N log N) 方法做同设定质量对比。强制保留三块会抬高 Recall@8,诊断数字含有这部分保底重叠。