大模型理论篇-Mixture-of-Experts架构
【大模型理论篇】–Mixture of Experts架构
Mixture of Experts（MoE，专家混合）【1】架构是一种神经网络架构，旨在通过有效分配计算负载来扩展模型规模。MoE架构通过在 推理和训练过程中仅使用部分“专家”（子模型），优化了资源利用率，从而能够处理复杂任务 。
在具体介绍MoE之前，先抛出MoE的一些表现【2】：
- 与密集模型相比，预训练速度快得多。
- 与具有相同参数数量的模型相比，推理速度更快。
- 由于所有Experts都加载在内存中，所以需要大量的显存。
- 在微调方面面临许多挑战，但最近关于 MoE 指令微调的工作很有前景。
1.MOE是什么？
MoE 是一种包含多个被称为“专家”子网的架构，每个子网专注于 数据任务或方面。其优势显著，在保持甚至提升模型质量的同时，能够以比相同或更大规模的传统模型更少的计算量进行预训练。例如，Mixtral 8x7B 在众多评估数据集上就超越了 LLaMA 2 70B。
这个概念可以类比为劳动分工，每个专家专注于某个大问题中的特定一小部分任务，从而生成更快、更准确的结果。
混合专家模型由三个关键组件组成：
- 专家（Experts）：专门针对特定任务的子模型。
- 门控网络（Gating Network）：一个选择器，它将输入数据路由到相关的专家。
- 稀疏激活（Sparse Activation）：只有少数专家针对每个输入被激活的方法，优化了计算效率。
传统transform架构：
MOE架构：
1.1.专家
在 MoE 架构中，专家是指训练好的子网络（神经网络或层），它们专门处理特定的数据或任务。例如，在图像分类任务中，一个专家可能专门识别纹理，而另一个专家可能识别边缘或形状。这种分工有助于整个模型更高效地处理问题，因为每个专家只处理它最适合的数据类型。
1.2.门控网络（路由器）
门控网络充当一个选择器，它决定将哪些输入数据发送给哪些专家。不是所有专家都同时工作，而是门控网络将数据路由到最相关的专家那里。类似于 token 选择路由策略，路由算法为每个 token 选择最佳的一个或两个专家。例如，在下图中，输入令牌 1，“我们”，被发送到第二个专家，而输入令牌 2，“喜欢”，被发送到第一个网络。
以下是一些主流的Token 路由技术：
- Top-k 路由：这是最简单的方法。门控网络选择亲和力得分（affinity score）最高的 top -k 个专家，并将输入数据发送给它们。
- 专家选择路由：与根据数据选择专家不同，这种方法由专家决定它们最能处理哪些数据。这种策略旨在实现最佳的负载均衡，并支持以多种方式将数据映射到专家。
1.3.稀疏激活
稀疏激活是 MoE 模型的关键部分和优势之一。与所有专家或参数对输入都活跃的密集模型不同，稀疏激活确保只有一小部分专家根据输入数据被激活。这种方法在保持性能的同时减少了计算需求，因为任何时候只有最相关的专家是活跃的。
- 稀疏路由：稀疏激活的一种特定技术，其中门控网络针对每个输入只激活少数专家。
1.4.总结
概括来看，在MoE模型中， 用混合专家层替换 Transformer 模型的每个前馈网络层，混合专家层由一个门控网络和一定数量的专家组成 。在MoE模型的核心是“专家”子网络。这些专家是更大神经网络中的独立模块，每个模块都能够处理输入数据。其概念是不同的专家专注于输入数据的不同方面，从而使模型能够有效利用专门的知识。
门控机制是一个关键组件，负责将输入引导至适当的专家网络。它根据一组门控值来决定每个专家的参与程度。门控机制可以实现为稠密或稀疏结构，其中稀疏结构由于选择性地激活部分专家，因而在计算上更加高效。
在稀疏MoE模型中，路由算法在决定哪些专家被激活以处理特定输入时起着重要作用。这些算法的复杂性各不相同，需考虑平衡模型的准确性和计算效率。路由算法的选择对模型的性能和推理速度有明显的影响。
2.MOE运作方式
首先引入混合专家的想法，其实很直接，就是希望不同的专家角色，关注不同的领域。例如 某部分的专家网络专注于数学能力的提升。另一部分专注编程能力的加强。这由此引出一个关键概念：稀疏性【2】。稀疏性利用了条件计算的理念。
在Dense模型中，所有输入都会使用所有参数，而稀疏性则允许仅运行系统的一部分。条件计算的思想（网络的一部分在每个示例的基础上被激活）使得可以在不增加计算量的情况下扩大模型的规模，因此，每个 MoE 层中使用了成千上万的专家。不过更多的专家可以提高样本效率并加快速度，但这些收益逐渐递减（尤其是在超过256或512个专家之后），而且推理时需要更多的显存。
在语言建模和 Transformer 的架构背景下，引入混合专家模型，需要注意专家们一次帮助预测一个token。路由器决定在每个token处哪个专家处于活动状态。专家网络通过门控线性组合来组合以预测下一个token。
输出：
Softmax 函数 是一种常用于多分类问题中的激活函数，它能够将一个向量转换成概率分布的形式。具体来说，Softmax 函数可以将任意实数向量映射到（0,1）区间内的概率分布，并且向量中的所有元素之和为1。
由于专家网络在一个 Transformer 模块内，路由是逐个token进行的。当token X 进入时，Transformer 模块的工作是预测下一个token Y。Transformer 模块通过自注意力机制可以访问所有先前的上下文。然后自注意力机制的输出作为条件，路由器判断输入到哪k个专家网络。
那么路由是如何工作的？实际上，路由只是一个位于注意力头和 K 个专家集合之间的 softmax 层。这个层学习哪种专家最适合哪种类型的token。
门控网络：
举个例子：
例如，序列[“I”、“love”、“LLM”、“#s”]中的每个标记都会通过路由函数，形成一个专家概率分布。 然后，我们为每个单个标记选择前 K 名专家–同一序列中的标记并不总是发送给相同的专家。
3.MOE代码实现
class Expert(nn.Module):
def __init__(self, input_size, output_size):
super(Expert, self).__init__()
self.linear = nn.Linear(input_size, output_size)
def forward(self, data):
x = self.linear(data)
return x
class GatingNetwork(nn.Module):
def __init__(self, input_size, num_experts):
super(GatingNetwork, self).__init__()
self.linear1 = nn.Linear(input_size, 4)
self.relu = nn.ReLU()
self.linear2 = nn.Linear(4, num_experts)
self.softmax = nn.Softmax(dim=-1)
def forward(self, data):
x = self.linear1(data)
x = self.relu(x)
x = self.linear2(x)
x = self.softmax(x)
return x
class MixtureOfExperts(nn.Module):
def __init__(self, num_experts=2):
super(MixtureOfExperts, self).__init__()
self.expert1 = Expert(2,1)
self.expert2 = Expert(2,1)
self.gating = GatingNetwork(2, num_experts)
self.sigmoid = nn.Sigmoid()
def forward(self, data):
expert1_output = self.expert1(data)
expert2_output = self.expert2(data)
gating_output = self.gating(data)
s = (expert1_output*gating_output[:,0][:,None] +
expert2_output*gating_output[:,1][:,None])
a = self.sigmoid(s)
return a
def backward(self, y_hat, labels, criterion, optimizer):
optimizer.zero_grad()
loss = criterion(y_hat, labels)
loss.backward()
optimizer.step()
return loss.item()
# Define the model, loss, and optimizer
moe = MixtureOfExperts()
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(moe.parameters(),lr=0.01)
# Define the learning rate scheduler
scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer)
# Convert data and labels to float tensors
data_tensor = data.float()
labels_tensor = labels.view(-1, 1).float()
# Training loop
num_epochs = 500
for epoch in tqdm(range(num_epochs)):
# Forward pass
y_hat = moe.forward(data)
# Backward pass and optimization
loss_value = moe.backward(y_hat, labels_tensor, criterion, optimizer)
# Decay the learning rate
scheduler.step()
expert1_weights = moe.expert1.linear.weight.detach()[0,0]
expert2_weights = moe.expert2.linear.weight.detach()[0,0]
expert1_bias = moe.expert1.linear.bias.detach()
expert2_bias = moe.expert2.linear.bias.detach()
gating_weights = moe.gating.linear2.weight.detach().flatten()
x_line = np.linspace(min(data[:, 0]), max(data[:, 0]), 100)
y_line1 = expert1_weights * x_line + 5
y_line2 = expert2_weights * x_line + 5
class_0 = data[labels == 0]
class_1 = data[labels == 1]
plt.scatter(class_0[:, 0], class_0[:, 1], label='Class 0', alpha=0.5)
plt.scatter(class_1[:, 0], class_1[:, 1], label='Class 1', alpha=0.5)
plt.plot(x_line, y_line1, label='Expert 1', alpha = 1)
plt.plot(x_line, y_line2, label='Expert 2', alpha = 1)
plt.xlabel('Feature 1')
plt.ylabel('Feature 2')
plt.grid(True, c='gray')
plt.legend()
plt.autoscale()
plt.show()4.参考文献
【1】Fedus, William, Barret Zoph, and Noam Shazeer. “Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.” Journal of Machine Learning Research 23.120 (2022): 1-39.
【2】
【3】
【4】
【5】
【6】
【7】
【8】
- 【9】
- THE SPARSELY-GATED MIXTURE-OF-EXPERTS LAYER
【10】
【11】
【12】
【13】