X-EcoMLA: Upcycling Pre-Trained Attention into MLA for Efficient and Extreme KV Compression

• URL: http://arxiv.org/abs/2503.11132v2
• Date: Sat, 29 Mar 2025 04:43:11 GMT
• Title: X-EcoMLA: Upcycling Pre-Trained Attention into MLA for Efficient and Extreme KV Compression
• Authors: Guihong Li, Mehdi Rezagholizadeh, Mingyu Yang, Vikram Appia, Emad Barsoum,
• Abstract summary: Multi-head latent attention (MLA) is designed to optimize KV cache memory through low-rank key-value joint compression.<n>We show that our proposed method can effectively compress the KV cache while preserving the performance on the benchmarks.
• Score: 23.023849840907594
• License: http://creativecommons.org/licenses/by-nc-nd/4.0/
• Abstract: Multi-head latent attention (MLA) is designed to optimize KV cache memory through low-rank key-value joint compression. Rather than caching keys and values separately, MLA stores their compressed latent representations, reducing memory overhead while maintaining the performance. While MLA improves memory efficiency without compromising language model accuracy, its major limitation lies in its integration during the pre-training phase, requiring models to be trained from scratch. This raises a key question: can we use MLA's benefits fully or partially in models that have already been pre-trained with different attention mechanisms? In this paper, we propose X-EcoMLA to deploy post training distillation to enable the upcycling of Transformer-based attention into an efficient hybrid MLA variant through lightweight post-training adaptation, bypassing the need for extensive pre-training. We demonstrate that leveraging the dark knowledge of a well-trained model can enhance training accuracy and enable extreme KV cache compression in MLA without compromising model performance. The experimental results show that our proposed method can effectively compress the KV cache while preserving the performance on the benchmarks; specifically, for Llama3.2-1B-Instruct baseline, a 6.4x compression achieves the same average score by using only 3.6B training tokens and 70 GPU hours on AMD MI300, whereas a 10.6x compression have less than 0.1\% average score drop with 7B training tokens and 140 GPU hours.

Related papers

• Efficient Token Compression for Vision Transformer with Spatial Information Preserved [59.79302182800274]
  Token compression is essential for reducing the computational and memory requirements of transformer models. We propose an efficient and hardware-compatible token compression method called Prune and Merge.
  arXiv  Detail & Related papers  (2025-03-30T14:23:18Z)
• DBudgetKV: Dynamic Budget in KV Cache Compression for Ensuring Optimal Performance [125.81664663201282]
  We introduce a new KV cache compression method dubbed DBudgetKV. It features an attention-based metric to signal when the remaining KV cache is unlikely to match the full-cache performance, then halting the pruning process. Our method is easy to integrate within LLM inference, not only optimizing memory space, but also showing reduced inference time compared to existing methods.
  arXiv  Detail & Related papers  (2025-02-24T06:33:39Z)
• Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs [74.74225314708225]
  Multi-head Latent Attention (MLA) is an innovative architecture designed to ensure efficient and economical inference.<n>This paper proposes the first data-efficient fine-tuning method for transitioning from Multi-Head Attention to MLA.
  arXiv  Detail & Related papers  (2025-02-20T18:50:42Z)
• TransMLA: Multi-Head Latent Attention Is All You Need [22.354283924006786]
  We introduce Multi-head Latent Attention (MLA) to solve communication bottlenecks in large language models.<n>We show that GQA can always be represented by MLA while maintaining the same KV cache overhead, but the converse does not hold.<n>We plan to develop MLA-specific inference acceleration techniques to preserve low latency in transformed models.
  arXiv  Detail & Related papers  (2025-02-11T18:20:18Z)
• Lillama: Large Language Models Compression via Low-Rank Feature Distillation [8.090496457850852]
  Lillama is a compression method that distills activations with low-rank weights. It compresses Mixtral-8x7B within minutes on a single A100 GPU, removing 10 billion parameters while retaining over 95% of its original performance. It generalizes well to non-transformer architectures, compressing Mamba-3B by 20% while maintaining 99% performance.
  arXiv  Detail & Related papers  (2024-12-21T18:04:01Z)
• KVSharer: Efficient Inference via Layer-Wise Dissimilar KV Cache Sharing [58.29726147780976]
  We propose a plug-and-play method called textit KVSharer, which shares the KV cache between layers to achieve layer-wise compression. Experiments show that textit KVSharer can reduce KV cache computation by 30%, thereby lowering memory consumption. We verify that textit KVSharer is compatible with existing intra-layer KV cache compression methods, and combining both can further save memory.
  arXiv  Detail & Related papers  (2024-10-24T08:06:41Z)
• MatryoshkaKV: Adaptive KV Compression via Trainable Orthogonal Projection [14.073722038551125]
  KV cache has become a de facto technique for the inference of large language models. This paper uses low-rank projection matrices to transform the cache features into spaces with reduced dimensions. We find that our method can sustain over 90% performance with an average KV cache compression rate of 60%.
  arXiv  Detail & Related papers  (2024-10-16T08:34:51Z)
• LoRC: Low-Rank Compression for LLMs KV Cache with a Progressive Compression Strategy [59.1298692559785]
  Key-Value ( KV) cache is crucial component in serving transformer-based autoregressive large language models (LLMs) Existing approaches to mitigate this issue include: (1) efficient attention variants integrated in upcycling stages; (2) KV cache compression at test time; and (3) KV cache compression at test time. We propose a low-rank approximation of KV weight matrices, allowing plug-in integration with existing transformer-based LLMs without model retraining. Our method is designed to function without model tuning in upcycling stages or task-specific profiling in test stages.
  arXiv  Detail & Related papers  (2024-10-04T03:10:53Z)
• Token Compensator: Altering Inference Cost of Vision Transformer without Re-Tuning [63.43972993473501]
  Token compression expedites the training and inference of Vision Transformers (ViTs) However, when applied to downstream tasks, compression degrees are mismatched between training and inference stages. We propose a model arithmetic framework to decouple the compression degrees between the two stages.
  arXiv  Detail & Related papers  (2024-08-13T10:36:43Z)
• Compact Language Models via Pruning and Knowledge Distillation [61.56557874432008]
  Minitron models exhibit up to a 16% improvement in MMLU scores compared to training from scratch. Deriving 8B and 4B models from an already pretrained 15B model using our approach requires up to 40x fewer training tokens per model compared to training from scratch.
  arXiv  Detail & Related papers  (2024-07-19T21:47:57Z)
• Asymmetric Masked Distillation for Pre-Training Small Foundation Models [52.56257450614992]
  Self-supervised foundation models have shown great potential in computer vision thanks to the pre-training paradigm of masked autoencoding. This paper focuses on pre-training relatively small vision transformer models that could be efficiently adapted to downstream tasks. We propose a new asymmetric masked distillation (AMD) framework for pre-training relatively small models with autoencoding.
  arXiv  Detail & Related papers  (2023-11-06T14:44:34Z)
• Large Product Key Memory for Pretrained Language Models [12.932177565788974]
  Product key memory (PKM) enables to improve prediction accuracy by increasing model capacity efficiently with insignificant computational overhead. Motivated by the recent success of pretrained language models (PLMs), we investigate how to incorporate large PKM into PLMs that can be fine for a wide variety of downstream NLP tasks.
  arXiv  Detail & Related papers  (2020-10-08T10:19:50Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.