Zebra-Llama：迈向极高效的混合模型
机器学习
2026-01-21 v2 计算与语言
摘要
随着将大语言模型（LLMs）部署到多样化应用的需求不断增长，提高推理效率对于可持续和民主化的访问至关重要。然而，针对新的用户特定需求重新训练LLM的成本极高且在环境上不可持续。在这项工作中，我们提出了一种实用且可扩展的替代方案：由现有预训练模型组合高效混合语言模型。我们的方法Zebra-Llama通过结合状态空间模型（SSM）和多头潜在注意力（MLA）层，引入了一个包含1B、3B和8B变体的混合模型家族，使用经过改进的初始化和后训练流程，高效地从预训练Transformer迁移知识。Zebra-Llama仅使用7-11B训练token（相比预训练所需的数万亿token）和一个8B教师模型，实现了Transformer级别的准确率和接近SSM的效率。此外，Zebra-Llama大幅减小了KV缓存大小——分别降至原始的3.9%、2%和2.73%（对应1B、3B和8B变体），同时在LM Harness任务上保持了100%、100%和>97%的平均零样本性能。与MambaInLLaMA、X-EcoMLA、Minitron和Llamba等模型相比，Zebra-Llama在显著更少的训练token、更小的教师模型和大幅减少的KV缓存内存条件下，始终提供具有竞争力或更优的准确率。值得注意的是，Zebra-Llama-8B在少样本准确率上超越Minitron-8B达7%，同时使用了8倍的训练token、超过12倍更小的KV缓存以及更小的教师模型（8B对比15B）。它还实现了高达2.6-3.8倍的吞吐量（token/s），上下文长度可达32k。我们将在接收后发布代码和模型检查点。
引用
@article{arxiv.2505.17272,
title = {Zebra-Llama: Towards Extremely Efficient Hybrid Models},
author = {Mingyu Yang and Mehdi Rezagholizadeh and Guihong Li and Vikram Appia and Emad Barsoum},
journal= {arXiv preprint arXiv:2505.17272},
year = {2026}
}