Atenção de Consultas Agrupadas com Custo Ótimo para LLMs de Contexto Longo

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

March 12, 2025

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

cs.AI

Resumo

A construção de modelos de linguagem grandes (LLMs) baseados em Transformers eficazes e eficientes tornou-se recentemente um foco de pesquisa, exigindo a maximização das capacidades linguísticas do modelo e a minimização dos custos de treinamento e implantação. Os esforços existentes descreveram principalmente relações complexas entre o desempenho do modelo, o tamanho dos parâmetros e o tamanho dos dados, além de buscarem a alocação ideal de computação para treinar LLMs. No entanto, eles negligenciam os impactos do comprimento do contexto e da configuração das cabeças de atenção (o número de cabeças de consulta e chave-valor na atenção de consulta agrupada) no treinamento e na inferência. Neste artigo, comparamos sistematicamente modelos com diferentes tamanhos de parâmetros, comprimentos de contexto e configurações de cabeças de atenção em termos de desempenho do modelo, custo computacional e custo de memória. Em seguida, estendemos os métodos de escalonamento existentes, que se baseiam apenas no tamanho dos parâmetros e na computação de treinamento, para orientar a construção de LLMs com custo ideal durante o treinamento e a inferência. Nossos estudos quantitativos de escalonamento mostram que, ao processar sequências suficientemente longas, um modelo maior com menos cabeças de atenção pode alcançar uma perda menor enquanto incorre em custos computacionais e de memória mais baixos. Nossas descobertas fornecem insights valiosos para o desenvolvimento de LLMs práticos, especialmente em cenários de processamento de contexto longo. Disponibilizaremos publicamente nosso código e dados.

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.