Atención de Consultas Agrupadas Óptima en Costo para LLMs de Contexto Largo

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

Resumen

La construcción de modelos de lenguaje grandes (LLMs) basados en Transformers eficaces y eficientes se ha convertido recientemente en un foco de investigación, requiriendo maximizar las capacidades lingüísticas del modelo y minimizar los costos de entrenamiento y despliegue. Los esfuerzos existentes han descrito principalmente relaciones complejas entre el rendimiento del modelo, el tamaño de los parámetros y el tamaño de los datos, así como han buscado la asignación óptima de recursos computacionales para entrenar LLMs. Sin embargo, han pasado por alto los impactos de la longitud del contexto y la configuración de las cabezas de atención (el número de cabezas de consulta y clave-valor en la atención de consulta agrupada) en el entrenamiento y la inferencia. En este artículo, comparamos sistemáticamente modelos con diferentes tamaños de parámetros, longitudes de contexto y configuraciones de cabezas de atención en términos de rendimiento del modelo, costo computacional y costo de memoria. Luego, extendemos los métodos de escalado existentes, que se basan únicamente en el tamaño de los parámetros y el cómputo de entrenamiento, para guiar la construcción de LLMs óptimos en costo durante el entrenamiento y la inferencia. Nuestros estudios cuantitativos de escalado muestran que, al procesar secuencias suficientemente largas, un modelo más grande con menos cabezas de atención puede lograr una pérdida menor mientras incurre en costos computacionales y de memoria más bajos. Nuestros hallazgos proporcionan insights valiosos para el desarrollo de LLMs prácticos, especialmente en escenarios de procesamiento de contextos largos. Publicaremos nuestro código y datos de manera abierta.

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.