一本读懂大模型MoE混合专家模型原理

MoE,全称 Mixture of Experts,混合专家模型。有很多人猜测 GPT-4 就是使用了 MoE 模型,让很多模型并行起来,加速了推理。阿里巴巴的通义千问大模型(Qwen)也是使用了 MoE,效果非常好。

一本读懂大模型MoE混合专家模型原理

本篇文章只给完全不明白 MoE 的同学介绍一下它的基本原理。

NO.1 动机:为什么要提出 MoE?

MoE 之前的模型有什么问题?

在模型预测下一个 token 时,需要很多神经元进行计算,比如下图所示的 8×8 神经元(别数了,图片是 7×8,意思一下)。
但是呢根据经验,对结果起作用的重要的神经元,可能很少(如下图黄色的点)。

一本读懂大模型MoE混合专家模型原理

原来的权重矩阵,起关键作用的可能是很少一部分。

所以呢,我们就想,是不是可以把这些点进行重排(如下图),这样我们就能得到 8 个子模型(对应 8 列)。
这样的话如果你想查询某个特定的专家知识,你就查询矩阵对应的列,然后只跟这一列的神经元计算就好了,这样就避免了与整个矩阵进行计算。

一本读懂大模型MoE混合专家模型原理

如果我们能把相关的神经元排成一排就好了。

NO.2 MoE 是什么?

接上步,怎么把相关的神经元都排到同一列呢?方法是,设计一个 Router(中文直译是路由器),让它来决定要把数据交给哪个子模型(也叫“专家”)做推理(如下图)。这就是 MoE 的关键操作。

一本读懂大模型MoE混合专家模型原理

这个 Router 也是一个映射,把一些特征向量输入给它,它会输出一个 one-hot 向量,来决定哪个子网络处理数据。

一本读懂大模型MoE混合专家模型原理

所以推理时,inputs 会先经过 Router,确认去哪个子网络之后,再经过那个选中的自网络,得到最终的 output tokens(如下图)。
训练时,根据 loss 更新对应的子网络模型(也就是“专家”)参数和 Router 参数即可。此时的 loss 可以简略写成 Loss=Actual−Predictedtokens。

一本读懂大模型MoE混合专家模型原理

如此这般,虽然模型与原来有同样的参数量,但是在推理时,由于只有一个“专家”参与了计算,所以推理速度要比原来快很多。

NO.3 这么做有什么问题?

然而,按照上面的方法训练会有一个很大的问题,那就是,训练出来的结果有可能一个专家非常强而其他专家非常弱。

因为训练过程中,当某个专家比较强时,Router 就更倾向于把数据分给这个专家,使得这个专家更强,导致各个专家不能平衡地学习,从而影响了模型效果。怎么办呢?

一种方法是给 Router 增加噪声,比如 Router 预测出来是 [0,0,0,0,0,0,1,0],那我们增加噪声让它变成 [0.04, 0.01, 0.04, 0.06, 0.03, 0.05, 0.75, 0.02],然后在这个分布上随机采样。

这样可以让数据遍布所有专家,避免了某些专家没有数据而另一个专家被着重训练的“旱的旱死、涝的涝死”的情况。

另一种方法是,给 Router 的选择赋一个惩罚。在原来的 loss 上增加一项,如果 Router 没有把数据平均分配给所有专家,我们就惩罚它:Loss=Actual−Predictedtokens+Penality 。这样就能促使 Router 尽量均匀把数据分配给每个专家。

NO.4 MoE 有什么影响?

(1)MoE 能减少内存消耗吗?

答案是不能。假设你在边缘设备上(比如你的笔记本电脑)运行 LLaMA 7b 推理,你把它量化成 Int4 模型。

如果把模型改成 MoE 形式,由于模型参数量与原来是一致的,所以你推理时占用的显存不会变。但是由于你只需要在一个专家上计算,所以计算推理速度会更快。

(2)MoE 会减少计算开销吗?

答案是能。你可以把不同专家部署在不同 GPU 上,比如这里我们把 8 个专家放在 8 个 GPU 上。

如果有好几个用户同时发来消息时,可以经过 Router,把它们分组给不同专家,然后由各自 GPU 上的专家分别计算,算完了再收集起来做下一次 Router 判断,再分配和计算,重复这个过程。

(3)MoE 会让训练变慢吗?

会的。因为给 Router 增加噪声后,会把本该送给某个专家的数据送给了另一个专家,让专家们训练和收敛速度变慢。

【另一个改进】所以有人提出了加速训练的方法,叫 fast feed forward networks(快速推理网络) 或者 binary tree networks(二叉树网络)。

它不再只用一个 Router 给 8 个专家发数据,而是用多个 Router 逐层往下发数据,每个 Router 只需要预测出一个概率值 p,由 p 来决定是往左下发还是往右下发(0/1)。

这样就能天然有一个比较均匀的前馈分布,而不需要人为增加噪声和惩罚。这样就能让 MoE 训练既能跟标准 GPT 一样快,又能在推理时比原来的 GPT 更快。

一本读懂大模型MoE混合专家模型原理

NO.5 MoE 的实验结果

下表中,一个专家的宽度是 32 个,所以 width = 64 表示有两个专家,width = 128 表示有 4 个专家。

在 fast feedforward 中 w=128 则说明二叉树有两层,一共 3 个 Router 下发数据给 4 个专家。

下表中的 MA 表示记忆能力,ETT 表示训练的 epoch 数量,GA 表示生成能力指标。

一本读懂大模型MoE混合专家模型原理

普通模型、MoE、FastFeedforward 训练速度对比

从表格中有数字的第一行中可以看出,单单 MoE(中间列) 比 feedforward(左边列) 训练 ETT 多了很多轮,但效果 MA 远不如 feedforward。

但是 fast feedforward 的训练论数和指标效果,跟 feadforward 差不多。说明 fast feedforward 在训练速度上非常关键。

下面这张表展示了 fast feedforward 的效果,l=32 表示有128/32=8个专家,l=1表示有 128/1=128 个专家。

一本读懂大模型MoE混合专家模型原理

可以看到到,随着 Router 二叉树层数的增加,速度越来越快,效果有些许下降。

这里有一点要提醒,上面这两个表的实验是在很小的网络上进行的,不确定 GPT-3 是否有这样的敏感度。

NO.6 总结 MoE

宏观上看,对于 MoE,每个神经元原本可以连接到下一层的其他所有神经元,但是 MoE 把它们分组了,每个神经元只能连接到本组内的下一层的神经元。

有一点点像 multi-heads,但是 multi-heads 分组是为了增强各组内的特征,并没有减少计算量;MoE 在推理时显著减少的计算量。

结果是,MoE 只损失一点点精度,就能让推理大约快好几倍,非常有利于大模型部署,赞!

一本读懂大模型MoE混合专家模型原理

红色箭头表示原始模型的神经元连接,黄色箭头表示 MoE 的神经元连接。

到现在,MoE 已经发展出了各种方法,比如 LLaMA-MoE 模型,大家可以自行了解。

一本读懂大模型MoE混合专家模型原理

 

一本读懂大模型MoE混合专家模型原理

极牛网精选文章《一本读懂大模型MoE混合专家模型原理》文中所述为作者独立观点,不代表极牛网立场。如若转载请注明出处:https://geeknb.com/27815.html

(30)
打赏 微信公众号 微信公众号 微信小程序 微信小程序
主编的头像主编认证作者
上一篇 2024年9月20日 上午11:02
下一篇 2024年9月20日 下午12:03

相关推荐

发表回复

登录后才能评论
扫码关注
扫码关注
分享本页
返回顶部