Zebra-Llama: Towards Extremely Efficient Hybrid Models

We sincerely thank the reviewer for the comprehensive review and valuable feedback, particularly your recognition of our method's strengths in KV cache saving, inference speedup, and the effectiveness of our intermediate layer distillation and layer selection strategy. We have carefully addressed the insightful comments regarding comparing MMLU scores, comparisons with Llamba, model scalability, and performance on other datasets and supported our response with additional experimental evidence. We appreciate your consideration of these updates and would be grateful if you could kindly revisit your evaluation of our paper in light of the improvements.

---

[Answer 1] Comparing MMLU score with Llamba for the 8B model

We acknowledge the observed MMLU score gap between our 8B Zebra-Llama and Llamba-8B. As shown in [1], the MMLU task format can be particularly challenging for State Space Models (SSMs) to adapt to, often requiring more training efforts to improve. After a thorough review of the Llamba paper, we believe this gap likely comes from two key factors:

1. Fewer training tokens and smaller teacher model: Our 8B Zebra-Llama model used the 8B base model as its teacher and was trained on 11 billion tokens. Llamba's 8B model, conversely, was distilled from a more powerful 70B teacher model in its final stage and trained on an additional 1 billion tokens (12 billion total). We believe this combination of a stronger teacher and a larger training budget contributed to Llamba's better MMLU performance.
2. More curated dataset selection Our Zebra-Llama model utilizes the same training datasets as Mamba-In-Llama (OpenHermes-2.5, GenQA, and Infinity-Instruct). The Llamba paper, however, emphasizes the critical role of dataset selection/pre-processing and demonstrates that leveraging a carefully filtered version of Fineweb-Edu could significantly boost the MMLU scores after distillation. We hypothesize that applying a similar rigorous data curation strategy would significantly enhance Zebra-Llama's MMLU performance and lead to higher scores than Llamba. As presented in [Answer 2], our MMLU score significantly surpasses that of Llamba when using the same dataset and training pipeline.

[1] Waleffe, Roger, et al. "An empirical study of mamba-based language models." arXiv preprint arXiv:2406.07887 (2024).

---

[Answer 2] Evaluating Llamba and Zebra-Llama Under Consistent Training Conditions

Thanks for the thoughtful and constructive comment. Following the reviewer’s suggestion, we conducted a controlled comparison between Zebra-Llama and Llamba using the same dataset and training pipeline. Due to time constraints, we were able to complete this evaluation at the 1B model scale.

For both Zebra-Llama and Llamba, we adopt the training pipeline introduced in Llamba paper which contains three stages. For the dataset, as the filtered FineWeb-Edu is not publicly available, we adopt the same training dataset in Zebra-Llama with ~6B tokens in total. We maintained the same token distribution across stages as Llamba: 225M tokens for Stage 1, 2.025B for Stage 2, and 3.75B for Stage 3. All other training parameters were kept consistent.

• Target Model: Llama-3.2-1B-Instruct
• Teacher Model: Llama-3.2-1B-Instruct
• Initial Learning Rate: 8e-5
• Batch Size: 96

The results, presented in the table below, demonstrate that our Zebra-Llama surpasses Llamba (pure SSM) across all 8 tasks with only 3.91% KV cache, highlighting the advantages of integrating efficient attention modules (i.e., MLA) with Mamba.

---

[Answer 3] Scaling to larger models or other model families

We thank the reviewer for the comments. Theortically it is possible to scale our method to larger models but computational resource and time constraints prevented us from conducting such experiments. Scaling to larger models remains in our future plan.

However, we did perform additional experiments on Qwen2.5-0.5B-Instruct and Qwen2.5-1.5B-Instruct to demonstrate that our proposed method (now Zebra-Qwen) performs well across other model families. Using the same training process as for the Llama series, our 0.5B Zebra-Qwen model achieved better performance than the target model with only 6.25% KV cache. For the 1.5B Zebra-Qwen, even when using the same model as the teacher, we achieved similar performance to the target model with 12.5% KV cache size.

---

[Answer 4] Extending Evaluation to More Challenging Tasks

Thanks for the comments and we agree on the importance of including more challenging tasks. As the reviewer suggested, we include the results on GSM8K and GPQA and summarize the results in the following.

Key observations from the results:

• Our 1B model excels, achieving the best average score and outperforming the target Llama-3.2-1B-Instruct model with only 7.81% KV cache.
• Our 3B model shows a 6.7% and 2.63% performance drop compared to the target model and X-EcoMLA, respectively, with 21.32x and 2x KV cache compression. Even at 2.01% KV cache, our overall performance surpasses Mamba-In-Llama and Llamba.
• Our 8B model exhibits a 12% and 7.5% performance drop compared to the target model and MambaInLlama, with 18.28x and 9.14x KV cache compression, respectively. For the performance gap, the reasoning task GSM8K is identified as a primary contributor. We hypothesize this is due to our efficient training scheme (8B teacher + 11B training tokens). Although Mamba-In-Llama exhibits higher GSM8K score, they adopted a 70B teacher and 20B training tokens. Since math reasoning tasks are generally harder, we expect Zebra-Llama to achieve further improvements with more training tokens.
• Our Zebra-Llama consistantly outperforms Llamba across all model sizes

To test the potential of math reasoning of Zebra-Llama, we conducted additional experiments on our 1B model where we insert more samples from OpenMathInstruct dataset during the SFT stage. As shown in the table, GSM8K accuracy consistently gets improved as we add more math tokens. We believe the math reasoning performance of Zebra-Llama will be much improved with more data engineering efforts.

---

[Answer 5] Splitting of $W_{ukv}$

The main goal for the split of $W^{UKV}$ is to align with operation when creating the $W^{UKV}$. The process begins by concatenating the original trained key ($W_K\in\mathbb{R}^{d\times n_{kv}d_{h}}$) and value ($W^{V}\in\mathbb{R}^{d\times n_{kv}d_{h}}$) weight matrices [$W^K$ ; $W^V$] and creating a single low-rank approximation of them called $W^{UKV}\in \mathbb{R}^{r_{kv}\times 2d_{kv}n_{kv}}$ via SVD.

[$W^{K}$ ; $W^{V}$] = $U_{kv}\Sigma_{kv}V_{kv}^T$ = $W^{DKV}$ $W^{UKV}$

Considering how we concatenate $W^{K}$ and $W^{V}$, we know that $W^{UKV}$ can be divided into two parts where the upper part maps the latent to keys and the lower part maps the latent to values.

$W^{UKV}$ = [ $\tilde{W}^{UK}$ ; $W^{UV}$ ]

It is straightforward to get $W^{UV}$. However, since MLA only decompress the nope dimensions for each key head from the latent, we need to select $d_{qk}$ dimensions out of $d_h$ from the upper part $\tilde{W}^{UK}$. In this work, we assume that we always take the first $d_{qk}$ dimensions from $d_h$ as the nope dimensions.

Such procedure can be expressed as:

$\tilde{W}^{UK}$ --> $\bar{W}^{UK}$ (reshape to $r_{kv}\times n_{kv} \times d_h$) --> $W^{UK}$ (take the first $d_{qk}$ from the last dimention and reshape to $r_{kv}\times d_{qv}n_{kv}$)

Thanks for the reply. I think those responses make sense.

Could you please run some additional benchmarks for long-context such as ruler benchmark? My concern is that the models may only perform well on short contexts.

Thank you.

We sincerely thank the reviewer for the constructive feedback. In this revised version, we have carefully addressed the main concerns raised in your review. Specifically, we clarified our key contributions [Answer 1], extended the evaluation of our SMART layer selection method to additional architectural families [Answer 2], completed the previously missing few-shot results [Answer 3], and provided further analysis to clarify the performance gap between the larger distilled and target models [Answer 4]. We hope these substantial updates address your concerns, and we would be grateful if you could revisit your evaluation based on the additional evidence presented.

---

[Answer 1] Our Key Contributions

Thank you for your comment. Given the importance of hybrid model architectures for balancing performance and efficiency, our paper introduces several key contributions spanning architecture design, training methodology, and empirical results:

Architecture: We propose the first hybrid model that combines MLA with Mamba2 layers to replace classical Transformer blocks. This design not only reduces memory usage but also addresses the quadratic attention bottleneck, enabling high-throughput inference with minimal performance degradation.
Training Methodology: We placed particular emphasis on improving initialization and layer selection, as both have a significant impact on model performance. Notably, layer selection has received limited attention in prior work. To this end, we developed careful initialization strategies to enable stable integration of Mamba2 and MLA layers, introduced auxiliary intermediate layer distillation, and proposed our SMART layer selection algorithm. Empirical Results: Using our training pipeline, we successfully trained a family of hybrid models that achieve competitive performance with baseline Transformers across a range of benchmarks. Notably, our models reduce KV cache usage to as low as 2–5% and improve inference throughput by up to 2.1×, while also requiring less training data and shorter training time, surpassing prior state-of-the-art hybrid methods.

---

[Answer 2] SMART Layer Selection on other Architectural Families

Thanks for the comment. We agree that the proposed SMART layer selection is only tested in Llama families and it is necessary to test it with another model series such as Qwen. To show the generalizability of our SMART layer selection algorithm, we apply it to another scenario where we start from the Qwen-2.5-0.5B model (24 layers in total) and we would like to pick up 4 layers as MLA and 20 layers as Mamba2. Since we cannot include figures in the rebuttal. we include the layer sensitivity scores for each layer in the table below. The bold layers are the ones selected by SMART {Layer Indices: 0, 8, 15, 22}.

We also evaluated several alternative layer selection schemes, and the results remained consistent with those reported in the paper (Table X), further demonstrating the generalizability of the proposed SMART layer selection algorithm.

---

[Answer 3] Complete Few-shot Results

We focused on reporting few-shot performance for the 8B model in the main text due to space and time constraints; however, we acknowledge the importance of completeness and now provide the full few-shot performance results for both 1B and 3B models (and will include them in the final version of our paper). As shown, our Zebra-Llama models consistently demonstrate competitive few-shot performance across multiple benchmarks at smaller scales as well.

---

[Answer 4] Performance Gap between Larger Distilled and Target models

Thank you for the comment. The larger performance gap observed for bigger hybrid models (especially the 8B models) is likely due to two main factors:

• the use of the same size 8B teacher for training our Zebra-Llama 8B models—unlike the 1B and 3B models, which had larger teachers—and
• the lack of scaling in training data proportional to model capacity. Due to limitations in time and computational resources, we kept the DPO dataset size fixed and were only able to use up to 11B tokens for SFT when training the 8B models. We believe that with a larger teacher and more extensive training (both in terms of steps and data volume), the performance of the 8B hybrid models can be further improved. We consider this an important direction for future work.

We sincerely thank the reviewer for acknowledging the effectiveness of our method in reducing inference overhead with minimal performance loss, as well as recognizing the strength of our experimental evaluation and the significance of the SMART layer selection strategy. In the rebuttal, we have carefully addressed your comments regarding [Answer 1] task-wise comparison between hybrid and base models, and [Answer 2] the extension of our evaluation to more challenging benchmarks in the following.

---

[Answer 1] Task-Wise Comparison of Hybrid and Base Models

Thank you for the helpful observation regarding performance differences between our hybrid models and their corresponding base models. We carefully reviewed the results across all tasks in Table 1 for the 1B, 3B, and 8B models. Across 7 out of 8 benchmarks, our hybrid models consistently achieve either superior or very close performance compared to their base models. The only notable deviation occurs on the MMLU benchmark, which we discuss in more detail below. When comparing the different size Zebra-Llama models to their corresponding base models, we observed that MMLU performance reaches approximately 84% of the base model at the 1B scale, 86% at the 3B scale, and 83% at the 8B scale, which is pretty consistent performance across different sizes. We explain the potential reasons for the observed MMLU performance gap below:

1) MMLU Task Formatting Difficulty: The observed gap in MMLU performance between hybrid models and pure Transformer baselines may stem from formatting sensitivity in MMLU’s multiple-choice structure. The study by [2] shows that state-space models (SSMs) like Mamba struggle with the standard MMLU format, which involves selecting a single letter (A/B/C/D) corresponding to the correct answer. While SSMs contain the same underlying knowledge as Transformers, they require more training to learn the task format, particularly the routing of information from multiple choices into a single output token—something Transformers handle more naturally via self-attention. This suggests that Zebra-Llama's hybrid use of Mamba/MLA layers may inherit some of this difficulty (inefficient decoding under MMLU’s standard prompting format), especially in tasks like MMLU that rely heavily on understanding structured input-output formatting, rather than open-ended generation or reasoning.

[2] Waleffe, Roger, et al. "An empirical study of mamba-based language models." arXiv preprint arXiv:2406.07887 (2024).

2) Importance of training data selection: We believe that more carefully curated dataset selection can significantly enhance performance, particularly on knowledge-intensive tasks like MMLU. Our Zebra-Llama models were trained using the same datasets as Mamba-In-Llama (which was our main initial baseline)—namely, OpenHermes-2.5, GenQA, and Infinity-Instruct. However, as highlighted in the Llamba paper [3], training on a carefully filtered version of FineWeb-Edu leads to substantial improvements in MMLU scores following distillation. To investigate this further, we trained a Llamba model using our dataset (without FineWeb-Edu). The MMLU score dropped from 38.11 (as reported in the original paper) to 27.35. In contrast, our Zebra-Llama model, trained on the same dataset and using the same training pipeline, significantly outperformed Llamba on MMLU (with the score of 37.1). This result suggests that Zebra-Llama would likely achieve even higher scores than Llamba if trained with similarly curated data, and we consider this an exciting direction for future work. Table: MMLU Comparison with Different Training Datasets

[3] Bick, Aviv, et al. "Llamba: Scaling distilled recurrent models for efficient language processing." arXiv preprint arXiv:2502.14458 (2025).

---

[Answer 2] Extending Evaluation to More Challenging Benchmarks

Thank you for the thoughtful suggestion. We fully agree that evaluating on more challenging benchmarks is important for assessing deeper reasoning capabilities, especially in light of the architectural changes introduced by our hybrid design. As recommended, we have added results on GSM8K for math reasoning and GPQA for graduate-level scientific question answering. These benchmarks help provide a more complete view of the model’s capabilities beyond the core benchmarks in the following table. While we were not able to complete testing on DROP due to time and infrastructure constraints, we do include results on RACE, which is also a well-established benchmark for reading comprehension and reasoning over longer passages. We summarize the additional results below (and we will add them to the final paper):

Key observations from the results:

• Our 1B model excels, achieving the best average score and outperforming the target Llama-3.2-1B-Instruct model with only 7.81% KV cache.
• Our 3B model shows a 1.46% and 0.55% performance drop compared to the target model and X-EcoMLA, respectively, with 21.32x and 2x KV cache compression. Even at 2.01% KV cache, our overall performance surpasses Mamba-In-Llama and Llamba.
• Our 8B model exhibits a 7.11% and 4.2% performance drop compared to the target model and MambaInLlama, with 18.28x and 9.14x KV cache compression, respectively. For the performance gap, the reasoning task GSM8K is identified as a primary contributor. We hypothesize this is due to our efficient training scheme (8B teacher + 11B training tokens). Although Mamba-In-Llama exhibits higher GSM8K score, they adopted a 70B teacher and 20B training tokens. Since math reasoning tasks are generally harder, we expect Zebra-Llama to achieve further improvements with more training tokens.
• Our Zebra-Llama consistantly outperforms Llamba across all model sizes

We sincerely thank the reviewer for the thoughtful feedback and kind recognition of our contributions. In response to your comments, we have done our best to address the raised concerns and will incorporate the corresponding improvements in the final version of the paper. Specifically, we have provided further insights into [Answer 1] the scalability of the proposed workflow across different model sizes, and [Answer 2] its performance on long chain-of-thought (CoT) reasoning tasks.

---

[Answer 1] Scalability of the Proposed Workflow across Smaller/Larger Model Sizes

Thank you for the suggestion. While the core components of our workflow—such as SMART layer selection, distillation, and hybrid architectural design—are scalable and model-agnostic, we focused on 1B, 3B, and 8B models in this work for two main reasons:

Comparison with existing baselines: To the best of our knowledge, the largest post-trained hybrid models available in the literature are at the 8B scale. To ensure meaningful comparisons, we aligned our model sizes with those baselines.
Practical training constraints: The 8B model is the largest that could be reasonably trained on a single node (with 8 decent GPUs). Due to time and resource limitations, we were not able to explore training at larger scales.
However, based on your comment and within the limited time available for the rebuttal, we trained a smaller language model—Qwen-0.5B—to verify the scalability of our workflow to both smaller model sizes and a different model family.

Our results show that the proposed method (referred to as Zebra-Qwen) performs well on models from the Qwen family. For training, we applied the same process as used for the LLaMA series. Notably, the 0.5B Zebra-Qwen model achieved better performance than the original target model while using only 6.25% KV cache.

---

[Answer 2] Performance on Long CoT Reasoning

Thank you for the insightful question and for referencing the paper [1]. The paper shows that small language models (≤3B) can suffer performance degradation when trained on limited long Chain-of-Thought (CoT) data, raising concerns about the effectiveness of such training for small models. To investigate whether our hybrid architecture exacerbates this issue, we conducted a controlled evaluation comparing Zebra-Llama-1B, 8MLA-8M2 on GSM8K, training the model with 8K, 16K, and 32K samples from OpenR1-Math-220k, following a setup similar to the referenced paper. Due to time constraints, our initial evaluation was conducted on a subset comprising the first 150 questions of GSM8K, with training on three distinct dataset sizes: 8k, 16k, and 32k samples. We follow the same evaluation protocol as [1], sampling 4 answers per question (temperature = 0.6, top p = 0.95) and report the average accuracy. Our experiments show that: as Zebra-Llama is further trained with long CoT data, the average response length dramatically increases from approximately 390 to over 4000 tokens. However, this increased "thinking" does not translate to improved accuracy. The GSM8k accuracy decreases from 39.5% to 25.83% when further trained with 8k CoT samples and further declines to 25% with 16k CoT samples. This observation strongly aligns with the hypothesis in [1] that when small Language Models (LLMs) are trained on a limited amount of long CoT data, they tend to learn superficial reasoning patterns. This can lead to an accumulation of errors during the extended "thinking" process, ultimately degrading accuracy. When we increased the CoT training dataset to 32k samples, the reasoning ability of the model starts to recover and the token efficiency increases, which also aligns with the findings in [1].

In summary, we have observed a similar trend that training small LLM on limited long CoT data hurts the reasoning accuracy, despite a marked increase in the response length. We hypothesize this phenomenon is general for small LLMs across various model structures, including hybrid models like Zebra-Llama. We are going to add this insightful analysis to the final version of the paper.

[1] Luo, Renjie, et al. "Through the Valley: Path to Effective Long CoT Training for Small Language Models." arXiv preprint arXiv:2506.07712 (2025).