Mamba时代来了？自2017年开创性研究论文《Attention is All You Need》问世以来，transformer架构就一直主导着生成式人工智能领域。然而，transformer架构实际上有两个显著缺点：内存占用大和推理速度慢。这使得在没有大量硬件资源的情况下运行长上下文窗口或大量并行批处理变得具有挑战性，从而限制了广泛的实验和部署。随着上下文长度的增加，推理速度会变慢，因为每个token都依赖于它之前的整个序列，将长上下文用例置于高效生产的范围之外。
然而，transformer并不是生成式人工智能唯一的前进方向。最近，AI21 Labs推出并开源了一种名为“Jamba”的新方法，在多个基准上超越了transformer。Mamba的SSM架构可以很好地解决transformer的内存资源和上下文问题。然而，Mamba方法很难提供与transformer模型相同的输出水平。Jamba将基于结构化状态空间模型（SSM）的Mamba模型与transformer架构相结合，旨在将SSM和transformer的最佳属性结合在一起。
Jamba还可以作为NVIDIA NIM推理微服务从NVIDIA API目录进行访问，企业应用程序开发人员可以使用NVIDIA AI Enterprise软件平台进行部署。Jamba模型具有以下特点：第一个基于Mamba的生产级模型，采用新颖的SSM-Transformer混合架构；与Mixtral 8x7B相比，长上下文上的吞吐量提高了3倍；提供对256K上下文窗口的访问；公开了模型权重；同等参数规模中唯一能够在单个GPU上容纳高达140K上下文的模型。
模型架构如下图所示，Jamba的架构采用块层（blocks-and-layers）方法，使Jamba能够集成两种架构。每个Jamba块包含一个注意力层或一个Mamba层，后跟一个多层感知器（MLP），从而形成transformer层。Jamba利用MoE来增加模型参数的总数，同时简化推理中使用的活跃参数的数量，从而在计算需求没有相应增加的情况下获得更高的模型容量。为了在单个80GB GPU上最大限度地提高���型的质量和吞吐量，研究团队优化了所使用的MoE层和专家的数量，为常见推理工作负载留出了足够的内存。Jamba的MoE层允许它在推理时仅利用可用的52B参数中的12B，并且其混合架构使这些12B活跃参数比同等大小的纯transformer模型更有效。此前，没有人将Mamba扩展到3B参数之外。Jamba是同类模型中第一个达到生产级规模的混合架构。
吞吐量和效率初步评估实验表明，Jamba在吞吐量和效率等关键衡量指标上表现出色。在效率方面，Jamba在长上下文上的吞吐量达到了Mixtral 8x7B的3倍。Jamba比Mixtral 8x7B等大小相当的基于Transformer的模型更高效。在成本方面，Jamba可以在单个GPU上容纳140K上下文。与当前类似大小的其他开源模型相比，Jamba能提供更多的部署和实验机会。需要注意的是，Jamba目前不太可能取代当前基于Transformer的大型语言模型（LLM），但它可能会成为某些领域的补充。