Slim attention: cut your context memory in half without loss -- K-cache is all you need for MHA (2503.05840v2)

Abstract: Slim attention shrinks the context memory size by 2x for transformer models with MHA (multi-head attention), which can speed up inference by up to 2x for large context windows. Slim attention is an exact, mathematically identical implementation of the standard attention mechanism and therefore doesn't compromise model accuracy. In other words, slim attention losslessly compresses the context memory by a factor of 2. For encoder-decoder transformers, the context memory size can be reduced even further: For the Whisper models for example, slim attention reduces the context memory by 8x, which can speed up token generation by 5x for batch size 64 for example. And for the T5-11B model for example, the memory can be reduced by 32x because its MHA projection dimension is larger than the embedding dimension. See https://github.com/OpenMachine-ai/transformer-tricks for code and more transformer tricks, and https://www.youtube.com/watch?v=uVtk3B6YO4Y for this paper's YouTube video.

Slim Attention: Efficient Context Memory Management for Transformers

The paper "Slim Attention: Cut Your Context Memory in Half Without Loss of Accuracy — K-cache is All You Need for MHA" presents an innovative approach to optimizing memory usage in transformer models equipped with multi-head attention (MHA) mechanisms. The authors propose a methodology termed "slim attention," which effectively compresses the context memory size by a factor of two without compromising accuracy, thereby enhancing the efficiency of large-context transformer models.

Mechanism Overview

Slim attention utilizes a mathematical equivalent approach to the standard multi-head self-attention mechanism, where keys (K) are stored and values (V) are computed on-the-fly using a derived transformation matrix $W_{KV}$. This is achieved through $V = K W_{KV}$, allowing models to eliminate storing the V-cache and thus halving the context memory footprint. The transformation matrix $W_{KV} = W_K^{-1} W_V$ is precomputed offline, assuming the invertibility of the key weight matrix $W_K$.

Strong Numerical Results

The paper highlights compelling numerical outcomes, including memory reduction factors of 2x and 8x for Whisper models, and up to 32x for rare configurations such as the T5-11B model. This methodology significantly reduces the resource requirements associated with long context lengths. The Phi-3-mini-128k model serves as a specific case paper, where slim attention reduces context memory size from 25GB to 12.5GB, resulting in a speedup of token generation by up to 2x.

Implications and Future Developments

Practically, slim attention is relevant for memory-bound systems where inference speed is constrained by bandwidth limitations. It promises noticeable advancements in applications demanding lengthy context windows, such as natural language processing, speech recognition, and machine translation models. Theoretically, the approach underscores the potential of algebraic optimization techniques in transformer architecture, paving the way for more compact and efficient models.

Future developments could focus on integrating slim attention into popular AI frameworks and further refinement of mathematical transformations to support non-square projection matrices. Additionally, extending support to more widespread architectures such as group query attention (GQA) or multi-query attention (MQA) can broaden the applicability of slim attention.

Conclusion

Slim attention represents a substantial improvement in resource efficiency for transformers, reaffirming the importance of post-training optimizations and memory management strategies. By leveraging mathematical transformations, it offers a pathway towards more scalable and efficient transformer models without the need for extensive retraining, aligning with the evolving demands of AI applications with extensive context lengths.