论文
CRISP:具有结构质量驱动路由的悬崖感知输入自适应稀疏预填充
CRISP: Cliff-awaRe Input-adaptive Sparse Prefilling with Structural-Mass-Motivated Routing
摘要
长上下文 LLM 推理的注意力预填充阶段呈二次方扩展,使自注意力成为严重的计算瓶颈。传统的稀疏注意力方法通过固定模式或离线分析来缓解这一问题,但缺乏适应依赖于输入的注意力结构的灵活性。最近的动态方法通过实时将头路由到稀疏模式来解决这个问题,但依赖于间接路由代理,其开销和预算分配机制忽略了 post-softmax 质量层次结构。我们提出了 CRISP(悬崖感知输入自适应稀疏预填充),它识别并解决了这种动态路由范例中的两个结构性挑战。首先,我们表明可以直接从代理注意图的结构中读取路由决策。我们用 C_struct 替换 Jensen-Shannon Divergence (JSD) 路由,C_struct 是一种结构代理,可以测量 Vertical-Slash 兼容位置的质量并重现 JSD 的路由决策,同时消除池化 matmul 和随后的 KL 发散开销。其次,我们形式化了 post-softmax 质量悬崖,并从理论上证明了严格累积覆盖阈值在长上下文中会累积 O(n) 背景噪声。 CRISP 通过基于本底噪声的接收器感知阈值来对此进行导航。根据经验,在两个模型系列的 InfiniteBench、RULER 和 LongBench 中,CRISP 是整体上最强的稀疏方法,并且匹配或超过了检索密集基准上的精确密集注意力,在基线上恢复了高达 +28.0 pp 的检索任务,并在 512k 标记下实现了高达 5.30 倍的注意力加速,这主要是由我们在选择过程中的 O(n) 噪声消除同时保持结构完整性驱动的。