Оптимизированное по стоимости групповое внимание с запросами для языковых моделей с длинным контекстом

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

March 12, 2025

Авторы: Yingfa Chen, Yutong Wu, Xu Han, Zhiyuan Liu, Maosong Sun

cs.AI

Аннотация

Создание эффективных и производительных крупномасштабных языковых моделей (LLM) на основе архитектуры Transformer стало в последнее время ключевым направлением исследований, требующим максимизации языковых возможностей моделей при минимизации затрат на обучение и развертывание. Существующие работы в основном описывают сложные взаимосвязи между производительностью модели, размером параметров и объемом данных, а также исследуют оптимальное распределение вычислительных ресурсов для обучения LLM. Однако они упускают из виду влияние длины контекста и конфигурации механизма внимания (количество голов запросов и ключей-значений в группированном внимании) на обучение и вывод. В данной работе мы систематически сравниваем модели с различными размерами параметров, длинами контекста и конфигурациями механизма внимания с точки зрения производительности модели, вычислительных затрат и затрат на память. Затем мы расширяем существующие методы масштабирования, основанные исключительно на размере параметров и вычислительных ресурсах для обучения, чтобы направлять создание экономически оптимальных LLM как в процессе обучения, так и при выводе. Наши количественные исследования масштабирования показывают, что при обработке достаточно длинных последовательностей более крупная модель с меньшим количеством голов внимания может достичь меньшей потери, одновременно снижая вычислительные затраты и затраты на память. Наши результаты предоставляют ценные инсайты для разработки практических LLM, особенно в сценариях обработки длинных контекстов. Мы опубликуем наш код и данные в открытом доступе.

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.