Jamba

是什么

Jamba 是 AI21 Labs 于2024年3月28日发布的开放权重语言模型。其名称代表 Joint Attention and Mamba^[1][2][4]。它是一种混合 SSM-Transformer 模型，即同时使用两类层：Mamba 状态空间层（state-space layer）和 Transformer 注意力层^[1]。据介绍，它是 Mamba 架构首个达到生产规模的实现^[2]。

工作原理

Transformer 模型使用注意力机制，让输出的每个部分都能参照输入的其他部分。注意力效果好，但随着输入变长，其计算成本迅速增长。Mamba 是一种状态空间设计，它逐步处理序列，并保留一份紧凑的运行摘要，因此成本随长度增长得更为平缓。

Jamba 同时使用两者：每八层中有一层是 Transformer 注意力层，其余为 Mamba 层，并且每层后都接一个多层感知机（MLP）^[1]。多层感知机是一种基础神经网络模块，负责变换每个位置上的数值。这种搭配让模型在需要注意力的任务上保留注意力层，同时由 Mamba 层处理大部分序列。

关键细节

Jamba 总参数为 520 亿，因为采用了混合专家（mixture of experts）层，每个词元只激活其中 120 亿参数^[1][2]。在混合专家模型中，路由器会把每个词元分配给少数几个专门的子网络，因此对于任一词元，只有一部分参数参与计算。

它支持 25.6 万词元（256k）的上下文窗口。在单块 80 GB 显卡上，最多可容纳 14 万词元，AI21 称这在同级别模型中独一无二^[1]。AI21 还声称，在长上下文上，其吞吐量是 Mixtral 8x7B 的3倍^[1]。吞吐量指系统在给定时间内能处理的文本量，这些数字均为 AI21 自己的说法。

在基础模型的基准测试中，Jamba 的成绩为：MMLU 67.4%，HellaSwag 87.1%，ARC Challenge 64.4%，WinoGrande 82.5%，GSM8K（思维链提示）59.9%^[2]。

模型以 Apache 2.0 许可证在 Hugging Face 上发布。它是未经指令微调、也未经安全审核的预训练基础模型，并被列入 NVIDIA API 目录^[1][2]。

这一设计的论文题为《Jamba: A Hybrid Transformer-Mamba Language Model》，2024年3月发布于 arXiv^[3]。

在发展中的位置

Jamba 接替了 Jurassic-2，成为 AI21 的模型系列^[2]。2024年8月，经过指令微调的 Jamba 1.5 Mini 与 Large 随之推出^[2]。