Mamba超强进化体一举颠覆Transformer！单张A100跑140K上下文

1,106 0 0

文章摘要

【关键词】 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模型更大的输出，在语言生成任务上也更优。