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, Zebra-Llama, 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. Zebra-Llama achieves Transformer-level accuracy with near-SSM efficiency using only 7-11B training tokens (compared to trillions of tokens required for pre-training) and an 8B teacher. Moreover, Zebra-Llama 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 >97% of average zero-shot performance on LM Harness tasks. Compared to models like MambaInLLaMA, X-EcoMLA, Minitron, and Llamba, Zebra-Llama consistently delivers competitive or superior accuracy while using significantly fewer tokens, smaller teachers, and vastly reduced KV cache memory. Notably, Zebra-Llama-8B surpasses Minitron-8B in few-shot accuracy by 7% while using 8x fewer training tokens, over 12x smaller KV cache, and a smaller teacher (8B vs. 15B). It also achieves 2.6x-3.8x higher throughput (tokens/s) than MambaInLlama up to a 32k context length. We will release code and model checkpoints upon acceptance.

翻译：随着在多样化应用中部署大语言模型（LLM）的需求日益增长，提升其推理效率对于实现可持续且普惠的访问至关重要。然而，为满足新的用户特定需求而重新训练LLM成本极其高昂，且在环境上不可持续。在本工作中，我们提出了一种实用且可扩展的替代方案：利用已有的预训练模型组合构建高效的混合语言模型。我们的方法Zebra-Llama通过结合状态空间模型（SSMs）与多头潜在注意力（MLA）层，引入了一个包含1B、3B和8B参数的混合模型系列，并采用精炼的初始化与后训练流程，以高效地从预训练的Transformer模型中迁移知识。Zebra-Llama仅使用7-11B训练词元（相较于预训练所需的万亿级词元）和一个8B教师模型，便达到了Transformer级别的准确度，同时具备接近SSM的效率。此外，Zebra-Llama显著降低了KV缓存大小——其1B、3B和8B变体分别降至原始大小的3.9%、2%和2.73%——同时在LM Harness任务上保持了100%、100%和>97%的平均零样本性能。与MambaInLLaMA、X-EcoMLA、Minitron和Llamba等模型相比，Zebra-Llama在准确度上持续提供具有竞争力或更优的结果，同时使用的训练词元显著更少，教师模型更小，且KV缓存内存大幅降低。值得注意的是，Zebra-Llama-8B在少样本准确度上超越了Minitron-8B达7%，同时使用的训练词元减少了8倍，KV缓存缩小了超过12倍，且教师模型更小（8B vs. 15B）。在高达32k的上下文长度下，其吞吐量（词元/秒）也比MambaInLlama高出2.6倍至3.8倍。我们将在论文被接受后发布代码和模型检查点。