Attention Groupée par Requêtes à Coût Optimal pour les LLM à Contexte Long

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

March 12, 2025

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

cs.AI

Résumé

La construction de modèles de langage de grande taille (LLMs) basés sur les Transformers, à la fois efficaces et performants, est récemment devenue un axe de recherche majeur, nécessitant de maximiser les capacités linguistiques des modèles tout en minimisant les coûts d'entraînement et de déploiement. Les travaux existants ont principalement décrit les relations complexes entre la performance du modèle, la taille des paramètres et la quantité de données, ainsi que recherché l'allocation optimale de calcul pour entraîner les LLMs. Cependant, ils négligent les impacts de la longueur du contexte et de la configuration des têtes d'attention (le nombre de têtes de requête et de clé-valeur dans l'attention par requêtes groupées) sur l'entraînement et l'inférence. Dans cet article, nous comparons systématiquement des modèles avec différentes tailles de paramètres, longueurs de contexte et configurations de têtes d'attention en termes de performance du modèle, de coût computationnel et de coût mémoire. Ensuite, nous étendons les méthodes de mise à l'échelle existantes, qui se basent uniquement sur la taille des paramètres et le calcul d'entraînement, pour guider la construction de LLMs optimaux en coût lors de l'entraînement et de l'inférence. Nos études quantitatives de mise à l'échelle montrent que, lors du traitement de séquences suffisamment longues, un modèle plus grand avec moins de têtes d'attention peut atteindre une perte plus faible tout en engendrant des coûts computationnels et mémoire réduits. Nos résultats fournissent des insights précieux pour le développement de LLMs pratiques, en particulier dans des scénarios de traitement de contextes longs. Nous rendrons publiquement disponibles notre code et nos données.

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.