TLDR: Compressed Convolutional Attention (CCA) is a new method that addresses the high computational and memory costs of Multi-headed Attention (MHA) in large language models. Unlike previous methods like GQA and MLA that primarily reduce KV-cache, CCA down-projects queries, keys, and values and performs the entire attention operation in a compressed latent space. This design, combined with innovations like convolutional mixing, q-k-mean, and value-shift, dramatically cuts parameters, KV-cache, and FLOPs. Its variant, Compressed Convolutional Grouped Query Attention (CCGQA), further enhances efficiency by integrating GQA-style head sharing. Experiments show CCA and CCGQA consistently outperform GQA and MLA in model quality at equal KV-cache compression, leading to faster training and inference prefill, making long-context transformers more feasible.
Large language models, the powerful AI systems behind many modern applications, rely heavily on a core component called Multi-headed Attention (MHA). While MHA is incredibly expressive, allowing the model to weigh the importance of every part of an input sequence, this flexibility comes at a significant computational cost. Training and running these models, especially with long input texts, becomes very expensive due to MHA’s quadratic increase in computation with sequence length and a linearly growing memory footprint for storing ‘keys’ and ‘values’ (the KV-cache).
This computational burden is particularly problematic during training and ‘prefill’ (the initial processing of an input sequence), and the large KV-cache slows down ‘decoding’ (the process of generating new text token by token). Researchers have explored various ways to address these bottlenecks.
Some approaches, like Grouped Query Attention (GQA) and Multi-Latent Attention (MLA), have focused on reducing the KV-cache size. GQA achieves this by sharing parameters across different ‘heads’ of attention, which helps speed up decoding. MLA takes a different route, compressing the keys and values into a smaller, shared latent space. While both GQA and MLA offer improvements in memory efficiency and decoding speed, they largely leave the computational cost of training and prefill unchanged, or even slightly increase it in MLA’s case.
A new research paper, Compressed Convolutional Attention: Efficient Attention in a Compressed Latent Space, introduces a novel method called Compressed Convolutional Attention (CCA) that tackles these challenges head-on. Developed by Tomas Figliolia, Nicholas Alonso, Rishi Iyer, Quentin Anthony, and Beren Millidge from Zyphra, CCA dramatically reduces parameters, KV-cache size, and computational operations (FLOPs) simultaneously.
How Compressed Convolutional Attention Works
The core idea behind CCA is elegant: instead of performing the attention operation in the full, high-dimensional space, it first ‘down-projects’ the queries, keys, and values into a much smaller, shared latent space. The entire attention calculation then happens within this compressed space. This simple design leads to substantial savings across the board.
To ensure that this compression doesn’t compromise performance, CCA incorporates three key innovations:
- Convolutional Mixing: CCA applies convolutional operations across both the sequence and channel dimensions of the compressed queries and keys. This helps to smooth information transfer and preserve important details, similar to how convolutions are used in other neural network architectures.
- Q-K-Mean: This operation adds the mean of the query and key values (before and after convolution) back into the convolved values. This helps in sharing information between queries and keys and allows the model to adjust the strength of the convolutions.
- Value-Shift: For the value projection, each attention head receives two types of value vectors: one from the current input and one from the previous input in the sequence. This introduces a useful inductive bias, where half the heads cannot see the present, which has been shown to improve performance in other models.
By performing attention entirely in this compressed latent space, CCA seamlessly integrates positional embeddings like RoPE, avoiding the complexities faced by some other methods.
Introducing CCGQA: Combining Strengths
The researchers also found that CCA can be combined with parameter-sharing techniques like GQA. This led to Compressed Convolutional Grouped Query Attention (CCGQA), which applies GQA-style key and value head sharing within the already compressed latent space. CCGQA offers even greater KV-cache reduction without sacrificing performance and allows for flexible tuning of compression rates for queries and keys.
Significant Performance Gains
Experiments show that CCGQA consistently outperforms both GQA and MLA when compared at equal KV-cache compression rates, across both dense and Mixture-of-Experts (MoE) models. On MoE models, CCGQA achieved an impressive 8x KV-cache compression with no drop in performance compared to standard MHA. CCA and CCGQA also significantly reduce the FLOP cost of attention, leading to much faster training and prefill times.
On H100 GPUs, a specialized fused kernel for CCA/CCGQA reduced prefill latency by approximately 1.7 times at a sequence length of 16,000 tokens relative to MHA, and accelerated backward pass computations by about 1.3 times. These speedups are crucial for making long-context models more practical and affordable to train and deploy.
Also Read:
- ChunkLLM: A New Approach to Faster and More Efficient Large Language Model Inference
- A New Approach to Efficient Large Language Models with Dynamic Expert Management
A Versatile and Promising Future
The development of CCA and CCGQA highlights that there is significant redundancy in traditional attention mechanisms. By performing attention in a compressed latent space, these methods offer a versatile solution that can be adjusted to balance computational and memory constraints without compromising model quality. This approach is also compatible with existing and future parallelism strategies, making it a strong candidate for improving the efficiency of next-generation large language models.
