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

• URL: http://arxiv.org/abs/2503.09579v2
• Date: Tue, 20 May 2025 09:31:37 GMT
• Title: Cost-Optimal Grouped-Query Attention for Long-Context Modeling
• Authors: Yingfa Chen, Yutong Wu, Chenyang Song, Zhen Leng Thai, Xingyu Shen, Xu Han, Zhiyuan Liu, Maosong Sun,
• Abstract summary: Grouped-Query Attention (GQA) is a widely adopted strategy for reducing the computational cost of attention layers in large language models.<n>We analyze the relationship among context length, model size, GQA configuration, and model loss.<n>We propose a recipe for deriving cost-optimal GQA configurations.
• Score: 58.10054123910204
• License: http://creativecommons.org/licenses/by-sa/4.0/
• Abstract: Grouped-Query Attention (GQA) is a widely adopted strategy for reducing the computational cost of attention layers in large language models (LLMs). However, current GQA configurations are often suboptimal because they overlook how context length influences inference cost. Since inference cost grows with context length, the most cost-efficient GQA configuration should also vary accordingly. In this work, we analyze the relationship among context length, model size, GQA configuration, and model loss, and introduce two innovations: (1) we decouple the total head size from the hidden size, enabling more flexible control over attention FLOPs; and (2) we jointly optimize the model size and the GQA configuration to arrive at a better allocation of inference resources between attention layers and other components. Our analysis reveals that commonly used GQA configurations are highly suboptimal for long-context scenarios. More importantly, we propose a recipe for deriving cost-optimal GQA configurations. Our results show that for long-context scenarios, one should use fewer attention heads while scaling up model size. Configurations selected by our recipe can reduce both memory usage and FLOPs by more than 50% compared to Llama-3's GQA, with *no degradation in model capabilities*. Our findings offer valuable insights for designing efficient long-context LLMs. The code is available at https://www.github.com/THUNLP/cost-optimal-gqa .

Related papers

• Ada-R1: Hybrid-CoT via Bi-Level Adaptive Reasoning Optimization [86.56120216550232]
  We propose a novel two-stage framework for adaptive and efficient reasoning.<n>First, we construct a hybrid reasoning model by merging long and short CoT models.<n>Second, we apply bi-level preference training to guide the model to select suitable reasoning styles.
  arXiv  Detail & Related papers  (2025-04-30T14:01:45Z)
• EfficientLLaVA:Generalizable Auto-Pruning for Large Vision-language Models [64.18350535770357]
  We propose an automatic pruning method for large vision-language models to enhance the efficiency of multimodal reasoning.<n>Our approach only leverages a small number of samples to search for the desired pruning policy.<n>We conduct extensive experiments on the ScienceQA, Vizwiz, MM-vet, and LLaVA-Bench datasets for the task of visual question answering.
  arXiv  Detail & Related papers  (2025-03-19T16:07:04Z)
• LESA: Learnable LLM Layer Scaling-Up [57.0510934286449]
  Training Large Language Models (LLMs) from scratch requires immense computational resources, making it prohibitively expensive.<n>Model scaling-up offers a promising solution by leveraging the parameters of smaller models to create larger ones.<n>We propose textbfLESA, a novel learnable method for depth scaling-up.
  arXiv  Detail & Related papers  (2025-02-19T14:58:48Z)
• Don't Do RAG: When Cache-Augmented Generation is All You Need for Knowledge Tasks [11.053340674721005]
  Retrieval-augmented generation (RAG) has gained traction as a powerful approach for enhancing language models by integrating external knowledge sources.<n>This paper proposes an alternative paradigm, cache-augmented generation (CAG) that bypasses real-time retrieval.
  arXiv  Detail & Related papers  (2024-12-20T06:58:32Z)
• Spectral Subspace Clustering for Attributed Graphs [3.974852803981998]
  Subspace clustering seeks to identify subspaces that segment a set of n data points into k (kn) groups. This paper presents two effective and efficient algorithms, S2CAG and M-S2CAG, for SCAG computation.
  arXiv  Detail & Related papers  (2024-11-17T13:22:15Z)
• Squeezed Attention: Accelerating Long Context Length LLM Inference [61.787865959140994]
  We propose Squeezed Attention to accelerate applications where a large portion of the input context is fixed.<n>During inference, we compare query tokens from the user input with the centroids to predict which keys from the fixed context are semantically relevant.<n>We also present a hierarchical version of our algorithm which can reduce the complexity of attention from linear to logarithmic with respect to the fixed context length.
  arXiv  Detail & Related papers  (2024-11-14T18:54:19Z)
• CORAG: A Cost-Constrained Retrieval Optimization System for Retrieval-Augmented Generation [22.918861762038116]
  Large Language Models (LLMs) have demonstrated remarkable generation capabilities but often struggle to access up-to-date information. Retrieval-Augmented Generation (RAG) addresses this issue by incorporating knowledge from external databases.
  arXiv  Detail & Related papers  (2024-11-01T17:11:16Z)
• Computational Bottlenecks of Training Small-scale Large Language Models [19.663560481459164]
  Small-scale large Language Models (SLMs) are gaining attention due to cost and efficiency demands from consumers. In this study, we explore the computational bottlenecks of training SLMs. We assess these factors on popular cloud services using metrics such as loss per dollar and tokens per second.
  arXiv  Detail & Related papers  (2024-10-25T10:30:21Z)
• Improving Small-Scale Large Language Models Function Calling for Reasoning Tasks [0.8425561594225592]
  This study introduces a novel framework for training smaller language models in function calling. It focuses on specific logical and mathematical reasoning tasks. The approach aims to improve performances of small-scale models for these tasks using function calling.
  arXiv  Detail & Related papers  (2024-10-24T16:27:35Z)
• Rational Metareasoning for Large Language Models [5.5539136805232205]
  Being prompted to engage in reasoning has emerged as a core technique for using large language models (LLMs)<n>This work introduces a novel approach based on computational models of metareasoning used in cognitive science.<n>We develop a reward function that incorporates the Value of Computation by penalizing unnecessary reasoning.
  arXiv  Detail & Related papers  (2024-10-07T23:48:52Z)
• Inference Scaling for Long-Context Retrieval Augmented Generation [37.15479223789199]
  scaling inference computation has unlocked the potential of long-context large language models (LLMs)<n>In this work, we investigate the combination of multiple strategies beyond simply increasing the quantity of knowledge, including in-context learning and iterative prompting.<n>We show that scaling inference compute on long-context LLMs achieves up to 58.9% gains on benchmark datasets compared to standard RAG.
  arXiv  Detail & Related papers  (2024-10-06T03:42:15Z)
• Efficient Training of Self-Supervised Speech Foundation Models on a Compute Budget [53.882492597728195]
  This paper investigates how to efficiently train speech foundation models with self-supervised learning (SSL) under a limited compute budget. We examine critical factors in SSL that impact the budget, including model architecture, model size, and data size.
  arXiv  Detail & Related papers  (2024-09-09T10:36:42Z)
• Achieving Peak Performance for Large Language Models: A Systematic Review [0.0]
  Large language models (LLMs) have achieved remarkable success in natural language processing (NLP) As models grow into the trillion- parameter range, computational and memory costs increase significantly. This makes it difficult for many researchers to access the resources needed to train or apply these models.
  arXiv  Detail & Related papers  (2024-09-07T13:57:41Z)
• Understanding the Performance and Estimating the Cost of LLM Fine-Tuning [9.751868268608675]
  Fine-tuning Large Language Models (LLMs) for specific tasks in a cost-effective manner. In this paper, we characterize sparse Mixture of Experts (MoE) based LLM fine-tuning to understand their accuracy and runtime performance. We also develop and validate an analytical model to estimate the cost of LLM fine-tuning on the cloud.
  arXiv  Detail & Related papers  (2024-08-08T16:26:07Z)
• ChatQA 2: Bridging the Gap to Proprietary LLMs in Long Context and RAG Capabilities [53.97515452727115]
  ChatQA 2 is a Llama 3.0-based model with a 128K context window.<n>We present a training recipe to extend the context window of Llama3-70B-base from 8K to 128K tokens.<n>We find that the performance of strong long-context LLMs using RAG improves when retrieving a larger number of chunks.
  arXiv  Detail & Related papers  (2024-07-19T17:35:47Z)
• Cost-efficient Knowledge-based Question Answering with Large Language Models [28.816821631082856]
  Knowledge-based question answering (KBQA) is widely used in many scenarios that necessitate domain knowledge. Large language models (LLMs) bring opportunities to KBQA, while their costs are significantly higher and absence of domain-specific knowledge during pre-training. We propose Coke, a novel cost-efficient strategy for KBQA with LLMs, modeled as a tailored multi-armed bandit problem.
  arXiv  Detail & Related papers  (2024-05-27T16:37:34Z)
• MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies [85.57899012821211]
  Small Language Models (SLMs) are a resource-efficient alternative to Large Language Models (LLMs) We introduce MiniCPM, specifically the 1.2B and 2.4B non-embedding parameter variants. We also introduce MiniCPM family, including MiniCPM-DPO, MiniCPM-MoE and MiniCPM-128K.
  arXiv  Detail & Related papers  (2024-04-09T15:36:50Z)
• E^2-LLM: Efficient and Extreme Length Extension of Large Language Models [74.1254067728251]
  We propose an Efficient and Extreme length extension method for Large Language Models, called E 2 -LLM, with only one training procedure and dramatically reduced cost. Comprehensive experimental results on multiple benchmark datasets demonstrate the effectiveness of our E 2 -LLM on challenging long-context tasks.
  arXiv  Detail & Related papers  (2024-01-13T02:11:20Z)
• On Optimal Caching and Model Multiplexing for Large Model Inference [66.50550915522551]
  Large Language Models (LLMs) and other large foundation models have achieved noteworthy success, but their size exacerbates existing resource consumption and latency challenges. We study two approaches for mitigating these challenges: employing a cache to store previous queries and learning a model multiplexer to choose from an ensemble of models for query processing.
  arXiv  Detail & Related papers  (2023-06-03T05:01:51Z)
• To Repeat or Not To Repeat: Insights from Scaling LLM under Token-Crisis [50.31589712761807]
  Large language models (LLMs) are notoriously token-hungry during pre-training, and high-quality text data on the web is approaching its scaling limit for LLMs. We investigate the consequences of repeating pre-training data, revealing that the model is susceptible to overfitting. Second, we examine the key factors contributing to multi-epoch degradation, finding that significant factors include dataset size, model parameters, and training objectives.
  arXiv  Detail & Related papers  (2023-05-22T17:02:15Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.