Attention isn’t all you need!Mamba混合大模型开源:三倍Transformer吞吐量

Mamba 时代来了?自 2017 年开创性研究论文《Attention is All You Need》问世以来,transformer 架构就一直主导着生成式人工智能领域。然而,transformer 架构实际上有两个显著缺点:内存占用大:Transformer 的内存占用量随上下文长度而变化。这使得在没有大量硬件资源的情况下运行长上下文窗口或大量并行批处理变得具有挑战性,从而限制了广泛的实验和部署。 随着上下文长度的增加,推理速度会变慢:Transformer 的注意力机制随序列长度呈二次方扩展,并且会降低吞

Mamba 时代来了?

自 2017 年开创性研究论文《Attention is All You Need》问世以来,transformer 架构就一直主导着生成式人工智能领域。

然而,transformer 架构实际上有两个显著缺点:

内存占用大:Transformer 的内存占用量随上下文长度而变化。这使得在没有大量硬件资源的情况下运行长上下文窗口或大量并行批处理变得具有挑战性,从而限制了广泛的实验和部署。 

随着上下文长度的增加,推理速度会变慢:Transformer 的注意力机制随序列长度呈二次方扩展,并且会降低吞吐量,因为每个 token 都依赖于它之前的整个序列,从而将长上下文用例置于高效生产的范围之外。

但 transformer 并不是生成式人工智能唯一的前进方向。最近,AI21 Labs 推出并开源了一种名为「Jamba」的新方法,在多个基准上超越了 transformer。

图片

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

图片

Mamba 的 SSM 架构可以很好地解决 transformer 的内存资源和上下文问题。然而,Mamba 方法很难提供与 transformer 模型相同的输出水平。 

Jamba 将基于结构化状态空间模型 (SSM) 的 Mamba 模型与 transformer 架构相结合,旨在将 SSM 和 transformer 的最佳属性结合在一起。

图片

Jamba 还可以作为 NVIDIA NIM 推理微服务从 NVIDIA API 目录进行访问,企业应用程序开发人员可以使用 NVIDIA AI Enterprise 软件平台进行部署。

‍总的来说,Jamba 模型具有以下特点:

第一个基于 Mamba 的生产级模型,采用新颖的 SSM-Transformer 混合架构;

与 Mixtral 8x7B 相比,长上下文上的吞吐量提高了 3 倍;

提供对 256K 上下文窗口的访问;

公开了模型权重;

同等参数规模中唯一能够在单个 GPU 上容纳高达 140K 上下文的模型。

模型架构

如下图所示,Jamba 的架构采用块层(blocks-and-layers)方法,使 Jamba 能够集成两种架构。每个 Jamba 块包含一个注意力层或一个 Mamba 层,后跟一个多层感知器(MLP),从而形成 transformer 层。

图片

Jamba 利用 MoE 来增加模型参数的总数,同时简化推理中使用的活跃参数的数量,从而在计算需求没有相应增加的情况下获得更高的模型容量。为了在单个 80GB GPU 上最大限度地提高模型的质量和吞吐量,研究团队优化了所使用的 MoE 层和专家的数量,为常见推理工作负载留出了足够的内存。

Jamba 的 MoE 层允许它在推理时仅利用可用的 52B 参数中的 12B,并且其混合架构使这些 12B 活跃参数比同等大小的纯 transformer 模型更有效。 

此前,没有人将 Mamba 扩展到 3B 参数之外。Jamba 是同类模型中第一个达到生产级规模的混合架构。

吞吐量和效率

初步评估实验表明,Jamba 在吞吐量和效率等关键衡量指标上表现出色。

在效率方面,Jamba 在长上下文上的吞吐量达到了 Mixtral 8x7B 的 3 倍。Jamba 比 Mixtral 8x7B 等大小相当的基于 Transformer 的模型更高效。

图片

在成本方面,Jamba 可以在单个 GPU 上容纳 140K 上下文。与当前类似大小的其他开源模型相比,Jamba 能提供更多的部署和实验机会。

图片

需要注意的是,Jamba 目前不太可能取代当前基于 Transformer 的大型语言模型 (LLM),但它可能会成为某些领域的补充。

参考链接:

https://www.ai21.com/blog/announcing-jamba

https://venturebeat.com/ai/ai21-labs-juices-up-gen-ai-transformers-with-jamba/

相关资讯

直接扩展到无限长,谷歌Infini-Transformer终结上下文长度之争

不知 Gemini 1.5 Pro 是否用到了这项技术。谷歌又放大招了,发布下一代 Transformer 模型 Infini-Transformer。Infini-Transformer 引入了一种有效的方法,可以将基于 Transformer 的大型语言模型 (LLM) 扩展到无限长输入,而不增加内存和计算需求。使用该技术,研究者成功将一个 1B 的模型上下文长度提高到 100 万;应用到 8B 模型上,模型能处理 500K 的书籍摘要任务。自 2017 年开创性研究论文《Attention is All Yo

机器人策略学习的Game Changer?伯克利提出Body Transformer

过去几年间,Transformer 架构已经取得了巨大的成功,同时其也衍生出了大量变体,比如擅长处理视觉任务的 Vision Transformer(ViT)。本文要介绍的 Body Transformer(BoT) 则是非常适合机器人策略学习的 Transformer 变体。我们知道,物理智能体在执行动作的校正和稳定时,往往会根据其感受到的外部刺激的位置给出空间上的响应。比如人类对这些刺激的响应回路位于脊髓神经回路层面,它们专门负责单个执行器的响应。起校正作用的局部执行是高效运动的主要因素,这对机器人来说也尤为重

论文分享 | 丢弃卷积,纯Transformer构建GAN网络

最近,计算机视觉(CV)领域的研究者对 Transformer 产生了极大的兴趣并陆续取得了不少突破。比如,2020 年 5 月,Facebook AI 的研究者推出了 Transformer 的视觉版本——Detection Transformer,在性能上媲美当时的 SOTA 方法,但架构得到了极大简化;10 月,谷歌提出了 Vision Transformer (ViT),可以直接利用 transformer 对图像进行分类,而不需要卷积网络。