Research map / Architectures and efficiency / Attention and transformers
Zebra-Llama: Towards Extremely Efficient Hybrid Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Composes 1B/3B/8B hybrid models from pre-trained Llama and Qwen Transformers by SVD-initializing MLA and Mamba2 layers, intermediate-layer distillation, sensitivity-based layer placement (SMART), and SFT+DPO on about 7-11B tokens, achieving large KV cache reductions with near-base accuracy.
Problem
Making LLM inference memory- and compute-efficient (small KV cache, high throughput) without the prohibitive cost of pre-training new architectures from scratch or retraining per user requirement.
Approach
Build hybrid models interleaving MLA and Mamba2 layers by upcycling a pre-trained Transformer: SVD-based MLA initialization, Mamba2 initialization from attention weights, intermediate layer distillation to refine both, then choose MLA layer positions via a sensitivity-aware rule (SMART) before end-to-end distillation and DPO.
Key results
- Zebra-Llama-1B (8MLA-8M2) reaches 52.94 zero-shot average vs 51.87 for Llama3.2-1B-Instruct with 7.81% KV cache.
- Zebra-Llama-8B (16MLA-16M2) reaches 64.97 zero-shot average vs 65.53 for Llama3.1-8B-Instruct with 5.47% KV cache, using 11B tokens and an 8B teacher.
- Zebra-Llama-8B achieves about 3.9x throughput over Llama-3.1-8B at 16k input tokens with 17% of peak memory on a single MI300X.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| compose 8 MLA and 8 Mamba2 layers | Zebra-Llama-1B · GSM8K 8-shot | accuracy: 38.51 -> 41.09 | Helped |
| increase MLA layers from 8 to 12 (ablation) | Zebra-Llama-1B · GSM8K 8-shot | accuracy: 41.09% to 29.57% | Hurt |
| select MLA layers with SMART rather than Uniform #1 (ablation) | Zebra-Llama-1B · SMART layer selection | average score: 48.84 -> 50.15 | Helped |
| select MLA layers with SMART rather than Uniform #2 (ablation) | Zebra-Llama-1B · SMART layer selection | average score: 49.72 -> 50.15 | Helped |
| select MLA layers with SMART rather than greedy maximum sensitivity (ablation) | Zebra-Llama-1B · SMART layer selection | average score: 48.98 -> 50.15 | Helped |
| use 8B rather than 1B teacher (ablation) | Zebra-Llama-1B 4MLA-12M2 · LM Harness zero-shot | average accuracy: 49.11 -> 52.03 | Helped |
| use 70B rather than 8B teacher (ablation) | Zebra-Llama-1B 4MLA-12M2 · LM Harness zero-shot | average accuracy: 52.03 -> 52.44 | Helped |
| use 70B rather than 8B teacher (ablation) | Zebra-Llama-3B 8MLA-20M2 · LM Harness zero-shot | average accuracy: 59.39 -> 59.37 | No clear effect |
| use 70B rather than 8B teacher (ablation) | Zebra-Llama-8B 8MLA-24M2 · LM Harness zero-shot | average accuracy: 63.08 -> 63.48 | Helped |
| compose 14 MLA and 14 Mamba2 layers | Zebra-Llama-3B · GSM8K 8-shot | accuracy: 62.77 | Mixed |
Surprising findings
- Zebra-Llama-1B (8MLA-8M2) slightly exceeds its Llama3.2-1B-Instruct teacher on zero-shot average despite 7.81% KV cache.
- Increasing the number of MLA layers to 12 in the 1B model drops GSM8K accuracy from 41.09 to 29.57.
Limitations
- Relies on a strong teacher model; teacher-free alternatives are unexplored.
- Trained with maximum sequence length 2048, so long-context performance (RULER) is limited.
- Performance gap on GSM8K for the 8B variants relative to the base model (about 7.1%).
Experimental setup
Datasets: OpenHermes-2.5 GenQA Infinity-Instruct Llama3-ultrafeedback orca_dpo_pairs ultrafeedback_binarized
Benchmarks: ARC-Challenge ARC-Easy HellaSwag MMLU OpenBookQA PIQA RACE WinoGrande GSM8K GPQA RULER TruthfulQA
Metrics: accuracy normalized accuracy KV cache size ratio throughput (tokens/s) peak memory
Baselines: MambaInLlama X-EcoMLA Llamba Minitron Llama3.2-1B-Instruct Llama3.2-3B-Instruct Llama3.1-8B-Instruct
Builds on: MambaInLlama X-EcoMLA MOHAWK Multi-head Latent Attention Mamba2
Closest papers by meaning
- Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning (ICLR 2024 · similarity 0.89)
- $\texttt{Model-GLUE}$: Democratized LLM Scaling for A Large Model Zoo in the Wild (NeurIPS 2024 · similarity 0.88)
- ChatQA 2: Bridging the Gap to Proprietary LLMs in Long Context and RAG Capabilities (ICLR 2025 · similarity 0.87)
- The Mamba in the Llama: Distilling and Accelerating Hybrid Models (NeurIPS 2024 · similarity 0.87)
- Nemotron-Flash: Towards Latency-Optimal Hybrid Small Language Models (NeurIPS 2025 · similarity 0.87)
- LLaMA-Omni: Seamless Speech Interaction with Large Language Models (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.