首个基于SSM-Transformer混合架构，开源商业大模型Jamba

3月29日，知名AI研究实验室AI21在官网开源了，首个基于SSM-Transformer混合架构的商业大模型——Jamba。

目前，ChatGPT、Stable Difusion 、Lyria等产品使用的皆是Transformer架构，虽然在捕捉序列内长距离依赖关系、泛化能力、特征提取等方面非常优秀，但在处理长序列、训练大参数模型时存在AI算力消耗大、过拟合、内存占用大等缺点。

Jamba则是在传统的Transformer架构之上，加入了结构化状态空间模型 (SSM) 技术，结合二者的优点来极大提升其性能。例如，Jamba的吞吐量是同类知名开源模型Mixtral 8x7B的3倍；也是同类模型中极少数能在单个GPU上容纳高达140K上下文的模型。

这也就是说，那些没有庞大GPU集群的小企业和个人开发者，通过Jamba也能开发出高性能、低消耗的生成式AI产品。

huggingface地址：https://huggingface.co/ai21labs/Jamba-v0.1

也可通过英伟达平台使用：https://developer.nvidia.com/blog/nvidia-nim-offers-optimized-inference-microservices-for-deploying-ai-models-at-scale/

Jamba架构简单介绍

Jamba加入的SSM技术是借鉴了2023年12月1日，由卡内基梅隆大学Albert Gu和普林斯顿大学Tri Dao提出的论文《Mamba: Linear-Time Sequence Modeling with Selective State Spaces》。

论文地址：https://arxiv.org/abs/2312.00752

Mamba的核心技术是使用“选择性状态空间”来进行序列推理，我们可以把状态空间看作是一个库房。

在建模过程中，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层的专家进行了大幅度优化，减轻了推理时内存占用大的问题。

AI21介绍

公开资料显示，AI21 Labs创立于2017年，总部位于特拉维夫和纽约。由Yoav Shoham、Ori Goshen和Amnon Shashua三人联合创立。其中Yoav是一位连续创业者，曾将两家企业出售给谷歌并担任斯坦福大学名誉教授；

Ori是希伯来大学的副教授，曾参与多个NLP项目并发表了几十篇学术论文；Amnon是希伯来大学的教授，同时是知名自动驾驶公司Mobileye的联合创始人兼CEO。

仅2023年，AI21一共融资了2.08亿美元，目前总融资额度3.26亿美元。

产品方面，AI21 Labs在2023年3月推出了大语言模型Jurassic-2，包含Large、Grande和Jumbo三种模型。

Jurassic-2除了在文本生成、API延迟、语言支持等全面增强之外，还开放了指令微调、数据微调，帮助企业、个人开发者打造量身定制的ChatGPT助手。

Jurassic-2特定微调的类型共包括语义搜索，了解查询的意图和上下文含义，并从文档中检索相关的文本片段；上下文问答，仅根据特定上下文提供答案，也可以从文档库中自动检索等。

目前，耐克、Zoom、沃尔玛、三星、阿迪达斯、airbnb等知名企业在使用AI21的大模型产品。

本文素材来源Jamba官网，如有侵权请联系删除

END