每个环节都详细解读，还有配套代码

克雷西 发自 凹非寺

量子位 | 公众号 QbitAI

传说中GPT-4的“致胜法宝”——MoE（混合专家）架构，自己也能手搓了！

Hugging Face上有一位机器学习大神，分享了如何从头开始建立一套完整的MoE系统。

这个项目被作者叫做MakeMoE，详细讲述了从注意力构建到形成完整MoE模型的过程。

作者介绍，MakeMoE是受到OpenAI创始成员Andrej Karpathy的makemore启发并以之为基础编写的。

makemore是一个针对自然语言处理和机器学习的教学项目，意在帮助学习者理解并实现一些基本模型。

同样，MakeMoE也是在一步步的搭建过程中，帮助学习者更深刻地理解混合专家模型。

那么，这份“手搓攻略”具体都讲了些什么呢？

从头开始搭建MoE模型

和Karpathy的makemore相比，MakeMoE用稀疏的专家混合体代替了孤立的前馈神经网络，同时加入了必要的门控逻辑。

同时，由于过程中需要用到ReLU激活函数，makemore中的默认初始化方式被替换成了Kaiming He方法。

想要创建一个MoE模型，首先要理解自注意力机制。

模型首先通过线性变换，将输入序列变换成用查询（Q）、键（K）和值（V）表示的参数。

这些参数随后被用于计算注意力分数，这些分数决定了在生成每个token时，模型对序列中每个位置的关注程度。

为了确保模型在生成文本时的自回归特性，即只能基于已经生成的token来预测下一个token，作者使用了多头因果自注意力机制。

这种机制通过一个掩码来实现将未处理的位置的注意力分数设置为负无穷大，这样这些位置的权重就会变为零。

多头因果则是让模型并行地执行多个这样的注意力计算，每个头关注序列的不同部分。

完成自注意力机制的配置后，就可以创建专家模块了，这里的“专家模块”是一种多层感知器。

每个专家模块包含一个线性层，它将嵌入向量映射到一个更大的维度，然后通过非线性激活函数（如ReLU），再通过另一个线性层将向量映射回原始的嵌入维度。

这样的设计使得每个专家能够专注于处理输入序列的不同部分，并通过门控网络来决定在生成每个token时应该激活哪些专家。

于是，接下来就要开始构建分配和管理专家的组件——门控网络。

这里的门控网络同样是通过一个线性层实现，该层将自注意力层的输出映射到专家模块的数量。

这个线性层的输出是一个分数向量，每个分数代表了对应专家模块对于当前处理的token的重要性。

门控网络会计算这个分数向量的top-k值并记录其索引，然后从中选择top-k个最大的分数，用来加权对应的专家模块输出。

为了在训练过程中增加模型的探索性，作者还引入了噪声，避免所有token都倾向于被相同的专家处理。

这种噪声通常通过在分数向量上添加随机的高斯噪声实现。

获得结果后，模型有选择地将前k个值与相应token的前k个专家的输出相乘，然后相加形成加权和，构成模型的输出。

最后，将这些模块在一起，就得到一个MoE模型了。

针对以上的整个过程，作者都提供了相应的代码，可以到原文中具体了解。

另外，作者还制作了端到端的Jupyter笔记，可以在学习各模块的同时直接运行。

感兴趣的话，就赶快学起来吧！

原文地址：
https://huggingface.co/blog/AviSoori1x/makemoe-from-scratch
笔记版本（GitHub）：
https://github.com/AviSoori1x/makeMoE/tree/main