Jamba: A Hybrid Transformer-Mamba Language Model
基本信息
- 作者：Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz, Mor Zusman, Yoav Shoham
- 提交时间：2024年3月28日（v1）；最后修订于2024年7月3日（v2）
- 分类：计算机科学 > 计算与语言（cs.CL）；机器学习（cs.LG）
- arXiv ID：2403.19887
- DOI：10.48550/arXiv.2403.19887
摘要
本文提出了 Jamba，一种基于新型混合 Transformer-Mamba 混合专家（Mixture-of-Experts, MoE） 架构的基础大语言模型。具体而言，Jamba 交错排列 Transformer 层和 Mamba 层，从而兼得两类模型家族的优势。在该架构的部分层中引入 MoE，以在控制活跃参数规模的同时提升模型容量。这种灵活的架构支持针对不同资源和目标进行特定配置。
在具体实现中，该模型可部署于单张 80GB GPU 上。与标准 Transformer 相比，Jamba 在大规模训练下具备更高的吞吐量和更小的内存占用，同时在标准语言模型基准和长上下文评估中取得了最先进的性能。值得注意的是，该模型在高达 256K token 的上下文长度下仍表现出色。
论文还研究了多项架构设计决策，例如如何组合 Transformer 层与 Mamba 层、如何进行专家混合，并表明其中一些决策对大规模建模至关重要。此外，作者描述了在训练和评估 Jamba 过程中发现的该架构的若干有趣特性，并计划发布多个消融实验的检查点，以鼓励对该新型架构的进一步探索。模型权重已在宽松许可证下公开发布。
关键要点
- 提出 混合 Transformer-Mamba 架构，交错堆叠两类层
- 引入 稀疏专家混合（MoE） 以提升模型容量
- 支持最长 256K token 的上下文长度
- 适配单张 80GB GPU
- 相比标准 Transformer 具有 更高吞吐量 和 更小内存占用
- 在标准基准和长上下文任务上达到 最先进性能
- 权重已公开，并计划发布消融实验检查点
相关链接