AI21 Labs 发布 Mamba 与 Transformer 混合模型 Jamba

背景

这一时期的许多语言模型都基于 Transformer 架构。它的注意力机制让输入的每个部分都能与其他部分相互参照。注意力效果好，但输入越长，计算成本上升得越快。Mamba 是一种较新的替代设计，它逐步处理序列，并保留一份紧凑的内容摘要，因此成本随长度增长得更为平缓。AI21 Labs 此前推出过 Jurassic 系列模型，此次发布的 Jamba 则把两种设计结合在一起。

发生了什么

2024年3月28日，AI21 Labs 发布 Jamba，它把 Mamba 状态空间层与 Transformer 注意力层结合在一起^[1]。公司将其称为首个达到生产规模的、基于 Mamba 的模型^[2]。模型权重以 Apache 2.0 许可证在 Hugging Face 上发布^[1]。

关键细节

Jamba 总参数为 520 亿，由于采用混合专家（mixture of experts）层，每个词元实际激活 120 亿参数^[1]。混合专家层会把每个词元分配给少数几个专门的子网络，因此模型存储的参数远多于单个词元所用的参数。

模型支持 25.6 万词元（256k）的上下文窗口。在单块 80 GB 显卡上，最多可容纳 14 万词元^[1]。AI21 称，在长上下文上，其吞吐量是 Mixtral 8x7B 的3倍^[1]。这些都是 AI21 自己公布的数字。

按 AI21 的描述，模型每八层中有一层是 Transformer 注意力层，其余为 Mamba 层，每层后各接一个多层感知机（MLP）。注意力层让模型能直接关联文本中相距较远的部分，Mamba 层则以较低成本承担大部分逐步处理序列的工作。

影响

此次发布把一种混合架构以开放权重的形式交给了公众，他人可以下载并测试。它的意义在于这种组合：一个生产规模的模型同时使用两类层，而不是只依赖注意力。2024年8月，经过指令微调的 Jamba 1.5 Mini 与 Large 随之推出^[2]。本条只涵盖首次发布，后续版本见各自的条目。