Jamba: A Hybrid Transformer-Mamba Language Model

我们提出了Jamba，这是一个基于新型混合Transformer-Mamba专家混合（MoE）架构的新型基础大型语言模型。具体而言，Jamba交错了Transformer和Mamba层的块，享受两个模型家族的好处。在其中一些层中添加MoE以增加模型容量，同时保持活跃参数的使用可管理性。这种灵活的架构允许资源和目标特定的配置。在我们实现的特定配置中，我们得到了一个强大的模型，适合单个80GB GPU。在大规模构建的情况下，与普通Transformer相比，Jamba提供了高吞吐量和小的内存占用，并且在标准语言模型基准和长上下文评估方面具有最先进的性能。值得注意的是，该模型在长达256K个标记的上下文长度方面呈现出强大的结果。我们研究了各种架构决策，例如如何组合Transformer和Mamba层，以及如何混合专家，并表明其中一些在大规模建模中至关重要。我们还描述了这些架构的几个有趣特性，这些特性在Jamba的训练和评估中已经显现，并计划发布各种消融运行的检查点，以鼓励进一步探索这种新型架构。我们以宽松的许可证公开发布我们实现的Jamba权重。

Jamba论文的目标是提出一种新的基于Transformer-Mamba混合模型的大型语言模型，以解决当前大型语言模型的资源和性能问题。

Jamba模型采用了一种新的混合Transformer-Mamba模型结构，并在其中加入了Mixture-of-Experts机制以增加模型容量，同时保持参数使用的可管理性。这种灵活的架构允许根据资源和目标进行特定的配置，提供了比传统Transformer更高的吞吐量和更小的内存占用，同时在标准的语言模型基准测试和长上下文评估中获得了最先进的性能。

Jamba模型在长达256K令牌上下文长度的情况下表现出强大的结果。论文还研究了Transformer和Mamba层如何组合以及如何混合专家等多种架构决策，并展示了这些决策在大规模建模中的重要性。此外，论文还公开了Jamba模型的权重，并计划公开各种消融运行的检查点，以鼓励进一步探索这种新型架构。

最近在这个领域中，还有一些相关的研究，例如GPT-3、Turing-NLG等。