机器之心 | 五倍吞吐量,性能全面包围Transformer:新架构Mamba引爆AI圈

本文主要是介绍机器之心 | 五倍吞吐量,性能全面包围Transformer:新架构Mamba引爆AI圈,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!

本文来源公众号“机器之心”,仅用于学术分享,侵权删,干货满满。

原文链接:五倍吞吐量,性能全面包围Transformer:新架构Mamba引爆AI圈

屹立不倒的 Transformer 迎来了一个强劲竞争者。

在别的领域,如果你想形容一个东西非常重要,你可能将其形容为「撑起了某领域的半壁江山」。但在 AI 大模型领域,Transformer 架构不能这么形容,因为它几乎撑起了「整个江山」。

自 2017 年被提出以来,Transformer 已经成为 AI 大模型的主流架构,但随着模型规模的扩展和需要处理的序列不断变长,Transformer 的局限性也逐渐凸显。一个很明显的缺陷是:Transformer 模型中自注意力机制的计算量会随着上下文长度的增加呈平方级增长,比如上下文增加 32 倍时,计算量可能会增长 1000 倍,计算效率非常低。

为了克服这些缺陷,研究者们开发出了很多注意力机制的高效变体,但这往往以牺牲其有效性特为代价。到目前为止,这些变体都还没有被证明能在不同领域发挥有效作用。

最近,一项名为「Mamba」的研究似乎打破了这一局面。

在这篇论文中,研究者提出了一种新的架构 ——「选择性状态空间模型( selective state space model)」。它在多个方面改进了先前的工作。

作者表示,「Mamba」在语言建模方面可以媲美甚至击败 Transformer。而且,它可以随上下文长度的增加实现线性扩展,其性能在实际数据中可提高到百万 token 长度序列,并实现 5 倍的推理吞吐量提升。

消息一出,人们纷纷点赞,有人表示已经迫不及待想要把它用在大模型上了。

作为通用序列模型的骨干,Mamba 在语言、音频和基因组学等多种模态中都达到了 SOTA 性能。在语言建模方面,无论是预训练还是下游评估,他们的 Mamba-3B 模型都优于同等规模的 Transformer 模型,并能与两倍于其规模的 Transformer 模型相媲美。

这篇论文的作者只有两位,一位是卡内基梅隆大学机器学习系助理教授 Albert Gu,另一位是 Together.AI 首席科学家、普林斯顿大学计算机科学助理教授(即将上任)Tri Dao。

左:Albert Gu;右:Tri Dao。

Albert Gu 表示,这项研究的一个重要创新是引入了一个名为「选择性 SSM」的架构,该架构是 Albert Gu 此前主导研发的 S4 架构(Structured State Spaces for Sequence Modeling ,用于序列建模的结构化状态空间)的一个简单泛化,可以有选择地决定关注还是忽略传入的输入。一个「小小的改变」—— 让某些参数成为输入的函数,结果却非常有效。

值得一提的是,S4 是一个非常成功的架构。此前,它成功地对  Long Range Arena (LRA) 中的长程依赖进行了建模,并成为首个在 Path-X 上获得高于平均性能的模型。更具体地说,S4 是一类用于深度学习的序列模型,与 RNN、CNN 和经典的状态空间模型(State Space Model,SSM)广泛相关。SSM 是独立的序列转换,可被整合到端到端神经网络架构中( SSM 架构有时也称 SSNN,它与 SSM 层的关系就像 CNN 与线性卷积层的关系一样)。Mamba 论文也讨论了一些著名的 SSM 架构,比如 Linear attention、H3、Hyena、RetNet、RWKV,其中许多也将作为论文研究的基线。Mamba 的成功让 Albert Gu 对 SSM 的未来充满了信心。

Tri Dao 则是 FlashAttention、Flash Attention v2、Flash-Decoding的作者。FlashAttention 是一种对注意力计算进行重新排序并利用经典技术(平铺、重新计算)加快速度并将内存使用从序列长度的二次减少到线性的算法。Flash Attention v2、Flash-Decoding 都是建立在 Flash Attention 基础上的后续工作,把大模型的长文本推理效率不断推向极限。在 Mamba 之前,Tri Dao 和 Albert Gu 也有过合作。

另外,这项研究的模型代码和预训练的检查点是开源的,参见以下链接:https://github.com/state-spaces/mamba.

论文链接:https://arxiv.org/ftp/arxiv/papers/2312/2312.00752.pdf

方法创新

论文第 3.1 节介绍了如何利用合成任务的直觉来启发选择机制,第 3.2 节解释了如何将这一机制纳入状态空间模型。由此产生的时变 SSM 不能使用卷积,导致了高效计算的技术难题。研究者采用了一种硬件感知算法,利用当前硬件的内存层次结构来克服这一难题(第 3.3 节)。第 3.4 节描述了一个简单的 SSM 架构,不需要注意力,甚至不需要 MLP 块。第 3.5 节讨论了选择机制的一些其他特性。

选择机制

研究者发现了此前模型的一个关键局限:以依赖输入的方式高效选择数据的能力(即关注或忽略特定输入)。

序列建模的一个基本方法是将上下文压缩到更小的状态,我们可以从这个角度来看待当下流行的序列模型。例如,注意力既高效又低效,因为它根本没有明确压缩上下文。这一点可以从自回归推理需要明确存储整个上下文(即 KV 缓存)这一事实中看出,这直接导致了 Transformer 缓慢的线性时间推理和二次时间训练。

递归模型的效率很高,因为它们的状态是有限的,这意味着恒定时间推理和线性时间训练。然而,它们的高效性受限于这种状态对上下文的压缩程度。

为了理解这一原理,下图展示了两个合成任务的运行示例:

研究者设计了一种简单的选择机制,根据输入对 SSM 参数进行参数化。这样,模型就能过滤掉无关信息,并无限期地记住相关信息。

将选择机制纳入模型的一种方法是让影响序列交互的参数(如 RNN 的递归动力学或 CNN 的卷积核)与输入相关。算法 1 和 2 展示了本文使用的主要选择机制。其主要区别在于,该方法只需将几个参数 ∆,B,C 设置为输入函数,并在整个过程中改变张量形状。这些参数现在都有一个长度维度 L ,意味着模型已经从时间不变变为时间可变。

硬件感知算法

上述变化对模型的计算提出了技术挑战。所有先前的 SSM 模型都必须是时间和输入不变的,这样才能提高计算效率。为此,研究者采用了一种硬件感知算法,通过扫描而不是卷积来计算模型,但不会将扩展状态具体化,以避免在 GPU 存储器层次结构的不同级别之间进行 IO 访问。由此产生的实现方法在理论上(与所有基于卷积的 SSM 的伪线性相比,在序列长度上呈线性缩放)和现有硬件上都比以前的方法更快(在 A100 GPU 上可快达 3 倍)。

架构

研究者将先前的 SSM 架构设计与 Transformer 的 MLP 块合并为一个块,从而简化了深度序列模型架构,形成了一种包含选择性状态空间的简单、同质的架构设计(Mamba)。

与结构化 SSM 一样,选择性 SSM 也是一种独立的序列变换,可以灵活地融入神经网络。H3 架构是著名的同质化架构设计的基础,通常由线性注意力启发的块和 MLP(多层感知器)块交错组成。

研究者简化了这一架构,将这两个部分合二为一,均匀堆叠,如图 3。他们受到门控注意力单元(GAU)的启发,该单元也对注意力做了类似的处理。

选择性 SSM 以及 Mamba 架构的扩展是完全递归模型,几个关键特性使其适合作为在序列上运行的通用基础模型的骨干:

  1. 高质量:选择性为语言和基因组学等密集模型带来了强大的性能。

  2. 快速训练和推理:在训练过程中,计算量和内存与序列长度成线性关系,而在推理过程中,由于不需要缓存以前的元素,自回归展开模型每一步只需要恒定的时间。

  3. 长上下文:质量和效率共同提高了实际数据的性能,序列长度可达 100 万。

实验评估

实证验证了 Mamba 作为通用序列基础模型骨干的潜力,无论是在预训练质量还是特定领域的任务性能方面,Mamba 都能在多种类型的模态和环境中发挥作用:

合成任务。在复制和感应头等重要的语言模型合成任务上,Mamba 不仅能轻松解决,而且能推断出无限长的解决方案(>100 万 token)。

音频和基因组学。在音频波形和 DNA 序列建模方面,Mamba 在预训练质量和下游指标方面都优于 SaShiMi、Hyena、Transformer 等先前的 SOTA 模型(例如,在具有挑战性的语音生成数据集上将 FID 降低了一半以上)。在这两种情况下,它的性能随着上下文长度的增加而提高,最高可达百万长度的序列。

语言建模。Mamba 是首个线性时间序列模型,在预训练复杂度和下游评估方面都真正达到了 Transformer 质量的性能。通过多达 1B 参数的缩放规律,研究者发现 Mamba 的性能超过了大量基线模型,包括 LLaMa 这种非常强大的现代 Transformer 训练配方。

与类似规模的 Transformer 相比,Mamba 具有 5 倍的生成吞吐量,而且 Mamba-3B 的质量与两倍于其规模的 Transformer 相当(例如,与 Pythia-3B 相比,常识推理的平均值高出 4 分,甚至超过 Pythia-7B)。

THE END !

文章结束,感谢阅读。您的点赞,收藏,评论是我继续更新的动力。大家有推荐的公众号可以评论区留言,共同学习,一起进步。

这篇关于机器之心 | 五倍吞吐量,性能全面包围Transformer:新架构Mamba引爆AI圈的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!



http://www.chinasem.cn/article/1122621

相关文章

Vue3 的 shallowRef 和 shallowReactive:优化性能

大家对 Vue3 的 ref 和 reactive 都很熟悉,那么对 shallowRef 和 shallowReactive 是否了解呢? 在编程和数据结构中,“shallow”(浅层)通常指对数据结构的最外层进行操作,而不递归地处理其内部或嵌套的数据。这种处理方式关注的是数据结构的第一层属性或元素,而忽略更深层次的嵌套内容。 1. 浅层与深层的对比 1.1 浅层(Shallow) 定义

Ilya-AI分享的他在OpenAI学习到的15个提示工程技巧

Ilya(不是本人,claude AI)在社交媒体上分享了他在OpenAI学习到的15个Prompt撰写技巧。 以下是详细的内容: 提示精确化:在编写提示时,力求表达清晰准确。清楚地阐述任务需求和概念定义至关重要。例:不用"分析文本",而用"判断这段话的情感倾向:积极、消极还是中性"。 快速迭代:善于快速连续调整提示。熟练的提示工程师能够灵活地进行多轮优化。例:从"总结文章"到"用

mybatis的整体架构

mybatis的整体架构分为三层: 1.基础支持层 该层包括:数据源模块、事务管理模块、缓存模块、Binding模块、反射模块、类型转换模块、日志模块、资源加载模块、解析器模块 2.核心处理层 该层包括:配置解析、参数映射、SQL解析、SQL执行、结果集映射、插件 3.接口层 该层包括:SqlSession 基础支持层 该层保护mybatis的基础模块,它们为核心处理层提供了良好的支撑。

AI绘图怎么变现?想做点副业的小白必看!

在科技飞速发展的今天,AI绘图作为一种新兴技术,不仅改变了艺术创作的方式,也为创作者提供了多种变现途径。本文将详细探讨几种常见的AI绘图变现方式,帮助创作者更好地利用这一技术实现经济收益。 更多实操教程和AI绘画工具,可以扫描下方,免费获取 定制服务:个性化的创意商机 个性化定制 AI绘图技术能够根据用户需求生成个性化的头像、壁纸、插画等作品。例如,姓氏头像在电商平台上非常受欢迎,

百度/小米/滴滴/京东,中台架构比较

小米中台建设实践 01 小米的三大中台建设:业务+数据+技术 业务中台--从业务说起 在中台建设中,需要规范化的服务接口、一致整合化的数据、容器化的技术组件以及弹性的基础设施。并结合业务情况,判定是否真的需要中台。 小米参考了业界优秀的案例包括移动中台、数据中台、业务中台、技术中台等,再结合其业务发展历程及业务现状,整理了中台架构的核心方法论,一是企业如何共享服务,二是如何为业务提供便利。

性能测试介绍

性能测试是一种测试方法,旨在评估系统、应用程序或组件在现实场景中的性能表现和可靠性。它通常用于衡量系统在不同负载条件下的响应时间、吞吐量、资源利用率、稳定性和可扩展性等关键指标。 为什么要进行性能测试 通过性能测试,可以确定系统是否能够满足预期的性能要求,找出性能瓶颈和潜在的问题,并进行优化和调整。 发现性能瓶颈:性能测试可以帮助发现系统的性能瓶颈,即系统在高负载或高并发情况下可能出现的问题

从去中心化到智能化:Web3如何与AI共同塑造数字生态

在数字时代的演进中,Web3和人工智能(AI)正成为塑造未来互联网的两大核心力量。Web3的去中心化理念与AI的智能化技术,正相互交织,共同推动数字生态的变革。本文将探讨Web3与AI的融合如何改变数字世界,并展望这一新兴组合如何重塑我们的在线体验。 Web3的去中心化愿景 Web3代表了互联网的第三代发展,它基于去中心化的区块链技术,旨在创建一个开放、透明且用户主导的数字生态。不同于传统

AI一键生成 PPT

AI一键生成 PPT 操作步骤 作为一名打工人,是不是经常需要制作各种PPT来分享我的生活和想法。但是,你们知道,有时候灵感来了,时间却不够用了!😩直到我发现了Kimi AI——一个能够自动生成PPT的神奇助手!🌟 什么是Kimi? 一款月之暗面科技有限公司开发的AI办公工具,帮助用户快速生成高质量的演示文稿。 无论你是职场人士、学生还是教师,Kimi都能够为你的办公文

性能分析之MySQL索引实战案例

文章目录 一、前言二、准备三、MySQL索引优化四、MySQL 索引知识回顾五、总结 一、前言 在上一讲性能工具之 JProfiler 简单登录案例分析实战中已经发现SQL没有建立索引问题,本文将一起从代码层去分析为什么没有建立索引? 开源ERP项目地址:https://gitee.com/jishenghua/JSH_ERP 二、准备 打开IDEA找到登录请求资源路径位置

Andrej Karpathy最新采访:认知核心模型10亿参数就够了,AI会打破教育不公的僵局

夕小瑶科技说 原创  作者 | 海野 AI圈子的红人,AI大神Andrej Karpathy,曾是OpenAI联合创始人之一,特斯拉AI总监。上一次的动态是官宣创办一家名为 Eureka Labs 的人工智能+教育公司 ,宣布将长期致力于AI原生教育。 近日,Andrej Karpathy接受了No Priors(投资博客)的采访,与硅谷知名投资人 Sara Guo 和 Elad G