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

相关文章

MySQL重复数据处理的七种高效方法

《MySQL重复数据处理的七种高效方法》你是不是也曾遇到过这样的烦恼:明明系统测试时一切正常,上线后却频频出现重复数据,大批量导数据时,总有那么几条不听话的记录导致整个事务莫名回滚,今天,我就跟大家分... 目录1. 重复数据插入问题分析1.1 问题本质1.2 常见场景图2. 基础解决方案:使用异常捕获3.

如何高效移除C++关联容器中的元素

《如何高效移除C++关联容器中的元素》关联容器和顺序容器有着很大不同,关联容器中的元素是按照关键字来保存和访问的,而顺序容器中的元素是按它们在容器中的位置来顺序保存和访问的,本文介绍了如何高效移除C+... 目录一、简介二、移除给定位置的元素三、移除与特定键值等价的元素四、移除满足特android定条件的元

基于Python实现高效PPT转图片工具

《基于Python实现高效PPT转图片工具》在日常工作中,PPT是我们常用的演示工具,但有时候我们需要将PPT的内容提取为图片格式以便于展示或保存,所以本文将用Python实现PPT转PNG工具,希望... 目录1. 概述2. 功能使用2.1 安装依赖2.2 使用步骤2.3 代码实现2.4 GUI界面3.效

Java图片压缩三种高效压缩方案详细解析

《Java图片压缩三种高效压缩方案详细解析》图片压缩通常涉及减少图片的尺寸缩放、调整图片的质量(针对JPEG、PNG等)、使用特定的算法来减少图片的数据量等,:本文主要介绍Java图片压缩三种高效... 目录一、基于OpenCV的智能尺寸压缩技术亮点:适用场景:二、JPEG质量参数压缩关键技术:压缩效果对比

C#使用SQLite进行大数据量高效处理的代码示例

《C#使用SQLite进行大数据量高效处理的代码示例》在软件开发中,高效处理大数据量是一个常见且具有挑战性的任务,SQLite因其零配置、嵌入式、跨平台的特性,成为许多开发者的首选数据库,本文将深入探... 目录前言准备工作数据实体核心技术批量插入:从乌龟到猎豹的蜕变分页查询:加载百万数据异步处理:拒绝界面

Spring Boot + MyBatis Plus 高效开发实战从入门到进阶优化(推荐)

《SpringBoot+MyBatisPlus高效开发实战从入门到进阶优化(推荐)》本文将详细介绍SpringBoot+MyBatisPlus的完整开发流程,并深入剖析分页查询、批量操作、动... 目录Spring Boot + MyBATis Plus 高效开发实战:从入门到进阶优化1. MyBatis

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思路代码实现代码解释复杂度分析三、