3月29日，AI21公司推出了突破性AI模型Jamba。这是世界上第一个基于SSM-Transformer混合架构的商业大模型。
Transformer架构是现阶段ChatGPT、Stable Difusion、Lyria等产品都在使用的架构，并且在捕捉序列内长距离依赖关系、泛化能力、特征提取等方面表现优异。可是，它在处理长序列、训练大参数模型时存在AI算力消耗大、过拟合、内存占用大等缺点。
AI21的研究人员针对这一情况推出了Jamba。Jamba的发布标志着LLM创新的两个重大里程碑：成功地将Mamba与Transformer架构结合在一起，并将混合SSM-Transformer模型提升到生产级规模和质量。
Jamba加入的SSM技术是借鉴了2023年12月1日，由卡内基梅隆大学Albert Gu和普林斯顿大学Tri Dao提出的论文《Mamba:Linear-Time Sequence Modeling with Selective State Spaces》。Mamba的核心技术是使用“选择性状态空间”来进行序列推理，我们可以把状态空间看作是一个库房。在建模过程中，Mamba可以根据用户输入的具体情况，有选择性地去库房里拿东西或者忽略，这种灵活性使得它能够更好地处理离散型数据。
与传统的序列模型不同的是，Mamba不需要使用复杂的自注意力机制或者MLP块（多层感知器）。主要通过选择性状态空间和MLP块的协同工作，实现了高效的推理过程，并且在处理长序列数据时非常高效，可以轻松处理100万tokens数据。但引入选择性机制后,状态空间模型不再满足时间不变性,所以，无法直接用高效的卷积来计算，Mamba设计了一种“硬件并行算法”。该扫描算法使用了GPU并行加速,充分利用了GPU内存层次结构，控制状态张量的具体化过程,只在更高带宽的内存层(如SMX寄存器)上暂存状态,避免了低效的全局内存访问。这使得模型可以更好地利用GPU效率，不会出现浪费的情况。
Jamba则在Mamba的基础之上进行了创新，采用了分块分层的方法成功融合了SSM和Transformer架构：每个Jamba模块都包含一个注意力层或一个Mamba层，然后是一个多层感知器，总体比例为每八个层中有一个Transformer层。这样可以在保持模型推理性能的前提下，极大的降低了AI算力同时提升吞吐量。例如，与同类的开源模型Mixtral 8x7B相比，Jamba的吞吐量是其3倍。
值得注意的是，Jamba还是一个专家混合模型（MoE）。其520亿参数中的120亿参数长期处于激活状态，并对所有MoE层的专家进行了大幅度优化，有效地减轻了推理时内存占用大的问题。
