長文脈LLMのためのコスト最適化グループ化クエリアテンション

Cost-Optimal Grouped-Query Attention for Long-Context LLMs

March 12, 2025

著者: Yingfa Chen, Yutong Wu, Xu Han, Zhiyuan Liu, Maosong Sun

cs.AI

要旨

Transformerベースの大規模言語モデル（LLM）を効果的かつ効率的に構築することが最近の研究焦点となっており、モデルの言語能力を最大化しつつ、学習とデプロイのコストを最小化することが求められています。既存の研究では、主にモデルの性能、パラメータサイズ、データサイズの間の複雑な関係を説明し、LLMを学習するための最適な計算リソースの割り当てを探求してきました。しかし、これらの研究は、コンテキスト長とアテンションヘッドの構成（グループ化クエリアテンションにおけるクエリとキー・バリューヘッドの数）が学習と推論に与える影響を見落としています。本論文では、異なるパラメータサイズ、コンテキスト長、アテンションヘッド構成を持つモデルを、モデル性能、計算コスト、メモリコストの観点から系統的に比較します。そして、パラメータサイズと学習計算量のみに基づいた既存のスケーリング手法を拡張し、学習と推論の両方においてコスト最適なLLMを構築するための指針を提供します。私たちの定量的なスケーリング研究は、十分に長いシーケンスを処理する場合、より少ないアテンションヘッドを持つ大規模モデルが、より低い損失を達成しつつ、計算コストとメモリコストを低減できることを示しています。これらの発見は、特に長いコンテキストを処理するシナリオにおいて、実用的なLLMを開発するための貴重な洞察を提供します。私たちはコードとデータを公開する予定です。

English

Building effective and efficient Transformer-based large language models (LLMs) has recently become a research focus, requiring maximizing model language capabilities and minimizing training and deployment costs. Existing efforts have primarily described complex relationships among model performance, parameter size, and data size, as well as searched for the optimal compute allocation to train LLMs. However, they overlook the impacts of context length and attention head configuration (the number of query and key-value heads in grouped-query attention) on training and inference. In this paper, we systematically compare models with different parameter sizes, context lengths, and attention head configurations in terms of model performance, computational cost, and memory cost. Then, we extend the existing scaling methods, which are based solely on parameter size and training compute, to guide the construction of cost-optimal LLMs during both training and inference. Our quantitative scaling studies show that, when processing sufficiently long sequences, a larger model with fewer attention heads can achieve a lower loss while incurring lower computational and memory costs. Our findings provide valuable insights for developing practical LLMs, especially in long-context processing scenarios. We will publicly release our code and data.