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

多头潜在注意力（MLA）是由DeepSeek提出的一种创新架构，旨在通过显著压缩键值（KV）缓存到潜在向量中，确保高效的、经济的推理过程。与MLA相比，使用多头注意力（MHA）及其变体如分组查询注意力（GQA）的标准大语言模型（LLM）表现出明显的成本劣势。使经过良好训练的大语言模型（例如Llama）无需从头开始预训练就能快速适应MLA，这既具有重要意义也充满挑战。本文提出了首个数据高效微调方法，用于从MHA过渡到MLA（MHA2MLA），该方法包含两个关键组件：对于部分旋转位置编码（partial-RoPE），我们移除了对注意力分数贡献较小的查询和键维度上的RoPE；对于低秩近似，我们基于预训练的键和值参数引入了联合奇异值分解（SVD）近似。这些精心设计的策略使得MHA2MLA仅使用一小部分数据（0.3%到0.6%）就能恢复性能，大幅降低推理成本，同时无缝结合诸如KV缓存量化等压缩技术。例如，Llama2-7B的KV缓存大小减少了92.19%，而LongBench性能仅下降了0.5%。

论文试图解决的问题是，如何使现有的大型语言模型（LLM）如Llama能够高效地转换到Multi-head Latent Attention (MLA)架构，以显著降低推理成本和压缩Key-Value (KV)缓存，而无需从头开始预训练。这是一个具有挑战性的问题，因为直接转换需要确保性能不会大幅下降，并且要尽可能减少所需的训练数据量。

关键思路是提出了一个数据高效型微调方法MHA2MLA，用于将多头注意力机制(MHA)转换为多头潜在注意力(MLA)。该方法包括两个核心组件：部分RoPE（移除对注意力分数贡献较小的查询和键维度上的RoPE），以及低秩近似（基于预训练的键和值参数引入联合SVD近似）。这些策略使得模型能够在使用极少量数据的情况下恢复性能，同时大幅减少了KV缓存的大小。

论文的亮点在于它不仅提出了一种新颖的方法来实现从MHA到MLA的有效转换，而且展示了通过仅用0.3%到0.6%的数据就能显著减少KV缓存大小并保持接近原始性能的结果。此外，该研究还探讨了与量化技术结合的可能性，进一步降低了存储需求。实验设计涵盖了对不同规模模型的测试，并使用了LongBench等基准进行评估。虽然未明确提及代码是否开源，但其方法论为未来研究提供了有价值的参考方向。

最近的相关研究包括探索更高效的注意力机制变体，如Grouped-Query Attention (GQA)，以及针对大规模语言模型优化的技术，如KV缓存压缩和量化。其他相关研究可能包括《Efficient Transformers: A Survey》、《Compressing Transformer Models for Inference》及《Optimizing Large Language Models with Structured Pruning and Quantization》等。