面向长上下文建模的成本最优分组查询注意力
计算与语言
2025-09-29 v3 人工智能
机器学习
摘要
分组查询注意力（GQA）是降低大语言模型（LLM）中注意力层计算成本的广泛采用策略。然而，当前GQA配置往往次优，因为它们忽略了上下文长度如何影响推理成本。由于推理成本随上下文长度增长，最具成本效率的GQA配置也应相应变化。在本工作中，我们分析了上下文长度、模型大小、GQA配置与模型损失之间的关系，并引入两项创新：（1）将总注意力头数与隐藏层维度解耦，实现对注意力浮点运算次数的更灵活控制；（2）联合优化模型大小与GQA配置，以在注意力层与其他组件之间实现更好的推理资源分配。我们的分析表明，常用的GQA配置在长上下文场景下高度次优。更重要的是，我们提出了一种推导成本最优GQA配置的方法。我们的结果表明，对于长上下文场景，应使用更少的注意力头同时扩大模型规模。通过我们方法选取的配置与Llama-3的GQA相比，可同时减少超过50%的内存使用和浮点运算次数，且模型能力无退化。我们的发现为设计高效的长上下文LLM提供了有价值的见解。代码可在 https://www.github.com/THUNLP/cost-optimal-gqa 获取。
引用
@article{arxiv.2503.09579,
title = {Cost-Optimal Grouped-Query Attention for Long-Context Modeling},
author = {Yingfa Chen and Yutong Wu and Chenyang Song and Zhen Leng Thai and Xingyu Shen and Xu Han and Zhiyuan Liu and Maosong Sun},
journal= {arXiv preprint arXiv:2503.09579},
year = {2025}
}
备注
EMNLP 2025 Main