提出动态调整注意力头数的优化方法,显著降低长文本建模的计算开销。
Cost-Optimal Grouped-Query Attention for Long-Context Modeling
- 解耦注意力头大小与隐藏层维度,灵活调控计算量。
- 针对不同上下文长度推荐更优的模型规模与注意力配置组合。
- 在不损失模型能力前提下,内存和算力消耗减少超50%。
分组查询注意力(GQA)是降低大语言模型注意力层计算成本的常用策略。然而,现有GQA配置常因忽略上下文长度对推理成本的影响而次优。由于推理成本随上下文长度增加,最经济的GQA配置也应随之变化。本文分析了上下文长度、模型规模、GQA配置与模型损失之间的关系,提出两项创新:(1) 将总注意力头大小与隐藏层维度解耦,实现对注意力浮点运算量更灵活的控制;(2) 联合优化模型规模与GQA配置,实现注意力层与其他组件间推理资源的更优分配。分析表明,当前常用GQA配置在长上下文场景中严重次优。更重要的是,本文提出一套推导成本最优GQA配置的方法:在长上下文场景下,应减少注意力头数并扩大模型规模。采用该方法的配置相比Llama-3的GQA,在无性能下降的前提下,可使内存使用和浮点运算量均降低超过50%。研究为高效长上下文大模型设计提供了关键洞见。代码已开源。
原文摘要 · Abstract (English)
Grouped-Query Attention (GQA) is a widely adopted strategy for reducing the computational cost of attention layers in large language models (LLMs). However, current GQA configurations are often suboptimal because they overlook how context length influences inference cost. Since inference cost grows with context length, the most cost-efficient GQA configuration should also vary accordingly. In this work, we analyze the relationship among context length, model size, GQA configuration, and model loss, and introduce two innovations: (1) we decouple the total head size from the hidden size, enabling more flexible control over attention FLOPs; and (2) we jointly optimize the model size and the GQA configuration to arrive at a better allocation of inference resources between attention layers and other components. Our analysis reveals that commonly used GQA configurations are highly suboptimal for long-context scenarios. More importantly, we propose a recipe for deriving cost-optimal GQA configurations. Our results show that for long-context scenarios, one should use fewer attention heads while scaling up model size. Configurations selected by our recipe can reduce both memory usage and FLOPs by more than 50% compared to Llama-3's GQA, with *no degradation in model capabilities*. Our findings offer valuable insights for designing efficient long-context LLMs. The code is available at https://www.github.com/THUNLP/cost-optimal-gqa .
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。