Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs (2502.14837v1)

Abstract: Multi-head Latent Attention (MLA) is an innovative architecture proposed by DeepSeek, designed to ensure efficient and economical inference by significantly compressing the Key-Value (KV) cache into a latent vector. Compared to MLA, standard LLMs employing Multi-Head Attention (MHA) and its variants such as Grouped-Query Attention (GQA) exhibit significant cost disadvantages. Enabling well-trained LLMs (e.g., Llama) to rapidly adapt to MLA without pre-training from scratch is both meaningful and challenging. This paper proposes the first data-efficient fine-tuning method for transitioning from MHA to MLA (MHA2MLA), which includes two key components: for partial-RoPE, we remove RoPE from dimensions of queries and keys that contribute less to the attention scores, for low-rank approximation, we introduce joint SVD approximations based on the pre-trained parameters of keys and values. These carefully designed strategies enable MHA2MLA to recover performance using only a small fraction (0.3% to 0.6%) of the data, significantly reducing inference costs while seamlessly integrating with compression techniques such as KV cache quantization. For example, the KV cache size of Llama2-7B is reduced by 92.19%, with only a 0.5% drop in LongBench performance.

Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs

The paper "Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs" presents a novel approach aimed at optimizing the inference process of LLMs, specifically focusing on reducing computational resource requirements through architectural modifications. The paper concentrates on adapting the Multi-Head Latent Attention (MLA) framework, initially developed by DeepSeek, to existing Transformer-based LLMs which predominantly utilize Multi-Head Attention (MHA).

Key Contributions

The research introduces a fine-tuning method, MHA2MLA, which facilitates the transition from MHA to MLA without the need for pre-training from scratch. The methodology encapsulates two primary components:

1. Partial Rotary Position Embedding (RoPE): This approach involves removing RoPE from dimensions that contribute less to attention scores, which aids in enhancing the compression efficiency of the Key-Value (KV) cache.
2. Low-rank Approximation through Joint Singular Value Decomposition (SVD): Utilizing SVD, the paper achieves a low-rank representation of keys and values, thereby compressing the KV cache significantly.

Experimental Results

The paper provides comprehensive empirical results across various model scales, ranging from 135M to 7B parameters. For instance, the paper reports a drastic reduction in the KV cache size—92.19% for Llama2-7B with only a 0.5% decrease in performance on LongBench. This level of compression is achieved using a fine-tuning corpus amounting to merely 3\textperthousand~to~6\textperthousand of the original dataset size, showcasing the data efficiency of the proposed method.

Implications and Future Directions

Practical Implications: The paper demonstrates that MHA2MLA effectively integrates with existing compression techniques such as KV cache quantization, resulting in exceedingly economical inference processes for LLMs. This compatibility with quantization suggests broader applicability in real-world scenarios where computational resources are constrained.

Theoretical Implications: By introducing a systematic approach to transition MHA-based architectures to MLA, the research potentially sets a precedent for future architectural innovations aimed at optimizing resource usage without sacrificing model performance.

Speculations on Future Developments: The ongoing trend towards more efficient LLMs could benefit from this research, particularly in applications where inference cost is a critical concern. As AI continues to evolve, methods like MHA2MLA might serve as foundational techniques for developing next-generation models that balance performance with computational efficiency.

In summary, the paper contributes substantially to the enhancement of LLM efficiency, aligning with industry goals to minimize energy consumption and operational costs. The proposed MHA2MLA framework underlines the possibility of achieving significant resource savings while maintaining performance integrity, heralding a promising direction for future AI research and applications.