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

[Answer 2] Comparison with other KV Cache compression techniques

We thank the reviewers for suggesting a comparison with other KV cache compression approaches. We evaluated our X-EcoMLA against the popular H2O method using the same base model (LLaMA3.2-1B-Instruct) and identical KV sizes. We report the results on lm-eval-hardness benchmark:

As the results show, across all compression levels, X-EcoMLA consistently outperforms H2O in average accuracy and most individual tasks. For instance, at a KV size of 9.4%, X-EcoMLA achieves 50.49% vs. H2O’s 45.05%, with strong gains in ARC, ARE, and PIQA. These results demonstrate that X-EcoMLA offers superior accuracy under the same memory constraints.

We will include this comparison in the final version of the paper and discuss the implications for memory-efficient inference.

[Answer 3] System-Level Inference Metrics

We appreciate the reviewer’s request for system-level metrics such as latency, throughput, and peak memory usage. Below, we report throughput (sequences/sec) and peak GPU memory (GB) for both Llama3.1-8B and X-EcoMLA-8B (kv_rank=128, 10.67× compression) across a range of batch sizes. All measurements were taken on the same hardware (8×AMD-MI300 GPUs) under identical settings.

✅ 1. Throughput Comparison

• X-EcoMLA-8B achieves roughly 1.7–2× higher throughput than the baseline at all batch sizes.
• The baseline model (Llama3.1-8B) runs out of memory (OOM) beyond batch 128, whereas X-EcoMLA continues to scale to batch 1024 without OOM.

✅ 2. Peak Memory Usage

• X-EcoMLA-8B reduces peak memory consumption significantly at small-to-medium batch sizes, enabling much larger batches.
• For example, at batch=128, Llama3.1-8B uses 143 GB (OOM beyond that), while X-EcoMLA-8B uses only 28 GB, a 5× reduction.

These results demonstrate that X-EcoMLA not only preserves model accuracy but also delivers substantial system-level gains—higher throughput, and lower memory footprint—making it highly suitable for latency- and memory-constrained deployment scenarios.

We appreciate the reviewer’s remarks regarding comparison with existing post-training low-rank methods, comparisons with other KV cache compression techniques, and system-level inference metrics. We have addressed all comments in the rebuttal with detailed analysis and supporting results. We kindly ask the reviewer to consider re-evaluating our work in light of these clarifications and improvements.

[Answer 1] Distinction from Existing Post‐Training Low-Rank Methods

We conducted experiments to compare our solution with two of the most recent SOTA methods — MHA2MLA [1] (released concurrently with our work) and PALU [2] — and to clarify how X-EcoMLA differs from and improves upon these approaches.

1. Comparison with MHA2MLA [1]

• Baseline Setup

  • Student / Teacher: SmolLM 1.7B-Instruct
  • MHA2MLA: We use the joint-SVD + $S_{\text{higher}}$ variant (as provided in their code) on the same SFT data.

• Key Differences

  • MHA2MLA stores separate RoPE vectors per head. Given a fixed storage budget (e.g., 32 dimensions total for Key-RoPE per token across all heads), each head’s RoPE dimension in MHA2MLA becomes $\frac{32}{{num_{heads}}}$.
  • X-EcoMLA follows the DeepSeek MLA structure, where all heads share one single Key-RoPE. Thus, with 32 storage units, each head in X-EcoMLA uses the full 32-dimensional RoPE, which is $8\times$ larger (for an 8-head MLA) than MHA2MLA’s per-head RoPE. This larger shared RoPE capacity preserves richer positional encoding under compression.

✅ SFT Performance Comparison (SmolLM 1.7B-Instruct)

• Observation: At 12.5% compression, X-EcoMLA (49.34) outperforms MHA2MLA (48.19) by +1.15 avg.
• At 50% compression, X-EcoMLA (50.15) also outperforms MHA2MLA (49.79) by +0.36 avg.

✅ Continual Pretraining Comparison (SmolLM 1.7B Base)

• Observation: Under 12.5% compression, X-EcoMLA (52.85) outperforms the MHA2MLA pre-trained model (51.69) by +1.16 avg.
• X-EcoMLA closes nearly the entire gap to the full baseline, despite using only 12.5% KV cache.

2. Comparison with PALU (ICLR 2025)

• PALU [2] is a SOTA low-rank projection approach for KV-cache compression. Palu decomposes the linear layers into low-rank matrices, caches compressed intermediate states, and reconstructs the full keys and values on the fly. We compare our X-EcoMLA with PALU on the Llama3-8B model.

• Result: At ~15.6% KV size, X-EcoMLA (67.34) nearly matches the full model (67.87) and outperforms PALU-based compression (66.19 / 64.45) by +1.15 / +2.89 avg. despite using ~3× less KV storage.

References

[1] Lin, Yi-Sheng, et al. “MHA2MLA: Upcycling Multi-Head Attention into Multi-Lingual Attention.” arXiv:2502.14837 (2025).
[2] Chang, Chi-Chih, et al. “PALU: KV-Cache Compression with Low-Rank Projection.” ICLR 2025.

We appreciate the reviewer’s detailed feedback on teacher model requirements, scalability across model sizes and attention types, and the impact of teacher size on KV cache compression. We have addressed these concerns in our rebuttal and supported our response with evidence of the method’s effectiveness. We respectfully encourage the reviewer to revisit the evaluation based on these updates.

---

[Answer 1] Dependency on Large Teacher Models in Low-Resource Settings

We agree that accessibility to large teachers may be constrained in low-resource settings, and we have explicitly addressed this scenario in Section B.2.1 and Figure 1 of the paper:

• Table 5 empirically examines two efficiency strategies: increasing training tokens versus employing larger teacher models. While larger teachers (e.g., LLaMA3.2-8B) yield stronger performance (e.g., 55.13 vs. 53.42), they achieve this using fewer tokens and less total training time compared to the 1B teacher trained on double the data.
• However, when such models are inaccessible, scaling up training data remains a viable alternative. For instance, a 1B model trained with 7B tokens and a 1B teacher nearly matches the performance of training with a 3B teacher on 3.5B tokens.
• Similarly, training with a 3B teacher and 7B tokens approaches the performance of using an 8B teacher with half the data.

In sum, while stronger teachers offer the most cost-effective improvements, our method also scales well with data when larger teachers are not feasible—preserving its practical utility across varying resource regimes.

---

[Answer 2] Generalizing to larger model sizes and other attention architectures

We emphasize that X-EcoMLA is architecture-agnostic and scales effectively to larger models. We support this claim with results on both non-GQA (MHA) and 8B-scale models:

✅Applicability to MHA — SmolLM 1.7B-Instruct

We applied X-EcoMLA to SmolLM, which uses MHA. The updated results below show that, even at 12.5% KV size, our model retains >99% of the original accuracy, and at 50% KV size, it is nearly identical to the baseline.

✅ Scalability to Larger Models — Llama 3-8B and Llama3.1-8B

X-EcoMLA maintains high performance even at extreme KV compression when applied to larger-scale models (8B). Notably, at KV-size 0.109, performance remains nearly on par with the full model.

These results confirm that X-EcoMLA generalizes well to both standard attention architectures (MHA) and larger-scale models, retaining competitive performance with aggressive KV compression.

---

[Answer 3] Analysis of Teacher Size for Different KV Cache Compression

Thank you for your question. Based on the results from Table 2 and Figure 1, we provide a rule-of-thumb guideline for selecting the minimum teacher size required to recover the base model’s accuracy (~52.77) under different KV cache compression ratios and training data budgets for the 1B model. We evaluate this trade-off under two data budgets: ~7B tokens and ~3.6B tokens.

• Extreme compression (≥ 10×) demands an 8 B teacher if you want to recover full-model performance—otherwise you lose 1–2 points.
• Moderate compression (∼ 6×) can be handled by a 3 B teacher (for large data budgets) but needs 8 B when data is scarcer.
• Mild compression (≤ 2×) is already supported by a 1 B teacher, even with only 3.6 B tokens.

We thank the reviewer for the constructive comments on Scalability to Larger Models, Long Context Evaluations, and Applicability to Base Models and Pretraining. We have addressed nearly all of these points and demonstrated the promising effectiveness of our solution. We would greatly appreciate it if the respected reviewer would consider re-evaluating our work and updating the score accordingly.

---

[Answer 1] Scalability to Larger Models — Llama 3-8B and Llama3.1-8B

We have extended X-EcoMLA to two 8 B-parameter models (Llama3-8B and Llama3.1-8B). The results below demonstrate that even under aggressive KV cache compression, X-EcoMLA maintains performance very close to the full-scale baseline. Notably, at KV-size 0.109, performance remains nearly on par with the full model.

In the final version of our paper, we will present additional results of X-EcoMLA on 8 B-parameter LLaMA models.

---

[Answer 2] Long Context Evaluations

We evaluated our MLA-optimized models on the LongBench benchmark, which covers a range of long-context understanding tasks such as LCC, Qasper, QMSum, Multi-News, and SamSum. We report the results under various KV-cache size for both llama3.2-1B and llama3.2-3B models:

Notably:

• X-EcoMLA-3B-A achieves 60.03 on LCC, outperforming the full-size LLaMA3.2-3B (52.11), despite using only 43% of the KV cache.
• Across tasks like Qasper, Multi-News, and SamSum, our compressed models match or slightly improve upon baseline performance, demonstrating strong long-context understanding even under aggressive memory constraints.

These results indicate that our method scales well to long-sequence scenarios and is particularly effective in memory-constrained environments.

---

[Answer 3] Applicability to Base Models and Pretraining Setting

While our primary evaluations focused on instruct-tuned models, based on your comment, we have also investigated the effectiveness of X-EcoMLA when applied to base models during continual pretraining.

We report results on SmolLM 1.7B in a continual pretraining setup, and compare with the following baselines:

1. The target full-attention model (uncompressed),
2. A new baseline: a 12.5% KV-compressed model using the prior MHA2MLA checkpoint, trained on 6B tokens (Hugging Face model link), and
3. Our X-EcoMLA model, continually pre-trained on the same 6B tokens (from the MHA2MLA setup), also with a 12.5% compressed KV-cache.

• Our X-EcoMLA base model, pretrained with 12.5% KV size, outperforms the MHA2MLA baseline by +1.16 average accuracy, and closes most of the gap to the full baseline model.

These results confirm that X-EcoMLA can be **effectively applied to the base model.

---

[Answer 4] Missing Citations and Use of Pretraining Data for Up-Cycling

We will include the suggested citations in the final version to better position our work within the up-cycling literature. Moreover, we explored minimal continued pretraining (see Answer #3), showing that X-EcoMLA supports both regimes.

Dear Reviewing Committee,

We kindly encourage the reviewers who have not yet revisited our rebuttal to do so. We believe that the additional results and clarifications more clearly highlight the value of our submission, and we sincerely hope they will be taken into consideration in the final evaluation. We would be more than happy to address any further questions or feedback.

Thank you again for your time and consideration.