文章摘要
【关 键 词】 AI21 Labs、Mamba架构、Jamba模型、Transformer、开源模型
AI21 Labs最近开源了Jamba,这是世界上第一个生产级的Mamba大模型,它在多项基准测试中表现出色,与目前最强的几个开源Transformer模型相媲美。
Jamba是基于全新的SSM-Transformer混合架构的首个生产级Mamba模型,其长文本处理吞吐量比Mixtral 8x7B提高了3倍,实现了256K超长上下文窗口,是同等规模中唯一一个能在单张GPU上处理140K上下文的模型。
Jamba的发布标志着LLM的两个重要里程碑:一是成功将Mamba与Transformer架构相结合,二是将新形态的模型(SSM-Transformer)成功提升到了生产级的规模和质量。尽管大家都认识到了Transformer架构存在的两个主要缺点:内存占用量大,推理速度会随着上下文的增长而变慢,但当前性能最强的大模型全是基于Transformer的。
Jamba通过Joint Attention和Mamba架构,整合了Transformer、Mamba、以及专家混合(MoE)的优势,同时优化了内存、吞吐量和性能。Jamba是第一个达到生产级规模(52B参数)的混合架构。AI21的Jamba架构采用blocks-and-layers的方法,使Jamba能够成功集成这两种架构。每个Jamba块都包含一个注意力层或一个Mamba层,然后是一个多层感知器(MLP)。
Jamba的另一个特点是利用MoE来增加模型参数的总数,同时简化推理中使用的活动参数的数量,从而在不增加计算要求的情况下提高模型容量。为了在单个80GB GPU上最大限度地提高模型的质量和吞吐量,研究人员优化了使用的MoE层和专家的数量,为常见的推理工作负载留出足够的内存。对比Mixtral 8x7B等类似大小的基于Transformer的模型,Jamba在长上下文上做到了3倍的加速。
Jamba的独特之处在于它同时结合了两种模型架构——Transformer和状态空间模型SSM。Transformer是复杂推理任务的首选架构,其最核心的定义特征就是「注意力机制」。而SSM结合了早前AI模型的多个优点,例如递归神经网络和卷积神经网络,因此能够实现长序列数据的处理,且计算效率更高。虽然SSM有自己的局限性,但一些早期的代表,比如由普林斯顿和CMU提出的Mamba,就可以处理比Transformer模型更大的输出,在语言生成任务上也更优。
原文和模型
【原文链接】 阅读原文 [ 1677字 | 7分钟 ]
【原文作者】 新智元
【摘要模型】 gpt-4-32k
【摘要评分】 ★★★★☆