用少量训练数据融合Transformer与状态空间模型,实现高效低内存的混合语言模型。
Zebra-Llama: Towards Extremely Efficient Hybrid Models
- 通过融合SSM与MLA层,从预训练Transformer迁移知识,仅需7-11B token微调。
- 在1B~8B规模下,KV缓存压缩至原大小的2%~3.9%,零样本性能保持97%以上。
- 相比同类模型,训练量少8倍、缓存小12倍,推理速度提升2.6倍以上。
随着大语言模型(LLMs)在多样化应用中的部署需求增长,提升推理效率对实现可持续、普惠化的访问至关重要。然而,为满足新用户需求而重新训练模型成本过高且环境不可持续。本文提出一种实用且可扩展的替代方案:从已有预训练模型构建高效混合语言模型。所提方法Zebra-Llama通过结合状态空间模型(SSMs)与多头隐式注意力(MLA)层,采用优化初始化和后训练流程,高效迁移预训练Transformer的知识,生成1B、3B、8B规模的混合模型。Zebra-Llama在仅使用7-11B训练标记(tokens)及8B教师模型的情况下,达到接近Transformer的精度,并实现近似SSM的推理效率。其KV缓存分别缩减至原始大小的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倍。代码与模型检查点将在论文接受后发布。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。