首个基于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官网,如有侵权请联系删除