Mamba超强进化体一举颠覆Transformer！单张A100跑140K上下文
发布时间:2024-03-29 11:18:00 来自:网易科技
新智元报道
编辑：编辑部
52B 的生产级 Mamba 大模型来了！这个超强变体 Jamba 刚刚打破世界纪录，它能正面硬刚 Transformer，256K 超长上下文窗口，吞吐量提升 3 倍，权重免费下载。
之前引爆了 AI 圈的 Mamba 架构，今天又推出了一版超强变体！
人工智能独角兽 AI21 Labs 刚刚开源了 Jamba，世界上第一个生产级的 Mamba 大模型！
Jamba 在多项基准测试中表现亮眼，与目前最强的几个开源 Transformer 平起平坐。
特别是对比性能最好的、同为 MoE 架构的 Mixtral 8x7B，也互有胜负。
具体来说它——
-
是基于全新 SSM-Transformer 混合架构的首个生产级 Mamba 模型
-
与 Mixtral 8x7B 相比，长文本处理吞吐量提高了 3 倍
-
实现了 256K 超长上下文窗口
-
是同等规模中，唯一一个能在单张 GPU 上处理 140K 上下文的模型
-
以 Apache 2.0 开源许可协议发布，开放权重
之前的 Mamba 因为各种限制，只做到了 3B，还被人质疑能否接过 Transformer 的大旗，而同为线性 RNN 家族的 RWKV、Griffin 等也只扩展到了 14B。
——Jamba 这次直接干到 52B，让 Mamba 架构第一次能够正面硬刚生产级别的 Transformer。
Jamba 在原始 Mamba 架构的基础上，融入了 Transformer 的优势来弥补状态空间模型（SSM）的固有局限性。
可以认为，这实际上是一种新的架构——Transformer 和 Mamba 的混合体，最重要的是，它可以在单张 A100 上运行。
它提供了高达 256K 的超长上下文窗口，单个 GPU 就可以跑 140K 上下文，而且吞吐量是 Transformer 的 3 倍！
与 Transformer 相比，看 Jamba 如何扩展到巨大的上下文长度，非常震撼
Jamba 采用了 MoE 的方案，52B 中有 12B 是活跃参数，目前模型在 Apache 2.0 下开放权重，可以在 huggingface 上下载。
模型下载：https://huggingface.co/ai21labs/Jamba-v0.1
LLM 新里程碑
Jamba 的发布标志着 LLM 的两个重要里程碑：
一是成功将 Mamba 与 Transformer 架构相结合，二是将新形态的模型（SSM-Transformer）成功提升到了生产级的规模和质量。
当前性能最强的大模型全是基于 Transformer 的，尽管大家也都认识到了 Transformer 架构存在的两个主要缺点：
内存占用量大：Transformer 的内存占用量随上下文长度而扩展。想要运行长上下文窗口，或大量并行批处理就需要大量硬件资源，这限制了大规模的实验和部署。 随着上下文的增长，推理速度会变慢：Transformer 的注意力机制导致推理时间相对于序列长度呈平方增长，吞吐会越来越慢。因为每个 token 都依赖于它之前的整个序列，所以要做到超长上下文就变得相当困难。
年前，来自卡内基梅隆和普林斯顿的两位大佬提出了 Mamba，一下子就点燃了人们的希望。
Mamba 以 SSM 为基础，增加了选择性提取信息的能力、以及硬件上高效的算法，一举解决了 Transformer 存在的问题。
这个新领域马上就吸引了大量的研究者，arXiv 上一时涌现了大量关于 Mamba 的应用和改进，比如将 Mamba 用于视觉的 Vision Mamba。
不得不说，现在的科研领域实在是太卷了，把 Transformer 引入视觉（ViT）用了三年，但 Mamba 到 Vision Mamba 只用了一个月。
不过原始 Mamba 的上下文长度较短，加上模型本身也没有做大，所以很难打过 SOTA 的 Transformer 模型，尤其是在与召回相关的任务上。
Jamba 于是更进一步，通过 Joint Attention and 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 将在不久之后加入 NVIDIA API 目录。
长上下文又出新选手
最近，各大公司都在卷长上下文。
具有较小上下文窗口的模型，往往会忘记最近对话的内容，而具有较大上下文的模型则避免了这种陷阱，可以更好地掌握所接收的数据流。
不过，具有长上下文窗口的模型，往往是计算密集的。
初创公司 AI21 Labs 的生成式模型就证明，事实并非如此。
Jamba 在具有至少 80GB 显存的单个 GPU（如 A100）上运行时，可以处理多达 140,000 个 token。
这相当于大约 105,000 字，或 210 页，是一本长度适中的长篇小说的篇幅。
相比之下，Meta Llama 2 的上下文窗口，只有 32,000 个 token，需要 12GB 的 GPU 显存。
按今天的标准来看，这种上下文窗口显然是偏小的。
对此，有网友也第一时间表示，性能什么的都不重要，关键的是 Jamba 有 256K 的上下文，除了 Gemini，其他人都没有这么长，——而 Jamba 可是开源的。
Jamba 真正的独特之处
从表面上看，Jamba 似乎并不起眼。
无论是昨天风头正盛的 DBRX，还是 Llama 2，现在都已经有大量免费提供、可下载的生成式 AI 模型。
而 Jamba 的独特之处，是藏在模型之下的： 它同时结合了两种模型架构—— Transformer 和状态空间模型 SSM 。
一方面，Transformer 是复杂推理任务的首选架构。 它最核心的定义特征，就是「注意力机制」。 对于每条输入数据，Transformer 会权衡所有其他输入的相关性，并从中提取以生成输出。
另一方面，SSM 结合了早前 AI 模型的多个优点，例如递归神经网络和卷积神经网络，因此能够实现长序列数据的处理，且计算效率更高。
虽然 SSM 有自己的局限性。但一些早期的代表 ，比如由普林斯顿和 CMU 提出的 Mamba，就可以处理比 Transformer 模型更大的输出，在语言生成任务上也更优。
对此，AI21 Labs 产品负责人 Dagan 表示——
虽然也有一些 SSM 模型的初步样例，但 Jamba 是第一个生产规模的商业级模型。
在他看来，Jamba 除了创新性和趣味性可供社区进一步研究，还提供了巨大的效率，和吞吐量的可能性。
目前，Jamba 是基于 Apache 2.0 许可发布的，使用限制较少但不能商用。后续的微调版本，预计会在几周内推出。
即便还处在研究的早期阶段，但 Dagan 断言，Jamba 无疑展示了 SSM 架构的巨大前景。
「这种模型的附加价值——无论是因为尺寸还是架构的创新——都可以很容易地安装到单个 GPU 上。」
他相信，随着 Mamba 的继续调整，性能将进一步提高。
参考资料：