Mamba与MoE架构强强联合,Mamba-MoE高效提升LLM计算效率和可扩展性

本文主要是介绍Mamba与MoE架构强强联合,Mamba-MoE高效提升LLM计算效率和可扩展性,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!

43008843de874a75be6866fb07305205.png

论文题目: MoE-Mamba: Efficient Selective State Space Models with Mixture of Experts
论文链接: https://arxiv.org/abs/2401.04081
代码仓库: GitHub - llm-random/llm-random

作为大型语言模型(LLM)基础架构的后起之秀,状态空间模型(State Space Models,SSMs)在序列数据建模领域中已取得了惊人的发展。其中Mamba模型改进了传统的SSM,其通过输入依赖的方式来调整SSM中的参数,允许模型自适应的根据输入数据选择性的传输或遗忘信息,来提高模型在密集型数据上的计算效率。与此同时,Mixture of Experts(MoE)框架也显着改进了基于Transformer的LLM,如何设计更高效的MoE混合策略也成为了LLM研究领域中的潮流方向。

本文介绍一篇来自IDEAS NCBR和华沙大学合作完成的文章,本文作者探索了如何将SSM与MoE结合起来,提出了一种名为MoE-Mamba的框架,MoE-Mamba继承了SSM序列模型的递归计算特性,在推理速度方面相比传统Transformer具有天然优势,同时预测精度优于 Mamba 和 Transformer-MoE。特别的是,MoE-Mamba可以在减少训练steps的情况下达到与普通 Mamba 相同的性能

01. 引言

SSMs模型相比Transformer具有很多优势,例如可并行训练、推理时的线性时间复杂度以及在长上下文任务上的依赖捕获能力。特别是近期提出的Mamba模型[1],基于选择性的SSM和硬件感知设计在多种语言理解任务上取得了优异的结果,被学术界广泛认为是下一个替代注意力Transformer结构的有力竞争者。Mamba对GPU显存的占用不依赖于上下文长度,因为其反向传播所需的中间状态不会被保存,而是在反向传播期间重新计算,强调了状态压缩的重要性。此外,专家混合模型MoE[2]被证明可以有效扩展和提升Transformer模型的性能,其可以大幅增加模型的参数数量,而不会对模型推理和训练所需的 FLOP 产生太大影响,例如近期开源的Mixtral8×7B [3],性能与LLaMa-2 70B相当,但是推理计算量仅为后者的1/6。

4d3db7d34a9f434ebf3abfe5e52080bb.png

本文引入了一种将Mamba与MoE层相结合的模型MoE-Mamba,MoE-Mamba实现了a+b>c的效果,上图展示了MoE-Mamba与其他baseline方法的性能对比效果,可以看到,MoE-Mamba可以在普通Mamba的基础上实现对模型训练步骤缩减的效果。如上图黑色虚线所示,MoE-Mamba在减少2.2倍训练steps的情况下达到了与普通Mamba相同的性能。后续的其他实验也表明,MoE-Mamba拥有大规模扩展模型参数规模的潜力

02. 本文方法

尽管 Mamba 的主要底层机制与 Transformer 中使用的注意力机制有很大不同,但 Mamba 仍然保留了 Transformer 模型的层次结构(即块叠加)。例如包含一层或多层的相同块依次堆叠,每一层的输出会被汇聚到残差信息流中再送入到下一个块中,残差流的最终状态随后被用于预测语言建模任务中的下一个token。下图展示了这些架构的细节对比,从左到右分别是vanilla Transformer、MoE-Transformer、Mamba、MoE-Mamba。

39316c309a36437cb2d00ccfbb5e5475.png

本文提出的MoE-Mamba充分利用了前两种架构的兼容性,例如,在原有Mamba结构的基础上仿照MoE-Transformer将两个mamba块中的其中一个替换成一个可选择的MoE块。这种将mamba层与MoE交错设置的模式可以有效地将序列的整个上下文集成到mamba块的内部表示中,从而将其与MoE层的条件处理分开。

b19196ff06eb4a16afb52ec48d98a517.png

此外,本文作者认为如果将Mamba块和MoE块在局部层的范围内进行并行执行也是一个非常有前景的改进方向,如上图左侧展示了一种并行的Mamba+MoE 架构,右侧展示了Mamba Block的构成。如果将Mamba Block中的输出投影也替换为MoE,模型可以选择更少的模块来匹配当前输入计算的需要,也能实现与原始Mamba架构相当的效果。当然,也可以进一步将MoE替换Conv Projection层来进一步减少计算量。

03. 实验效果

在本文的实验部分,作者比较了5种不同的设置:vanilla Transformer、Mamba、Mamba-MLP、MoE 和 MoE-Mamba。为了保证 Mamba 和本文MoE-Mamba模型中每个token的活动参数数量大致相同,作者将每个MoE前馈层的参数量进行了缩减。不包括嵌入层和非嵌入层,所有模型的每个 token 大约需要访问 26M 个参数,训练数据集使用C4,整体的训练量为大约6.5B个token和10万个训练steps。

4abaf58c60e6438a8b3019c786c358e1.png

上表展示了上述几种对比模型的训练效果,可以看到,MoE-Mamba 比普通 Mamba 模型有了显着的改进。值得注意的是,MoE-Mamba 只需 46% 的训练steps即可达到与 vanilla Mamba 相同的性能

38c3505059b24e9aa01b89eacc894728.png

MoE中的专家数量是MoE-Mamba中的一个重要超参数,为了评估Mamba能否随着专家数量的增加而扩展,作者在下图绘制了模型在不同专家数量情况下的运行情况,并且加入了Mamba 和 Mamba-MLP(后者相当于具有单个专家的 MoE-Mamba)作为参考。上表展示了模型在10万训练steps之后的结果,结果表明,MoE-Mamba可以很好的适应专家数量的变化,如果专家数量设置为8或超过8,MoE-Mamba可以获得比普通Mamba模型更好的最终性能

5dcdcb5239e04e78aac141a8506854b1.png

04. 总结

本文提出了一种将专家混合MoE技术与Mamba架构进行集成的模块,即MoE-Mamba。基于并行计算的天然优势,Mamba减轻了大模型中复杂循环顺序性质的影响,并且对硬件进行感知来实现参数扩展。Mamba相比普通注意力机制Transformer解决了序列模型中效率和有效性之间的基本权衡,强调了状态压缩的重要性。将Mamba与高度稀疏的MoE前馈层交错设置可以实现更高推理效率的LLM,但目前的组合方式仍然非常简单,作者也探索了一种局部并行的Mamba+MoE架构以实现更高的预测准确率和更稀疏的推理效果。期望后续有更多基于条件计算与状态空间模型技术相结合的技术出现,作者认为这条道路将能够更有效地扩展到更大的语言模型中。

参考

[1] Gu A, Dao T. Mamba: Linear-time sequence modeling with selective state spaces[J]. arXiv preprint arXiv:2312.00752, 2023.
[2] Sanseviero, O., Tunstall, L., Schmid, P., Mangrulkar, S., Belkada, Y., and Cuenca, P. Mixture of experts explained, 2023. URL https://huggingface.co/blog/moe.

[3] Mistral. Mixtral of experts, Dec 2023. URL https://mistral.ai/news/mixtral-of-experts/.


  关于TechBeat人工智能社区

TechBeat(www.techbeat.net)隶属于将门创投,是一个荟聚全球华人AI精英的成长社区。

我们希望为AI人才打造更专业的服务和体验,加速并陪伴其学习成长。

期待这里可以成为你学习AI前沿知识的高地,分享自己最新工作的沃土,在AI进阶之路上的升级打怪的根据地!

更多详细介绍>>TechBeat,一个荟聚全球华人AI精英的学习成长社区 

这篇关于Mamba与MoE架构强强联合,Mamba-MoE高效提升LLM计算效率和可扩展性的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!



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

相关文章

Java异常架构Exception(异常)详解

《Java异常架构Exception(异常)详解》:本文主要介绍Java异常架构Exception(异常),具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,望不吝赐教... 目录1. Exception 类的概述Exception的分类2. 受检异常(Checked Exception)

SpringBoot使用OkHttp完成高效网络请求详解

《SpringBoot使用OkHttp完成高效网络请求详解》OkHttp是一个高效的HTTP客户端,支持同步和异步请求,且具备自动处理cookie、缓存和连接池等高级功能,下面我们来看看SpringB... 目录一、OkHttp 简介二、在 Spring Boot 中集成 OkHttp三、封装 OkHttp

使用Python高效获取网络数据的操作指南

《使用Python高效获取网络数据的操作指南》网络爬虫是一种自动化程序,用于访问和提取网站上的数据,Python是进行网络爬虫开发的理想语言,拥有丰富的库和工具,使得编写和维护爬虫变得简单高效,本文将... 目录网络爬虫的基本概念常用库介绍安装库Requests和BeautifulSoup爬虫开发发送请求解

C++实现回文串判断的两种高效方法

《C++实现回文串判断的两种高效方法》文章介绍了两种判断回文串的方法:解法一通过创建新字符串来处理,解法二在原字符串上直接筛选判断,两种方法都使用了双指针法,文中通过代码示例讲解的非常详细,需要的朋友... 目录一、问题描述示例二、解法一:将字母数字连接到新的 string思路代码实现代码解释复杂度分析三、

Python如何计算两个不同类型列表的相似度

《Python如何计算两个不同类型列表的相似度》在编程中,经常需要比较两个列表的相似度,尤其是当这两个列表包含不同类型的元素时,下面小编就来讲讲如何使用Python计算两个不同类型列表的相似度吧... 目录摘要引言数字类型相似度欧几里得距离曼哈顿距离字符串类型相似度Levenshtein距离Jaccard相

轻松上手MYSQL之JSON函数实现高效数据查询与操作

《轻松上手MYSQL之JSON函数实现高效数据查询与操作》:本文主要介绍轻松上手MYSQL之JSON函数实现高效数据查询与操作的相关资料,MySQL提供了多个JSON函数,用于处理和查询JSON数... 目录一、jsON_EXTRACT 提取指定数据二、JSON_UNQUOTE 取消双引号三、JSON_KE

Tomcat高效部署与性能优化方式

《Tomcat高效部署与性能优化方式》本文介绍了如何高效部署Tomcat并进行性能优化,以确保Web应用的稳定运行和高效响应,高效部署包括环境准备、安装Tomcat、配置Tomcat、部署应用和启动T... 目录Tomcat高效部署与性能优化一、引言二、Tomcat高效部署三、Tomcat性能优化总结Tom

Python利用自带模块实现屏幕像素高效操作

《Python利用自带模块实现屏幕像素高效操作》这篇文章主要为大家详细介绍了Python如何利用自带模块实现屏幕像素高效操作,文中的示例代码讲解详,感兴趣的小伙伴可以跟随小编一起学习一下... 目录1、获取屏幕放缩比例2、获取屏幕指定坐标处像素颜色3、一个简单的使用案例4、总结1、获取屏幕放缩比例from

MySQL 缓存机制与架构解析(最新推荐)

《MySQL缓存机制与架构解析(最新推荐)》本文详细介绍了MySQL的缓存机制和整体架构,包括一级缓存(InnoDBBufferPool)和二级缓存(QueryCache),文章还探讨了SQL... 目录一、mysql缓存机制概述二、MySQL整体架构三、SQL查询执行全流程四、MySQL 8.0为何移除查

微服务架构之使用RabbitMQ进行异步处理方式

《微服务架构之使用RabbitMQ进行异步处理方式》本文介绍了RabbitMQ的基本概念、异步调用处理逻辑、RabbitMQ的基本使用方法以及在SpringBoot项目中使用RabbitMQ解决高并发... 目录一.什么是RabbitMQ?二.异步调用处理逻辑:三.RabbitMQ的基本使用1.安装2.架构