Attenzione a Query Raggruppate Ottimizzata in Termini di Costo per LLM con Contesti Lunghi

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

March 12, 2025

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

cs.AI

Abstract

La costruzione di modelli linguistici di grandi dimensioni (LLM) basati su Transformer efficaci ed efficienti è recentemente diventata un focus di ricerca, richiedendo la massimizzazione delle capacità linguistiche del modello e la minimizzazione dei costi di addestramento e distribuzione. Gli sforzi esistenti hanno principalmente descritto relazioni complesse tra le prestazioni del modello, la dimensione dei parametri e la dimensione dei dati, oltre a cercare l'allocazione ottimale delle risorse computazionali per addestrare gli LLM. Tuttavia, trascurano gli impatti della lunghezza del contesto e della configurazione delle teste di attenzione (il numero di teste di query e key-value nell'attenzione a query raggruppate) sull'addestramento e sull'inferenza. In questo articolo, confrontiamo sistematicamente modelli con diverse dimensioni dei parametri, lunghezze del contesto e configurazioni delle teste di attenzione in termini di prestazioni del modello, costi computazionali e costi di memoria. Successivamente, estendiamo i metodi di scalatura esistenti, basati esclusivamente sulla dimensione dei parametri e sulle risorse computazionali di addestramento, per guidare la costruzione di LLM ottimali in termini di costi sia durante l'addestramento che durante l'inferenza. I nostri studi quantitativi di scalatura dimostrano che, quando si elaborano sequenze sufficientemente lunghe, un modello più grande con meno teste di attenzione può raggiungere una perdita inferiore pur comportando costi computazionali e di memoria più bassi. Le nostre scoperte forniscono preziose intuizioni per lo sviluppo di LLM pratici, specialmente in scenari di elaborazione con contesti lunghi. Rilasceremo pubblicamente il nostro codice e i nostri dati.

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.