Kostenoptimale gegroepeerde-query-attentie voor lange-context LLM's

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

Samenvatting

Het ontwikkelen van effectieve en efficiënte Transformer-gebaseerde grote taalmodelen (LLM's) is recentelijk een onderzoeksfocus geworden, waarbij het maximaliseren van de taalvaardigheden van het model en het minimaliseren van de trainings- en implementatiekosten centraal staan. Bestaande inspanningen hebben voornamelijk complexe relaties beschreven tussen modelprestaties, parametergrootte en datagrootte, en hebben gezocht naar de optimale compute-toewijzing voor het trainen van LLM's. Ze hebben echter de invloed van contextlengte en de configuratie van aandachtskoppen (het aantal query- en key-value koppen in grouped-query attention) op training en inferentie over het hoofd gezien. In dit artikel vergelijken we systematisch modellen met verschillende parametergroottes, contextlengtes en aandachtskopconfiguraties wat betreft modelprestaties, rekenkosten en geheugenkosten. Vervolgens breiden we de bestaande schaalingsmethoden, die uitsluitend gebaseerd zijn op parametergrootte en trainingscompute, uit om de constructie van kost-optimale LLM's tijdens zowel training als inferentie te begeleiden. Onze kwantitatieve schaalingsstudies tonen aan dat, bij het verwerken van voldoende lange sequenties, een groter model met minder aandachtskoppen een lager verlies kan bereiken terwijl het lagere reken- en geheugenkosten met zich meebrengt. Onze bevindingen bieden waardevolle inzichten voor de ontwikkeling van praktische LLM's, vooral in scenario's met lange-contextverwerking. We zullen onze code en data openbaar beschikbaar stellen.

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.