Twilight: Adaptive Attention Sparsity with Hierarchical Top-p Pruning¶

会议: NeurIPS 2025
arXiv: 2502.02770
代码: https://github.com/tsinghua-ideal/Twilight
领域: 模型压缩
关键词: 注意力稀疏, 自适应预算, Top-p采样, 长上下文加速, KV缓存压缩

一句话总结¶

提出 Twilight，借鉴 top-p 采样（nucleus sampling）的思想替代固定预算 top-k 做注意力稀疏——动态选择注意力权重累积达 p% 的最少 Token，自适应不同注意力头的分布特征，在保持精度的同时比 SOTA 稀疏注意力再提速 1.4x。

研究背景与动机¶

领域现状：利用注意力稀疏性加速长上下文 LLM 推理是热门方向。H2O、StreamingLLM、Quest、Double Sparsity 等方法选择"关键 Token"子集计算注意力

现有痛点：所有现有方法使用固定预算 top-k——保留 k 个注意力权重最高的 Token。但最优 k 因层、注意力头、序列位置和输入内容而异： - 聚焦注意力（focused）：权重集中在少数 Token，固定 k 导致过选——选了太多不重要的 Token - 分散注意力（diffuse）：权重均匀分布，固定 k 导致欠选——遗漏了重要 Token - 不同算法需不同程度的"过选冗余"来补偿估计误差

核心矛盾：固定预算无法适应注意力分布的动态变化

切入角度：这个问题与 LLM 采样中 top-k vs top-p 的困境完全类比——top-k 采样也因无法适应不同的词分布而被 nucleus sampling 替代

核心 idea：将 top-p 引入注意力稀疏——不固定保留 k 个 Token，而是保留注意力权重累积到 p 的最少 Token 数量，自动适应每个头的分布特征

方法详解¶

整体框架¶

Twilight 采用"选择-剪枝"（Select-then-Prune）两阶段架构：(1) Token Selector：任何现有稀疏注意力算法，用保守的较大预算（如 1/4 稀疏度）做初步选择；(2) Twilight Pruner：在选择结果上做 top-p 剪枝，只保留注意力权重累积到 p 的最少 Token 子集。最终的稀疏注意力只在 top-p Token 上计算。

关键设计¶

Oracle Top-p 稀疏注意力的形式化：
- 功能：给定阈值 p，选择最小 Token 集合使注意力权重之和 ≥ p
- 形式化：\(\mathcal{I} = \arg\min_\mathcal{I} |\mathcal{I}|\) s.t. \(\sum_{i \in \mathcal{I}} \mathbf{W}[i] \geq p\)
- 理论保证：输出误差上界为 \((1-p) \cdot \|\mathbf{V}\|_F\)——与 Token 数量无关，只与 p 相关
- vs top-k：top-k 的误差取决于分布——聚焦分布误差小、分散分布误差大。top-p 自动适应
高效 Top-p 实现（二分搜索）：
- 功能：在 GPU 上并行高效实现 top-p 选择
- 核心思路：不排序（排序不适合 GPU 并行），而是用二分搜索找到阈值 m，使 \(\sum_{W[i] \geq m} W[i] \geq p\)。在 FlashInfer 的 top-p 采样内核上修改实现
- 设计动机：避免 \(O(n \log n)\) 排序开销，二分搜索只需 \(O(\log (\max W / \epsilon))\) 轮，每轮 \(O(n)\) 并行操作
4-bit K 缓存量化：
- 功能：用 INT4 量化的 K 缓存快速估计注意力权重
- 核心思路：top-p 需要注意力权重的数值精度（不只是排序正确性），实验证明 4-bit 是精度和效率的最优平衡——2-bit 权重和下降严重，4-bit 和 8-bit 几乎相同
- 额外内存：1/8 的 KV 缓存开销（16-bit → 4-bit，只量化 K 的一半）
负载均衡（Head 动态性感知）：
- 功能：处理不同头预算不同导致的计算负载不平衡
- 核心思路：复用 FlashInfer 的负载均衡算法，将头维度展平后重新分配计算资源
- 设计动机：某些头可能只需 2% Token，其他需 50%——传统实现给所有头相同计算资源导致浪费

损失函数 / 训练策略¶

完全无训练——Twilight 是推理时的加速框架
p 是唯一超参数，通常设为 0.8-0.9
可增强任何现有的 top-k 稀疏注意力算法

实验关键数据¶

准确率评估¶

基准	方法	Base Budget	有效 Budget (top-p后)	性能保留
PG-19 (PPL)	Quest	1024	~100 (p=0.85)	~99.5%
RULER (128K)	DS	1/4	~1/64	~98%
LongBench	H2O+Twilight	1/4	自适应	优于固定 H2O

效率评估（端到端解码延迟）¶

方法	单层注意力加速	端到端加速	最大加速
Full Attention	1x	1x	-
Quest (top-k)	~8x	~3x	-
Quest + Twilight (top-p)	~11x	~4x	15.8x
相对 SOTA 稀疏注意力	1.4x	1.35x	-

消融实验¶

配置	精度	效率	说明
p=0.7	略有下降	最高	过于激进
p=0.85	几乎无损	高	推荐设置
p=0.95	最好	中等	保守
2-bit K cache	下降	-	精度不足
4-bit K cache	几乎无损	好	最优平衡

关键发现¶

同一模型中不同注意力头的预算差异可达 50 倍（2% vs ~100% Token）——固定预算注定浪费或不足
top-p 在正确性上几乎与 oracle top-k（已知完美注意力权重的 top-k）持平，但自适应性远超人工调参的 top-k
Twilight 对所有测试的基础算法（Quest、DS、H2O）都带来一致提速——即插即用增强
长上下文场景增益更大——因为长序列中注意力分布变化更剧烈

亮点与洞察¶

top-p → 注意力剪枝的类比是本文最深刻的洞察——两个看似不相关的问题（采样和注意力选择）共享相同的数学结构：从概率分布中选择累积质量达阈值的最小集合
"增强器而非替代品"的定位非常聪明——不与现有方法竞争，而是在它们之上叠加自适应性。这使得 Twilight 可以跟随该领域的所有未来进展
理论误差界 \((1-p) \cdot \|\mathbf{V}\|_F\) 简洁且与 Token 数量无关——为注意力近似提供了一个通用的精度-效率权衡旋钮

局限与展望¶

top-p 需要（近似的）注意力权重而非仅排序——比 top-k 对权重估计精度要求更高
额外的 INT4 K 缓存增加 1/8 内存开销（某些内存受限场景可能不可接受）
二分搜索实现中 \(\epsilon\) 参数需要选择——太大可能导致实际选择不满足 p 约束
未探索 prefill 阶段的 top-p 稀疏——当前主要关注解码阶段

评分¶

新颖性: ⭐⭐⭐⭐⭐ top-p 到注意力稀疏的概念迁移深刻且自然
实验充分度: ⭐⭐⭐⭐⭐ 长/中上下文准确率 + 延迟/吞吐量效率 + 多基础算法 + 详细消融
写作质量: ⭐⭐⭐⭐⭐ 问题形式化+类比推理+系统实现都很清晰
价值: ⭐⭐⭐⭐⭐ 解决了固定预算的根本问题，即插即用增强所有稀疏注意力方法