Kostenoptimale gruppierte Abfrage-Aufmerksamkeit für LLMs mit langem Kontext

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

March 12, 2025

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

cs.AI

Zusammenfassung

Die Entwicklung effektiver und effizienter Transformer-basierter großer Sprachmodelle (LLMs) hat sich in letzter Zeit zu einem Forschungsschwerpunkt entwickelt, bei dem es darum geht, die Sprachfähigkeiten der Modelle zu maximieren und gleichzeitig die Trainings- und Bereitstellungskosten zu minimieren. Bisherige Bemühungen haben hauptsächlich komplexe Zusammenhänge zwischen Modellleistung, Parametergröße und Datenmenge beschrieben sowie nach der optimalen Rechenressourcenverteilung für das Training von LLMs gesucht. Dabei wurden jedoch die Auswirkungen der Kontextlänge und der Konfiguration der Aufmerksamkeitsköpfe (die Anzahl der Query- und Key-Value-Köpfe in der gruppierten Query-Aufmerksamkeit) auf Training und Inferenz übersehen. In diesem Artikel vergleichen wir systematisch Modelle mit unterschiedlichen Parametergrößen, Kontextlängen und Aufmerksamkeitskopf-Konfigurationen hinsichtlich Modellleistung, Rechenkosten und Speicherbedarf. Anschließend erweitern wir die bestehenden Skalierungsmethoden, die ausschließlich auf Parametergröße und Trainingsrechenleistung basieren, um die Konstruktion kosteneffizienter LLMs sowohl während des Trainings als auch der Inferenz zu leiten. Unsere quantitativen Skalierungsstudien zeigen, dass bei der Verarbeitung ausreichend langer Sequenzen ein größeres Modell mit weniger Aufmerksamkeitsköpfen einen geringeren Verlust erzielen kann, während gleichzeitig die Rechen- und Speicherkosten reduziert werden. Unsere Erkenntnisse bieten wertvolle Einblicke für die Entwicklung praktischer LLMs, insbesondere in Szenarien mit langen Kontexten. Wir werden unseren Code und unsere Daten öffentlich zugänglich machen.

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.