通过分层复用注意力,让大模型长文本推理快得多且不丢精度。
HyLRA: Hybrid Layer Reuse Attention for Efficient Long-Context Inference
- 按层分析注意力敏感度,区分必须全算和可复用的层
- 在关键层保留完整计算,其他层复用前层最相关索引
- 比现有方法快6%~46%,精度损失小于1%,适合长文本任务
大语言模型在长上下文推理中受限于注意力机制的二次计算复杂度及键值缓存的巨大内存开销。现有稀疏注意力方法虽试图利用内在稀疏性缓解问题,但多依赖固定模式或激进剪枝,难以在效率与准确率间取得平衡。本文提出 {f HyLRA}(混合层复用注意力),基于分层稀疏性分析发现:注意力机制具有双重特性—— extit{层内敏感性}(特定层需全注意力防止特征失真)与 extit{层间相似性}(连续层共享大量关键标记)。基于此,HyLRA 采用离线动态规划生成最优分层策略:对敏感层保留全注意力以保障鲁棒性,对容错层则直接复用前一层的 top-$k$ 索引,跳过二次计算。该方法将计算聚焦于最关键的标记,有效突破密集注意力的二次瓶颈。大量实验表明,HyLRA 在保持相近性能(<1% 准确率下降)的前提下,推理吞吐提升 6%–46%,持续优于当前最优稀疏注意力方法。代码已开源。
原文摘要 · Abstract (English)
Long-context inference in Large Language Models (LLMs) is bottlenecked by the quadratic computation complexity of attention and the substantial memory footprint of Key-Value (KV) caches. While existing sparse attention mechanisms attempt to mitigate this by exploiting inherent sparsity, they often rely on rigid patterns or aggressive pruning, failing to achieve an optimal balance between efficiency and accuracy. In this paper, we introduce {\bf HyLRA} ({\bf Hy}brid {\bf L}ayer {\bf R}euse {\bf A}ttention), a novel framework driven by layer-wise sparsity profiling. Our empirical analysis uncovers a dual characteristic in attention mechanics: \textit{intra-layer sensitivity}, where specific layers necessitate full attention to prevent feature distortion, and \textit{inter-layer similarity}, where consecutive layers share substantial critical tokens. Based on these observations, HyLRA employs an offline dynamic programming approach to derive an optimal layer-wise policy. This hybrid strategy retains full attention for sensitive layers to ensure robustness, while enabling tolerant layers to bypass quadratic calculations by directly reusing top-$k$ indices from preceding layers. This approach allows LLMs to restrict computation to the most critical tokens, effectively overcoming the quadratic bottleneck of dense attention. Extensive evaluations demonstrate that HyLRA improves inference throughput by 6\%--46\% while maintaining comparable performance (with $<1\%$ accuracy degradation), consistently outperforming state-of-the-art sparse attention methods. HyLRA is open source at \href{https://anonymous.4open.science/r/unified-cache-management-CF80/}{\texttt{/r/unified-cache-management-CF80/}}
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。