MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding

MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding

The paper "MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding" presents an intriguing methodology aimed at addressing memory bottlenecks in auto-regressive inference of transformers. This is particularly relevant given the growing size of transformers, which necessitates more efficient memory management strategies for practical deployment.

Introduction and Background

The transformer architecture, renowned for its efficacy in NLP tasks, leverages multi-head attention mechanisms that become increasingly resource-intensive as models scale up. The auto-regressive nature of transformer decoding exacerbates memory bandwidth issues due to Key-Value (KV) caching, which scales linearly with model size, batch size, and sequence length. This challenge has been previously addressed through approaches like Multi-Query Attention (MQA) and Grouped-Query Attention (GQA), both of which reduce the number of KV heads to alleviate memory usage.

While MQA and GQA offer notable memory savings, they are limited to sharing KV heads within the same transformer layer. The novelty in this paper lies in extending KV sharing across multiple layers, termed Multi-Layer Key-Value (MLKV) sharing, thereby reducing memory usage beyond the capabilities of MQA and GQA.

Proposed Method: Multi-Layer Key-Value Sharing (MLKV)

The key contribution of MLKV is its ability to share KV heads not just among attention heads within a layer, but also across different layers of the transformer. This approach theorizes that groups of successive layers can leverage similar KV heads due to the analogous nature of computations performed across these layers. Thus, MLKV can reduce the total number of KV heads in the transformer to lower than the number of layers.

The paper details the mathematical formulation of MLKV and provides an illustrative comparison with existing KV sharing methods, highlighting potential memory savings. The theoretical improvements are summarized in Table~\ref{tab:cache-sizes} of the document, indicating significant reductions in KV cache size.

Experimental Setup and Results

The authors utilized the Pythia-160M model trained on a deduplicated version of The Pile dataset for empirical validation. Various MLKV configurations were uptrained from this baseline to evaluate memory savings and performance trade-offs. The uptraining process involved continuing pre-training with 5% of the original dataset, ensuring alignment with strategies proposed in the GQA literature.

Benchmark Performance: The evaluations were performed on standard NLP benchmarks including ARC-e, LAMBADA, PIQA, and SciQ. As reflected in Table~\ref{tab:benchmark-results}, the results indicate that while MLKV's benchmark performance generally decreases with a reduction in KV heads, configurations with head counts as low as one-third the number of layers (i.e., MLKV-6) still maintained competitive accuracy, providing a reasonable performance trade-off.

Inference Metrics: Inference time memory usage and throughput were paramount to the implications of MLKV. Figure~\ref{fig:memvsbatch} illustrates that MLKV configurations substantially decrease memory consumption during generation. Specifically, MLKV-6 and MLKV-2 configurations demonstrate significantly higher memory efficiency while maintaining acceptable accuracy levels compared to the baseline and other KV sharing methods.

Implications and Limitations: MLKV's potential for reducing KV cache size up to a factor of 6x compared to MQA presents a substantial improvement for transformer models deployed in memory-constrained environments. This advancement is especially relevant for applications requiring large sequence lengths or batch sizes, where memory bottlenecks are more pronounced.

However, it is worth noting that the experiments were conducted on relatively small models (160M parameters). The behavior of MLKV at billion-parameter scales, more commonly used in current NLP tasks, remains untested. Furthermore, MLKV has not been evaluated on encoder-decoder models, which could also benefit from this approach in their decoders. Future work should investigate MLKV implementations during the initial pre-training phases and explore its effects across a broader range of downstream tasks.

Conclusion

The proposed MLKV method extends the current capabilities of KV caching by enabling cross-layer KV head sharing. This innovation results in significantly reduced memory usage without substantial degradation in model performance, making it a practical solution for efficient transformer deployment at scale. The empirical results affirm MLKV's viability, particularly for models requiring rigorous memory optimizations. Future research should aim to validate these findings on larger models and diverse transformer architectures.

References

The bibliographical references pertain to foundational works in transformer models and KV caching strategies, which substantiate the methodologies and experimental comparisons in the paper.