Redian新闻
>
MoE与Mamba强强联合,将状态空间模型扩展到数百亿参数

MoE与Mamba强强联合,将状态空间模型扩展到数百亿参数

公众号新闻

机器之心报道

编辑:Panda

性能与 Mamba 一样,但所需训练步骤数却少 2.2 倍。


状态空间模型(SSM)是近来一种备受关注的 Transformer 替代技术,其优势是能在长上下文任务上实现线性时间的推理、并行化训练和强大的性能。而基于选择性 SSM 和硬件感知型设计的 Mamba 更是表现出色,成为了基于注意力的 Transformer 架构的一大有力替代架构。

近期也有一些研究者在探索将 SSM 和 Mamba 与其它方法组合起来创造更强大的架构,比如机器之心曾报告过《Mamba 可以替代 Transformer,但它们也能组合起来使用》。

近日,波兰一个研究团队发现,如果将 SSM 与混合专家系统(MoE/Mixture of Experts)组合起来,可望让 SSM 实现大规模扩展。MoE 是目前常用于扩展 Transformer 的技术,比如近期的 Mixtral 模型就使用了这一技术,参阅机器之心文章

这个波兰研究团队给出的研究成果是 MoE-Mamba,即将 Mamba 和混合专家层组合起来的模型。

论文地址:https://arxiv.org/pdf/2401.04081.pdf

MoE-Mamba 能同时提升 SSM 和 MoE 的效率。而且该团队还发现,当专家的数量发生变化时,MoE-Mamba 的行为是可预测的。

该团队也进行了实验论证,如图 1 所示,结果表明:相比于 Mamba,MoE-Mamba 达到同等性能时所需的训练步骤数少 2.2 倍,这彰显了新方法相较于 Transformer 和 Transformer-MoE 的潜在优势。这些初步结果也指出了一个颇具潜力的研究方向:SSM 也许可以扩展到数百亿参数!

相关研究

状态空间模型

状态空间模型(SSM)是一类用于序列建模的架构。这些模型的思想源自控制论领域,可被看作是 RNN 和 CNN 的组合。尽管它们具有相当大的优势,但也有一些问题,因此难以成为语言建模任务的主导架构。但是,近期的一些研究突破却让深度 SSM 可以扩展到数十亿参数,同时还能维持计算效率和强大的性能表现。

Mamba

Mamba 是基于 SSM 构建的模型,能实现线性时间的推理速度(对上下文长度而言),并且其还通过硬件感知型设计实现了高效的训练流程。Mamba 采用了一种工作高效型的并行扫描方法,可以减轻循环的序列性的影响,而融合 GPU 操作则可无需实现扩展状态。反向传播所必需的中间状态不会被保存下来,而是会在反向通过过程中被重新计算,由此可以降低内存需求。Mamba 优于注意力机制的优势在推理阶段尤其显著,因为其不仅能降低计算复杂度,而且内存使用量还不会取决于上下文长度。

Mamba 能解决序列模型的效率和效果之间的根本性权衡,这就凸显了状态压缩的重要性。高效的模型必需要小状态,而有效的模型所需的状态应当包含上下文的所有关键信息。不同于其它 SSM 对时间和输入不变性的需求,Mamba 引入了一种选择机制,可以控制信息沿序列维度传播的方式。这一设计选择的灵感来自对选择性复制和归纳头等合成任务的直观理解,让模型可以分辨和保留关键信息,同时滤除无关信息。

研究发现,Mamba 有能力高效地利用更长的上下文(长达 1M token),并且随着上下文长度增长,预训练困惑度也会得到改善。Mamba 模型是由堆叠的 Mamba 块构成的,在 NLP、基因组学、音频等多个不同领域都取得了非常好的结果,其性能可以媲美和超越已有的 Transformer 模型。因此,Mamba 成为了通用序列建模骨干模型的一个有力候选模型,参阅《五倍吞吐量,性能全面包围 Transformer:新架构 Mamba 引爆 AI 圈》。

混合专家

混合专家(MoE)这类技术能极大提升模型的参数数量,同时不会影响模型推理和训练所需的 FLOPs。MoE 最早由 Jacobs et al. 于 1991 年提出,并在 2017 年由 Shazeer et al. 开始用于 NLP 任务。

MoE 有一个优势:激活很稀疏 —— 对于处理的每个 token,只会用到模型的一小部分参数。由于其计算需求,Transformer 中的前向层已经变成了多种 MoE 技术的标准目标。

研究社区已经提出了多种方法用于解决 MoE 的核心问题,即将 token 分配给专家的过程,也称路由(routing)过程。目前有两种基本的路由算法:Token Choice 和 Expert Choice。其中前者是将每个 token 路由到一定数量(K)的专家,至于后者则是路由到每个专家的 token 数量是固定的。

Fedus et al. 在 2022 年的论文《Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity》中提出的 Switch 是一种 Token Choice 架构,其是将每个 token 路由到单个专家(K=1),而他们使用该方法将 Transformer 的参数规模成功扩增至了 1.6 万亿。波兰的这个团队在实验中也采用了这种 MoE 设计。

最近,MoE 也开始进入开源社区,比如 OpenMoE。

项目地址:https://github.com/XueFuzhao/OpenMoE

尤其值得一提的是 Mistral 开源的 Mixtral 8×7B,其性能可比肩 LLaMa 2 70B,同时所需的推理计算预算只有后者的约六分之一。

模型架构

尽管 Mamba 的主要底层机制与 Transformer 中使用的注意力机制大不相同,但 Mamba 保留了 Transformer 模型的高层级、基于模块的结构。使用这一范式,由相同模块构成的一层或多层会彼此堆叠在一起,而每一层的输出都会被添加到残差流(residual stream)中,见图 2。之后,这个残差流的最终值会被用于预测语言建模任务的下一个 token。

MoE-Mamba 利用了这两种架构的兼容能力。如图 2 所示,在 MoE-Mamba 中,每间隔一个 Mamba 层就会被替换成一个基于 Switch 的 MoE 前馈层。

不过该团队也注意到这一设计和《Mamba: Linear-time sequence modeling with selective state spaces》的设计有些相似;后者交替堆叠了 Mamba 层和前馈层,但得到的模型相比于单纯的 Mamba 还略有不及。该设计在图 1 中被记为 Mamba-MLP。

MoE-Mamba 分开了 Mamba 层执行的每个 token 的无条件处理和 MoE 层执行的有条件处理;其中的无条件处理可高效地将序列的整个上下文整合到一个内部表征中,而有条件处理可为每个 token 使用最相关的专家。这种将有条件处理和无条件处理交替起来的思路在一些基于 MoE 的模型中已经得到了应用,不过它们通常是交替基本的和 MoE 的前馈层。

主要结果

训练设置

该团队比较了 5 种不同设置:基本 Transformer、Mamba、Mamba-MLP、MoE 和 MoE-Mamba。

在大多数 Transformers 中,前馈层包含 8dm² 个参数,而 Mamba 论文中则让 Mamba 更小一些(约 6dm²),这样两个 Mamba 层的参数数量与一个前馈层和一个注意力层加起来差不多。为了让 Mamba 和新模型中每个 token 的活动参数数量大致一样,该团队将每个专家前向层的大小缩小到了 6dm²。除了嵌入层和解除嵌入(unembedding)层,所有模型都是每个 token 使用大约 2600 万参数。训练过程使用了 65 亿个 token,训练步骤数为 100k。

训练使用的数据集是 English C4 数据集,任务是预测下一 token。文本的 token 化使用了 GPT2 tokenizer。表 3 给出了超参数的完整列表。

结果

表 1 给出了训练结果。MoE-Mamba 的表现显著优于普通 Mamba 模型。

值得注意的是,MoE-Mamba 只用仅仅 46% 的训练步骤就达到了与普通 Mamba 同等的结果水平。由于学习率是针对普通 Mamba 进行调整的,因此可以预计,如果再针对 MoE-Mamba 对训练流程进行一番优化,MoE-Mamba 的表现还会更好。

消融研究

为了评估 Mamba 是否能随专家数量的增长而很好地扩展,研究者比较了使用不同数量专家的模型。

图 3 展示了使用不同数量的专家时的训练运行步骤情况。

表 2 给出了 100k 步骤后的结果。

这些结果表明新提出的方法能随专家数量而很好地扩展。如果专家的数量为 8 或更多,新模型的最终性能优于普通 Mamba。由于 Mamba-MLP 比普通 Mamba 差,可以预见使用少量专家的 MoE-Mamba 的性能表现会比 Mamba 差。当专家数为 32 时,新方法得到了最佳结果。

© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:[email protected]

微信扫码关注该文公众号作者

戳这里提交新闻线索和高质量文章给我们。
相关阅读
人气爆棚!家庭同乐日现场拍卖土地,成交率100%!Freecity与Castle Group强强联合,缔造悉尼西北明星项目!马斯克打脸OpenAI,全球最大巨无霸模型Grok-1开源!3140亿参数8个MoE,GitHub狂揽6k星2亿参数时序模型替代LLM?谷歌突破性研究被批「犯新手错误」支持百亿参数大模型、卢伟冰现场官宣小米首发,高通骁龙8s Gen3发布仅剩1席|别错过现场与MBB导师进行Mock Interview的难得机会,轻松应对令人头疼的案例面试!号称中文评测超越 GPT-4,百川智能发布超千亿参数大模型 Baichuan 3苹果iOS 18将搭载生成式AI;谷歌Bard最新排名超GPT-4;百川发超千亿参数大模型Baichuan 3丨AIGC大事日报20亿参数模型要“干掉”Llama13B,面壁智能:老子打的就是精锐大规模视觉模型新突破!70亿参数!Apple新作AIM:自回归图像模型日本最贵的“Ponshu”(日本酒)傅盛的AI大课(4):企业专用大模型,百亿参数就够了最新千亿大模型免费商用:1026亿参数,无需授权!诚邀开发者共同训练家居修理系列: 家用电器的修理(5)国产EDA强强联合,发力RISC-V芯片验证打通大模型训练任督二脉!国内首个千亿参数、全面开源大模型来了,还联手开发者共训OpenAI 忙着“宫斗”,手握 2.2 万块 H100 的竞争对手趁机发布新款大模型:1750 亿参数,性能仅次于 GPT-4打破MoE训练效率与性能瓶颈,华为盘古稀疏大模型全新架构LocMoE出炉最让我感动的伟大和奇迹AI早知道|微软Copilot将可免费使用GPT-4Turbo; 国产670亿参数大模型DeepSeek亮相iPhone 16或全系告别静音拨片/微信称好友上限约1万个/阿里云开源通义千问720亿参数模型| 灵感周报通义千问爆甩开源全家桶!最强720亿参数超车Llama 2,新上视觉模型看图直出代码傅盛揭穿OpenAI「阳谋」,引130万人围观!百亿参数模型帮老板打王者,一张3060就能跑6056 血壮山河 卢沟桥之变 25720亿参数大模型都拿来开源了!通义千问开源全家桶,最小18亿模型端侧都能跑国产720亿参数开源免费模型来了!对标Llama2 70B,一手实测在此强强联合,百度智能云牵手荣耀,启动大模型生态全面战略合作浪潮信息发布“源 2.0”,千亿参数大模型开源可商用AI早知道|360安全大模型3.0发布;通义听悟上线音视频问答助手;腾讯混元大模型参数规模扩展至万亿OpenAI忙着“宫斗”,手握2.2万块H100的竞争对手趁机发布新款大模型:1750亿参数,性能仅次于GPT-4大事件,QD中国与Lake Shore强强联合,低温测量新时代!强强联合!芯华章携手芯擎科技,软硬协同加速车规级芯片创新视觉全新主干!中科院&华为提出VMamba:视觉状态空间模型从《繁花》再看大剧营销:品牌如何和高品质内容强强联手?马斯克发全球最大开源模型!3140亿参数的Grok真打脸OpenAI了?风云人物郭文贵(1)
logo
联系我们隐私协议©2024 redian.news
Redian新闻
Redian.news刊载任何文章,不代表同意其说法或描述,仅为提供更多信息,也不构成任何建议。文章信息的合法性及真实性由其作者负责,与Redian.news及其运营公司无关。欢迎投稿,如发现稿件侵权,或作者不愿在本网发表文章,请版权拥有者通知本网处理。