注意力跨度混合：利用异构滑动窗口长度优化大语言模型推理效率
机器学习
2025-11-26 v3 人工智能
计算与语言
摘要
滑动窗口注意力为大型语言模型（LLM）在长上下文场景中的内存和吞吐量挑战提供了一种硬件高效的解决方案。现有方法通常在所有注意力头和输入规模上采用单一的窗口长度。然而，这种统一的方法未能捕捉LLM固有的异构注意力模式，忽略了它们各自不同的精度-延迟权衡。为应对这一挑战，我们提出了*注意力跨度混合*（MoA），该方法自动为不同的头和层定制不同的滑动窗口长度配置。MoA构建并导航一个包含各种窗口长度及其相对于输入规模的缩放规则的搜索空间。它对模型进行剖析，评估潜在配置，并为每个头确定最优的长度配置。MoA适应不同的输入规模，揭示出一些注意力头会扩展其关注范围以适应更长的输入，而其他头则持续专注于固定长度的局部上下文。实验表明，在平均滑动窗口长度相同的情况下，MoA将有效上下文长度提高了3.9倍，在Vicuna-{7B, 13B}和Llama3-{8B, 70B}模型上，检索准确率比统一窗口基线提升了1.5-7.1倍。此外，MoA缩小了与全注意力的性能差距，在三个长上下文理解基准测试中，将最大相对性能下降从9%-36%降至5%以内。MoA实现了1.2-1.4倍的GPU内存缩减，与FlashAttention2和vLLM相比，解码吞吐量分别提升了6.6-8.2倍和1.7-1.9倍，且性能影响极小。我们的代码可在 https://github.com/thu-nics/MoA 获取。
引用
@article{arxiv.2406.14909,
title = {Mixture of Attention Spans: Optimizing LLM Inference Efficiency with Heterogeneous Sliding-Window Lengths},
author = {Tianyu Fu and Haofeng Huang and Xuefei Ning and Genghan Zhang and Boju Chen and Tianqi Wu and Hongyi Wang and Zixiao Huang and Shiyao Li and Shengen Yan and Guohao Dai and Huazhong Yang and Yu Wang},
journal= {arXiv preprint arXiv:2406.14909},
year = {2025}
}
备注
Published at CoLM'25