Title: Zebra-Llama: Towards Extremely Efficient Hybrid Models
Venue: NeurIPS 2025
Accepted: yes
Keywords: Model Efficiency

Large Language Models
Abstract: With the growing demand for deploying large language models (LLMs) across diverse applications, improving their inference efficiency is crucial for sustainable and democratized access. However, retraining LLMs to meet new user-specific requirements is prohibitively expensive and environmentally unsustainable. In this work, we propose a practical and scalable alternative: composing efficient hybrid language models from existing pre-trained models.
Our approach, X-EcoMLA, introduces a family of 1B, 3B, and 8B hybrid models by combining State Space Models (SSMs) and Multi-head Latent Attention (MLA) layers, using a refined initialization and post-training pipeline to efficiently transfer knowledge from pre-trained Transformers.
X-EcoMLA achieves Transformer-level accuracy with near-SSM efficiency using only 7–11 billion training tokens (compared to the trillions required for pre-training) and an 8B teacher. Moreover, it dramatically reduces KV cache size—down to 3.9%, 2%, and 2.73% of the original for the 1B, 3B, and 8B variants, respectively—while preserving 100%, 100%, and over 97% of average zero-shot performance on LM Harness tasks.
Compared to models like MambaInLLaMA, X-EcoMLA, Minitron, and Llamba, our approach consistently delivers competitive or superior accuracy while using significantly fewer tokens, smaller teachers, and vastly reduced KV cache memory. Notably, X-EcoMLA-8B surpasses Minitron-8B in few-shot accuracy by 7%, while using 8× fewer training tokens, over 12× smaller KV cache, and a smaller teacher (8B vs. 15B). 
It also achieves 1.4x–3.3x higher throughput (tokens/s) than MambaInLlama. The source code is
released at https://github.com/AMD-AGI/AMD-Hybrid-Models.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces Zebra-Llama, a novel approach for creating efficient hybrid language models by combining existing pre-trained Transformer models with State Space Models and Multi-head Latent Attention layers. The method significantly reduces computational overhead and memory requirements, KV cache size, achieving performance competitive with traditional Transformers while drastically cutting down the resources required for model deployment and training.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\nS1. Zebra-Llama demonstrates remarkable efficiency in reducing KV cache size (up to ~50x compression), significantly improving inference throughput (2.6x-3.8x faster than comparable models).\n\nS2. The intermediate layer distillation and layer selection strategy strategies are methodologically robust, clearly contributing to model effectiveness. Ablation studies clearly demonstrates the effective of each idea (layer selection and intermediate layer distillation).\n\nWeaknesses:\n\nW1. Llamba-8B performs better than Zebra-Llama on MMLU. Could the authors provide insights into this result? It would be helpful to conduct a controlled study with Llamba distillation pipeline using the same training dataset, since Llamba is a pure RNN model and it only has fixed size which has more saving compared with the proposed approach.\n\nW2. Scaling to larger models and different model families: The paper could be improved by exploring different models, such as Qwen3, or by scaling to larger model sizes.\n\nQUESTIONS:\nAny insights about why split $W_{UKV}$ in such way?\n\nLIMITATIONS:\n1. **Most importantly**, Llamba-8B performs better than Zebra-Llama on MMLU. Could the authors provide insights into this result? It would be helpful to conduct a controlled study with Llamba distillation pipeline using the same training dataset, since Llamba is a pure RNN model and it only has fixed size which has more saving compared with the proposed approach.\n\n2. Scaling to larger models and different model families: The paper could be improved by exploring different models, such as Qwen3, or by scaling to larger model sizes.\n\n3. Consider evaluating on more challenging tasks, such as GSM8K and MATH, to better demonstrate the effectiveness of the proposed approach.\n\nI can increase my score if some of these concerns are addressed.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces Zebra-Llama, a hybrid language model designed for high inference efficiency. The hybrid architecture is composed from existing pre-trained Transformers. Through a post-training pipeline a new architecture is created that interleaves Mamba2 layers with Multi-head Latent Attention (MLA) layers. This process involves three key stages: refined weight initialization from a teacher model, intermediate layer distillation to align representations, and a sensitivity-aware layer selection strategy called SMART that places the MLA layers where they are most necessary. The authors demonstrate that this method dramatically reduces KV cache size and improves throughput, while maintaining the performance of the original teacher model.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n- The paper is generally well-written and easy to follow.\n- The approach to construct the hybrid model is intuitive and the KV cache reduction is significant.\n- The authors strengthen their claims by reporting direct inference throughput metrics.\n\n**Weaknesses:**\n- **Limited Novelty**: While the hybrid architecture seems to yield meaningful efficiency gains, the novelty of the approach is limited. The use of distillation to transfer knowledge from a pre-trained transformers to hybrid architectures is common practice and upcycling pre-trained attention into MLA to enable extreme KV cache compression was already proposed in X-EcoMLA, reducing the novelty of this paper's architectural contributions.\n\n- **Sensitivity Measure-Aware Replacement Method:** The Sensitivity Measure-Aware Replacement of Transformer Layers (SMART) strategy is central to the paper's contribution but is presented as a heuristic and is not well-motivated. This raises a major concern that the strategy may be overfitted to the distilled Llama architecture which is the only architecture family considered in this work. To substantiate the claim that SMART is a generally applicable technique, its effectiveness should be demonstrated on other architectural families, such as Qwen.\n\n- **Incomplete Results**: The few-shot performance is only shown for the 8B model. Why not reporting the results for all models?\n\nMinor points:\n*   Line 92: should be \"targeting extreme efficiency\".\n*   Table 1 Caption: There is a typo in the word \"except\".\n*   Line 307: The word \"accuracy\" is repeated.\n\nQUESTIONS:\nThe results in Table 1 indicate that the performance gap between the distilled hybrid models and the teacher baselines is larger for bigger models. Could you elaborate on why this might be the case?\n\nLIMITATIONS:\nThe authors sufficiently discuss the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-29"}

comments: SUMMARY:
This paper studies the composition of efficient hybrid language models from existing pre-trained models. It proposes a pipeline (including weight initialization, intermediate layer distillation, the SMART layer selection strategy, etc.) to mix SSM and MLA layers, in order to greatly improve inference efficiency while maintaining foundational capabilities without significant impact.

STRENGTHS AND WEAKNESSES:
#### **Strengths**

- The proposed method is very effective, capable of greatly reducing inference overhead while keeping performance almost unchanged. For instance, the 8B model achieves an impressive 36.6x KV cache compression with only a 2.8% drop in average zero-shot performance.
    
- The experimental evaluation is relatively comprehensive, and the ablation analysis, which validates key designs like the novel SMART strategy, is solid
    

#### **Weaknesses**

- The paper requires more discussion regarding the specific performance differences between the accelerated hybrid model and the original model. For example, while the average score is high, a closer look at Table 1 reveals the 8B model has a notable performance drop on the MMLU benchmark, and a discussion of such specific trade-offs would be valuable.
    
- In terms of capability measurement, the benchmarks used in Table 1 are not comprehensive enough. These more challenging benchmarks are crucial for probing deeper reasoning capabilities that may be affected by the architectural changes. I recommend further testing on benchmarks such as DROP (for reading comprehension), AIME (for math reasoning), GPQA (for graduate-level scientific QA), etc.

QUESTIONS:
See Weaknesses

LIMITATIONS:
None
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-30

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors introduce Zebra-Llama, a hybrid LLM designed to significantly reduce KV cache size while preserving the performance of the original pre-trained model on one-shot and few-shot benchmarks.\n\nThe proposed workflow consists of four main steps:\n\n- First, separate Multi-head Latent Attention (MLA) and Mamba2 models are initialized using weights from the pre-trained LLM.\n\n- Next, this initialization is refined by aligning the intermediate representations of the new models with those of the original.\n\n- The best layer-mixing scheme is then identified to create an optimal hybrid architecture.\n\n- Finally, the mixed model is further fine-tuned using Supervised Fine-Tuning (SFT) and Direct Preference Optimization (DPO).\n\nThis hybrid model achieves a significant reduction in test-time computation and outperforms other competitive hybrid models in few-shot settings. Further ablation studies validate the layer selection strategy, analyze the trade-offs between model configurations, and examine scaling behaviors as the size of the teacher model changes.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: \n\n- The paper is well written and quite involving. I appreciate the effort the authors made to deliver the implementation details. \n- The authors report a comprehensive series of experimental results and ablation studies. The use of SMART allocation scheme is also justified by the ablation study. \n\nWeakness: \n- It seems that the authors are using 1B/3B/8B models for the proposed scheme, how does the workflow perform when the model size is extremely large/small? \n- A typo in line 180. \"MLA initialization\" -> \"Mamba2 initialization\"\n\nQUESTIONS:\nGiven that long Chain-of-Thought (CoT) reasoning is susceptible to compounding errors [1], how does Zebra-Llama's performance on such tasks compare to its base model? Does its hybrid architecture of attention and state-space layers potentially exacerbate error propagation in multi-step reasoning?\n\n[1] Luo, Renjie, et al. \"Through the Valley: Path to Effective Long CoT Training for Small Language Models.\" arXiv preprint arXiv:2506.07712 (2025).\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-02"}