用轻量 sketch 预判注意力,动态压缩计算量,提速六倍且几乎不丢精度。
Scout Before You Attend: Sketch-and-Walk Sparse Attention for Efficient LLM Inference
- 通过哈达玛随机投影快速估算注意力得分,实现训练无关的稀疏化。
- 在20%注意力密度下保持近似无损精度,部分场景优于全连接注意力。
- 统一适用于预填充与解码阶段,适合追求高效推理的部署场景。
自注意力机制在长上下文大模型推理中占据主要计算与内存开销,贯穿预填充与解码阶段。为应对这一挑战,本文提出 Sketch&Walk Attention,一种无需训练的稀疏注意力方法。该方法利用哈达玛随机投影生成注意力得分的低成本近似值,并通过确定性“行走”机制在层间聚合这些估计,捕捉令牌间非直接交互的注意力影响。累积的行走得分用于选择 top-k 注意力块,实现动态稀疏化,且仅需单一训练无关算法,同时适配预填充与解码阶段,并配合定制化稀疏注意力核。在多种模型与任务上,Sketch&Walk 在 20% 注意力密度下保持近乎无损精度,某些场景甚至略超密集注意力,同时实现最高 6 倍推理加速。
原文摘要 · Abstract (English)
Self-attention dominates the computational and memory cost of long-context LLM inference across both prefill and decode phases. To address this challenge, we introduce Sketch&Walk Attention, a training-free sparse attention method that determines sparsity with lightweight sketches and deterministic walk. Sketch&Walk applies Hadamard sketching to get inexpensive approximations of attention scores, then aggregates these estimates across layers via a walk mechanism that captures attention influence beyond direct interactions between tokens. The accumulated walk scores are used to select top-k attention blocks, enabling dynamic sparsity with a single training-free algorithm that applies uniformly to both the prefill and decode phases, together with custom sparse attention kernels. Across a wide range of models and tasks, Sketch&Walk maintains near-lossless accuracy at 20% attention density and can slightly outperform dense attention in some settings, while achieving up to 6x inference speedup.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。