机器之心 | 五倍吞吐量,性能全面包围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

相关文章

Spring AI使用tool Calling和MCP的示例详解

《SpringAI使用toolCalling和MCP的示例详解》SpringAI1.0.0.M6引入ToolCalling与MCP协议,提升AI与工具交互的扩展性与标准化,支持信息检索、行动执行等... 目录深入探索 Spring AI聊天接口示例Function CallingMCPSTDIOSSE结束语

Spring Boot3.0新特性全面解析与应用实战

《SpringBoot3.0新特性全面解析与应用实战》SpringBoot3.0作为Spring生态系统的一个重要里程碑,带来了众多令人兴奋的新特性和改进,本文将深入解析SpringBoot3.0的... 目录核心变化概览Java版本要求提升迁移至Jakarta EE重要新特性详解1. Native Ima

Zabbix在MySQL性能监控方面的运用及最佳实践记录

《Zabbix在MySQL性能监控方面的运用及最佳实践记录》Zabbix通过自定义脚本和内置模板监控MySQL核心指标(连接、查询、资源、复制),支持自动发现多实例及告警通知,结合可视化仪表盘,可有效... 目录一、核心监控指标及配置1. 关键监控指标示例2. 配置方法二、自动发现与多实例管理1. 实践步骤

三频BE12000国补到手2549元! ROG 魔盒Pro WIFI7电竞AI路由器上架

《三频BE12000国补到手2549元!ROG魔盒ProWIFI7电竞AI路由器上架》近日,华硕带来了ROG魔盒ProWIFI7电竞AI路由器(ROGSTRIXGR7Pro),目前新... 华硕推出了ROG 魔盒Pro WIFI7电竞AI路由器(ROG STRIX GR7 Phttp://www.cppcn

MySQL深分页进行性能优化的常见方法

《MySQL深分页进行性能优化的常见方法》在Web应用中,分页查询是数据库操作中的常见需求,然而,在面对大型数据集时,深分页(deeppagination)却成为了性能优化的一个挑战,在本文中,我们将... 目录引言:深分页,真的只是“翻页慢”那么简单吗?一、背景介绍二、深分页的性能问题三、业务场景分析四、

全面掌握 SQL 中的 DATEDIFF函数及用法最佳实践

《全面掌握SQL中的DATEDIFF函数及用法最佳实践》本文解析DATEDIFF在不同数据库中的差异,强调其边界计算原理,探讨应用场景及陷阱,推荐根据需求选择TIMESTAMPDIFF或inte... 目录1. 核心概念:DATEDIFF 究竟在计算什么?2. 主流数据库中的 DATEDIFF 实现2.1

MySQL 多列 IN 查询之语法、性能与实战技巧(最新整理)

《MySQL多列IN查询之语法、性能与实战技巧(最新整理)》本文详解MySQL多列IN查询,对比传统OR写法,强调其简洁高效,适合批量匹配复合键,通过联合索引、分批次优化提升性能,兼容多种数据库... 目录一、基础语法:多列 IN 的两种写法1. 直接值列表2. 子查询二、对比传统 OR 的写法三、性能分析

Linux系统性能检测命令详解

《Linux系统性能检测命令详解》本文介绍了Linux系统常用的监控命令(如top、vmstat、iostat、htop等)及其参数功能,涵盖进程状态、内存使用、磁盘I/O、系统负载等多维度资源监控,... 目录toppsuptimevmstatIOStatiotopslabtophtopdstatnmon

Java操作Word文档的全面指南

《Java操作Word文档的全面指南》在Java开发中,操作Word文档是常见的业务需求,广泛应用于合同生成、报表输出、通知发布、法律文书生成、病历模板填写等场景,本文将全面介绍Java操作Word文... 目录简介段落页头与页脚页码表格图片批注文本框目录图表简介Word编程最重要的类是org.apach

Knife4j+Axios+Redis前后端分离架构下的 API 管理与会话方案(最新推荐)

《Knife4j+Axios+Redis前后端分离架构下的API管理与会话方案(最新推荐)》本文主要介绍了Swagger与Knife4j的配置要点、前后端对接方法以及分布式Session实现原理,... 目录一、Swagger 与 Knife4j 的深度理解及配置要点Knife4j 配置关键要点1.Spri