Jamba: A Hybrid Transformer-Mamba Language Model

Jamba: A Hybrid Transformer-Mamba Language Model

相关链接：arXiv
关键字：hybrid architecture、Transformer、Mamba、mixture-of-experts (MoE)、language model

摘要

我们介绍了Jamba，一种新的基于新颖混合Transformer-Mamba混合专家（MoE）架构的基础大型语言模型。具体来说，Jamba交错了Transformer和Mamba层的块，享受两个模型家族的好处。在这些层中的一些层添加了MoE，以增加模型容量，同时保持活跃参数使用量可控。这种灵活的架构允许针对资源和目标的特定配置。在我们实现的特定配置中，我们得到了一个强大的模型，适合在单个80GB GPU中运行。Jamba在大规模构建，相比传统的Transformer，提供了高吞吐量和小内存占用，并且同时在标准语言模型基准测试和长上下文评估中表现出最先进的性能。值得注意的是，该模型对长达256K个token的上下文呈现出强大的结果。我们研究了各种架构决策，例如如何结合Transformer和Mamba层，以及如何混合专家，并且展示了它们在大规模建模中的重要性。我们还描述了在训练和评估Jamba过程中揭示的这些架构的几个有趣特性，并且计划公开不同消融运行的检查点，以鼓励对这个新颖架构的进一步探索。我们根据宽松的许可证公开了我们实现的Jamba的权重。

核心方法

1. 混合架构：Jamba结合了Transformer层和Mamba层，这是一种最新的状态空间模型，以及混合专家（MoE）组件。
2. 灵活配置：通过调整Transformer层和Mamba层的比例，可以在内存使用、高效训练和长上下文能力之间进行平衡。
3. MoE层：MoE层允许在不增加计算需求的情况下增加模型容量，通过在每两层中应用MoE，使用16个专家，每个token使用前2个专家。
4. 训练和评估：Jamba在多种基准测试中进行了评估，显示出与具有相似参数数量的Mixtral-8x7B相当的性能，并且在长上下文评估中超过了Mixtral。

实验说明

Jamba模型在一个80GB的GPU上进行了实现和训练，以适应单个GPU的内存限制，同时在质量和吞吐量方面实现了最佳性能。Jamba的特定配置选择了适当的层数、Transformer和Mamba层的比例、MoE的使用频率以及每个MoE层的专家数量，以确保模型可以在单个GPU上运行，同时具有足够的内存来处理输入。

结论

Jamba是一种新颖的架构，它结合了注意力和Mamba层，以及MoE模块，并且我们公开了它的实现，达到了最先进的性能并支持长上下文。我们展示了Jamba在平衡性能和内存需求方面提供的灵活性，同时保持高吞吐量。我们对几种设计选择进行了实验，例如注意力层和Mamba层的比例，并讨论了在开发过程中所做的一些发现，这些发现将为未来的混合注意力-状态空间模型的工作提供信息。为了促进这样的研究，我们计划公开小规模训练运行的模型检查点。我们在此发布的最大的模型具有12B活跃参数和52B可用参数，支持长达256K个token的上下文，并且即使在处理140K-token文本时也能适合单个80GB GPU。