网站后台 模板,小程序网站app定制开发,wordpress文章页标题优化,哪些广告平台留号码注#xff1a;本文为 “稀疏混合专家架构语言模型#xff08;MoE#xff09;” 相关文章合辑。 手把手教你#xff0c;从零开始实现一个稀疏混合专家架构语言模型#xff08;MoE#xff09;
机器之心 2024年02月11日 12:21 河南
选自huggingface 机器之心编译 机器之心…注本文为 “稀疏混合专家架构语言模型MoE” 相关文章合辑。 手把手教你从零开始实现一个稀疏混合专家架构语言模型MoE
机器之心 2024年02月11日 12:21 河南
选自huggingface 机器之心编译 机器之心编辑部 本文介绍了实现一个稀疏混合专家语言模型MoE的方法详细解释了模型的实施过程包括采用稀疏混合专家取代传统的前馈神经网络实现 top-k 门控和带噪声的 top-k 门控以及采用 Kaiming He 初始化技术。作者还说明了从 makemore 架构保持不变的元素比如数据集处理、分词预处理和语言建模任务。最后还提供了一个 GitHub 仓库链接用于实现模型的整个过程是一本不可多得的实战教科书。 内容简介
在混合专家模型 Mixtral 发布后混合专家模型MoE越来越受到人们的关注。在稀疏化的混合专家语言模型中大部分组件都与传统的 transformers 相同。然而尽管看似简单但经验表明稀疏混合专家语言模型训练的稳定性还存在着一些问题。
像这样易于修改的小规模实现可能有助于快速试验新方法。Hugging Face 上的一篇博客介绍了一种可配置的小规模稀疏 MoE 实施方法也许有助于打算在这个方向深耕的研究者们进行快速试验自己的新方法并且给出了基于 PyTorch 的详细代码https://github.com/AviSoori1x/makeMoE/tree/main
机器之心对此进行了整理以飨读者。
本文在 makemore 架构的基础上进行了几处更改
使用稀疏混合专家代替单独的前馈神经网络Top-k 门控和有噪声的 Top-k 门控参数初始化使用了 Kaiming He 初始化方法但本文的重点是可以对初始化方法进行自定义这样就可以在 Xavier/Glorot 等初始化中进行选择。
同时以下模块与 makemore 保持一致
数据集、预处理分词部分以及 Andrej 最初选择的语言建模任务 - 生成莎士比亚文风的文本内容Casusal 自注意力机制训练循环推理逻辑 接下来逐步介绍实施方案先从注意力机制开始。
因果缩放点积注意力机制 下面这段代码展示了自注意力机制的基本概念并且侧重于使用经典的缩放点积自注意力scaled dot product self-attention.实现。在这一自注意力变体机制中查询矩阵、键矩阵和值矩阵都来自相同的输入序列。同时为了确保自回归语言生成过程的完整性特别是在纯解码器模型中使用了一种掩码机制。
这种掩码机制非常关键因为它可以掩盖当前 token 所处位置之后的任何信息从而引导模型只关注序列的前面部分。这种了遮挡 token 后面内容的注意力被称为因果自注意力。值得注意的是稀疏混合专家模型并不局限于仅有解码器的 Transformer 架构。事实上这一领域的许多重要的成果都是围绕 T5 架构展开的T5 架构也包含了 Transformer 模型中的编码器和解码器组件。
#This code is borrowed from Andrej Karpathys makemore repository linked in the repo.
The self attention layers in Sparse mixture of experts models are the same as
in regular transformer modelstorch.manual_seed(1337)
B,T,C 4,8,32 # batch, time, channels
x torch.randn(B,T,C)# lets see a single Head perform self-attention
head_size 16
key nn.Linear(C, head_size, biasFalse)
query nn.Linear(C, head_size, biasFalse)
value nn.Linear(C, head_size, biasFalse)
k key(x) # (B, T, 16)
q query(x) # (B, T, 16)
wei q k.transpose(-2, -1) # (B, T, 16) (B, 16, T) --- (B, T, T)tril torch.tril(torch.ones(T, T))
#wei torch.zeros((T,T))
wei wei.masked_fill(tril 0, float(-inf))
wei F.softmax(wei, dim-1) #B,T,Tv value(x) #B,T,H
out wei v # (B,T,T) (B,T,H) - (B,T,H)
out.shapetorch.Size([4, 8, 16])然后因果自注意力和多头因果自注意力的代码可整理如下。多头自注意力并行应用多个注意力头每个注意力头单独关注通道的一个部分嵌入维度。多头自注意力从本质上改善了学习过程并由于其固有的并行能力提高了模型训练的效率。下面这段代码使用了 dropout 来进行正则化来防止过拟合。
#Causal scaled dot product self-Attention Head
n_embd 64
n_head 4
n_layer 4
head_size 16
dropout 0.1class Head(nn.Module): one head of self-attention def __init__(self, head_size):super().__init__()self.key nn.Linear(n_embd, head_size, biasFalse)self.query nn.Linear(n_embd, head_size, biasFalse)self.value nn.Linear(n_embd, head_size, biasFalse)self.register_buffer(tril, torch.tril(torch.ones(block_size, block_size)))self.dropout nn.Dropout(dropout)def forward(self, x):B,T,C x.shapek self.key(x) # (B,T,C)q self.query(x) # (B,T,C)# compute attention scores (affinities)wei q k.transpose(-2,-1) * C**-0.5 # (B, T, C) (B, C, T) - (B, T, T)wei wei.masked_fill(self.tril[:T, :T] 0, float(-inf)) # (B, T, T)wei F.softmax(wei, dim-1) # (B, T, T)wei self.dropout(wei)# perform the weighted aggregation of the valuesv self.value(x) # (B,T,C)out wei v # (B, T, T) (B, T, C) - (B, T, C)return out多头自注意力的实现方式如下 #Multi-Headed Self Attention
class MultiHeadAttention(nn.Module): multiple heads of self-attention in parallel def __init__(self, num_heads, head_size):super().__init__()self.heads nn.ModuleList([Head(head_size) for _ in range(num_heads)])self.proj nn.Linear(n_embd, n_embd)self.dropout nn.Dropout(dropout)def forward(self, x):out torch.cat([h(x) for h in self.heads], dim-1)out self.dropout(self.proj(out))return out创建一个专家模块即一个简单的多层感知器
在稀疏混合专家架构中每个 transformer 区块内的自注意力机制保持不变。不过每个区块的结构发生了巨大的变化标准的前馈神经网络被多个稀疏激活的前馈网络即专家网络所取代。所谓「稀疏激活」是指序列中的每个 token 只被分配给有限数量的专家通常是一个或两个。
这有助于提高训练和推理速度因为每次前向传递都会激活少数专家。不过所有专家都必须存在 GPU 内存中因此当参数总数达到数千亿甚至数万亿时就会产生部署方面的问题。 #Expert module
class Expert(nn.Module): An MLP is a simple linear layer followed by a non-linearity i.e. each Expert def __init__(self, n_embd):super().__init__()self.net nn.Sequential(nn.Linear(n_embd, 4 * n_embd),nn.ReLU(),nn.Linear(4 * n_embd, n_embd),nn.Dropout(dropout),)def forward(self, x):return self.net(x)Top-k 门控的一个例子 门控网络也称为路由确定哪个专家网络接收来自多头注意力的 token 的输出。举个例子解释路由的机制假设有 4 个专家token 需要被路由到前 2 个专家中。首先需要通过线性层将 token 输入到门控网络中。该层将对应于Batch sizeTokensn_embed的输入张量从2432维度投影到对应于Batch size、Tokensnum_expert的新形状2、44。其中 n_embed 是输入的通道维度num_experts 是专家网络的计数。
接下来沿最后一个维度找出最大的前两个值及其相应的索引。
#Understanding how gating works
num_experts 4
top_k2
n_embed32#Example multi-head attention output for a simple illustrative example, consider n_embed32, context_length4 and batch_size2
mh_output torch.randn(2, 4, n_embed)topkgate_linear nn.Linear(n_embed, num_experts) # nn.Linear(32, 4)logits topkgate_linear(mh_output)
top_k_logits, top_k_indices logits.topk(top_k, dim-1) # Get top-k experts
top_k_logits, top_k_indices#output:
(tensor([[[ 0.0246, -0.0190],[ 0.1991, 0.1513],[ 0.9749, 0.7185],[ 0.4406, -0.8357]],[[ 0.6206, -0.0503],[ 0.8635, 0.3784],[ 0.6828, 0.5972],[ 0.4743, 0.3420]]], grad_fnTopkBackward0),tensor([[[2, 3],[2, 1],[3, 1],[2, 1]],[[0, 2],[0, 3],[3, 2],[3, 0]]]))通过仅保留沿最后一个维度进行比较的前 k 大的值来获得稀疏门控的输出。用负无穷值填充其余部分在使用 softmax 激活函数。负无穷会被映射至零而最大的前两个值会更加突出且和为 1。要求和为 1 是为了对专家输出的内容进行加权。
zeros torch.full_like(logits, float(-inf)) #full_like clones a tensor and fills it with a specified value (like infinity) for masking or calculations.
sparse_logits zeros.scatter(-1, top_k_indices, top_k_logits)
sparse_logits#output
tensor([[[ -inf, -inf, 0.0246, -0.0190],[ -inf, 0.1513, 0.1991, -inf],[ -inf, 0.7185, -inf, 0.9749],[ -inf, -0.8357, 0.4406, -inf]],[[ 0.6206, -inf, -0.0503, -inf],[ 0.8635, -inf, -inf, 0.3784],[ -inf, -inf, 0.5972, 0.6828],[ 0.3420, -inf, -inf, 0.4743]]], grad_fnScatterBackward0)gating_output F.softmax(sparse_logits, dim-1)
gating_output#ouput
tensor([[[0.0000, 0.0000, 0.5109, 0.4891],[0.0000, 0.4881, 0.5119, 0.0000],[0.0000, 0.4362, 0.0000, 0.5638],[0.0000, 0.2182, 0.7818, 0.0000]],[[0.6617, 0.0000, 0.3383, 0.0000],[0.6190, 0.0000, 0.0000, 0.3810],[0.0000, 0.0000, 0.4786, 0.5214],[0.4670, 0.0000, 0.0000, 0.5330]]], grad_fnSoftmaxBackward0)使用有噪声的 top-k 门控以实现负载平衡
# First define the top k router module
class TopkRouter(nn.Module):def __init__(self, n_embed, num_experts, top_k):super(TopkRouter, self).__init__()self.top_k top_kself.linear nn.Linear(n_embed, num_experts)def forward(self, mh_ouput):# mh_ouput is the output tensor from multihead self attention blocklogits self.linear(mh_output)top_k_logits, indices logits.topk(self.top_k, dim-1)zeros torch.full_like(logits, float(-inf))sparse_logits zeros.scatter(-1, indices, top_k_logits)router_output F.softmax(sparse_logits, dim-1)return router_output, indices接下来使用下面这段代码来测试程序
#Testing this out:
num_experts 4
top_k 2
n_embd 32mh_output torch.randn(2, 4, n_embd) # Example input
top_k_gate TopkRouter(n_embd, num_experts, top_k)
gating_output, indices top_k_gate(mh_output)
gating_output.shape, gating_output, indices
#And it works!!#output
(torch.Size([2, 4, 4]),tensor([[[0.5284, 0.0000, 0.4716, 0.0000],[0.0000, 0.4592, 0.0000, 0.5408],[0.0000, 0.3529, 0.0000, 0.6471],[0.3948, 0.0000, 0.0000, 0.6052]],[[0.0000, 0.5950, 0.4050, 0.0000],[0.4456, 0.0000, 0.5544, 0.0000],[0.7208, 0.0000, 0.0000, 0.2792],[0.0000, 0.0000, 0.5659, 0.4341]]], grad_fnSoftmaxBackward0),tensor([[[0, 2],[3, 1],[3, 1],[3, 0]],[[1, 2],[2, 0],[0, 3],[2, 3]]]))尽管最近发布的 mixtral 的论文没有提到这一点但本文的作者相信有噪声的 Top-k 门控机制是训练 MoE 模型的一个重要工具。从本质上讲不会希望所有的 token 都发送给同一组「受欢迎」的专家网络。人们需要的是能在开发和探索之间取得良好平衡。为此为了负载平衡从门控的线性层向 logits 激活函数添加标准正态噪声是有帮助的这使训练更有效率。 #Changing the above to accomodate noisy top-k gating
class NoisyTopkRouter(nn.Module):def __init__(self, n_embed, num_experts, top_k):super(NoisyTopkRouter, self).__init__()self.top_k top_k#layer for router logitsself.topkroute_linear nn.Linear(n_embed, num_experts)self.noise_linear nn.Linear(n_embed, num_experts)def forward(self, mh_output):# mh_ouput is the output tensor from multihead self attention blocklogits self.topkroute_linear(mh_output)#Noise logitsnoise_logits self.noise_linear(mh_output)#Adding scaled unit gaussian noise to the logitsnoise torch.randn_like(logits)*F.softplus(noise_logits)noisy_logits logits noisetop_k_logits, indices noisy_logits.topk(self.top_k, dim-1)zeros torch.full_like(noisy_logits, float(-inf))sparse_logits zeros.scatter(-1, indices, top_k_logits)router_output F.softmax(sparse_logits, dim-1)return router_output, indices再次尝试代码
#Testing this out, again:
num_experts 8
top_k 2
n_embd 16mh_output torch.randn(2, 4, n_embd) # Example input
noisy_top_k_gate NoisyTopkRouter(n_embd, num_experts, top_k)
gating_output, indices noisy_top_k_gate(mh_output)
gating_output.shape, gating_output, indices
#It works!!#output
(torch.Size([2, 4, 8]),tensor([[[0.4181, 0.0000, 0.5819, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],[0.4693, 0.5307, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],[0.0000, 0.4985, 0.5015, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],[0.0000, 0.0000, 0.0000, 0.2641, 0.0000, 0.7359, 0.0000, 0.0000]],[[0.0000, 0.0000, 0.0000, 0.6301, 0.0000, 0.3699, 0.0000, 0.0000],[0.0000, 0.0000, 0.0000, 0.4766, 0.0000, 0.0000, 0.0000, 0.5234],[0.0000, 0.0000, 0.0000, 0.6815, 0.0000, 0.0000, 0.3185, 0.0000],[0.4482, 0.5518, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]]],grad_fnSoftmaxBackward0),tensor([[[2, 0],[1, 0],[2, 1],[5, 3]],[[3, 5],[7, 3],[3, 6],[1, 0]]]))创建稀疏化的混合专家模块
在获得门控网络的输出结果之后对于给定的 token将前 k 个值选择性地与来自相应的前 k 个专家的输出相乘。这种选择性乘法的结果是一个加权和该加权和构成 SparseMoe 模块的输出。这个过程的关键和难点是避免不必要的乘法运算只为前 k 名专家进行正向转播。为每个专家执行前向传播将破坏使用稀疏 MoE 的目的因为这个过程将不再是稀疏的。
class SparseMoE(nn.Module):def __init__(self, n_embed, num_experts, top_k):super(SparseMoE, self).__init__()self.router NoisyTopkRouter(n_embed, num_experts, top_k)self.experts nn.ModuleList([Expert(n_embed) for _ in range(num_experts)])self.top_k top_kdef forward(self, x):gating_output, indices self.router(x)final_output torch.zeros_like(x)# Reshape inputs for batch processingflat_x x.view(-1, x.size(-1))flat_gating_output gating_output.view(-1, gating_output.size(-1))# Process each expert in parallelfor i, expert in enumerate(self.experts):# Create a mask for the inputs where the current expert is in top-kexpert_mask (indices i).any(dim-1)flat_mask expert_mask.view(-1)if flat_mask.any():expert_input flat_x[flat_mask]expert_output expert(expert_input)# Extract and apply gating scoresgating_scores flat_gating_output[flat_mask, i].unsqueeze(1)weighted_output expert_output * gating_scores# Update final output additively by indexing and addingfinal_output[expert_mask] weighted_output.squeeze(1)return final_output运行以下代码来用样本测试上述实现可以看到确实如此
import torch
import torch.nn as nn#Lets test this out
num_experts 8
top_k 2
n_embd 16
dropout0.1mh_output torch.randn(4, 8, n_embd) # Example multi-head attention output
sparse_moe SparseMoE(n_embd, num_experts, top_k)
final_output sparse_moe(mh_output)
print(Shape of the final output:, final_output.shape)Shape of the final output: torch.Size([4, 8, 16])需要强调的是如上代码所示从路由 / 门控网络输出的 top_k 本身也很重要。索引确定了被激活的专家是哪些 对应的值又决定了权重大小。下图进一步解释了加权求和的概念。 模块整合
将多头自注意力和稀疏混合专家相结合形成稀疏混合专家 transformer 块。就像在 vanilla transformer 块中一样也要使用残差以确保训练稳定并避免梯度消失等问题。此外要采用层归一化来进一步稳定学习过程。
#Create a self attention mixture of experts block, that may be repeated several number of times
class Block(nn.Module): Mixture of Experts Transformer block: communication followed by computation (multi-head self attention SparseMoE) def __init__(self, n_embed, n_head, num_experts, top_k):# n_embed: embedding dimension, n_head: the number of heads wed likesuper().__init__()head_size n_embed // n_headself.sa MultiHeadAttention(n_head, head_size)self.smoe SparseMoE(n_embed, num_experts, top_k)self.ln1 nn.LayerNorm(n_embed)self.ln2 nn.LayerNorm(n_embed)def forward(self, x):x x self.sa(self.ln1(x))x x self.smoe(self.ln2(x))return x最后将所有内容整合在一起形成稀疏混合专家语言模型。
class SparseMoELanguageModel(nn.Module):def __init__(self):super().__init__()# each token directly reads off the logits for the next token from a lookup tableself.token_embedding_table nn.Embedding(vocab_size, n_embed)self.position_embedding_table nn.Embedding(block_size, n_embed)self.blocks nn.Sequential(*[Block(n_embed, n_headn_head, num_expertsnum_experts,top_ktop_k) for _ in range(n_layer)])self.ln_f nn.LayerNorm(n_embed) # final layer normself.lm_head nn.Linear(n_embed, vocab_size)def forward(self, idx, targetsNone):B, T idx.shape# idx and targets are both (B,T) tensor of integerstok_emb self.token_embedding_table(idx) # (B,T,C)pos_emb self.position_embedding_table(torch.arange(T, devicedevice)) # (T,C)x tok_emb pos_emb # (B,T,C)x self.blocks(x) # (B,T,C)x self.ln_f(x) # (B,T,C)logits self.lm_head(x) # (B,T,vocab_size)if targets is None:loss Noneelse:B, T, C logits.shapelogits logits.view(B*T, C)targets targets.view(B*T)loss F.cross_entropy(logits, targets)return logits, lossdef generate(self, idx, max_new_tokens):# idx is (B, T) array of indices in the current contextfor _ in range(max_new_tokens):# crop idx to the last block_size tokensidx_cond idx[:, -block_size:]# get the predictionslogits, loss self(idx_cond)# focus only on the last time steplogits logits[:, -1, :] # becomes (B, C)# apply softmax to get probabilitiesprobs F.softmax(logits, dim-1) # (B, C)# sample from the distributionidx_next torch.multinomial(probs, num_samples1) # (B, 1)# append sampled index to the running sequenceidx torch.cat((idx, idx_next), dim1) # (B, T1)return idx参数初始化对于深度神经网络的高效训练非常重要。由于专家中存在 ReLU 激活因此这里使用了 Kaiming He 初始化。也可以尝试在 transformer 中更常用的 Glorot 初始化。杰里米 - 霍华德Jeremy Howard的《Fastai》第 2 部分有一个从头开始实现这些功能的精彩讲座https://course.fast.ai/Lessons/lesson17.html
Glorot 参数初始化通常被用于 transformer 模型因此这是一个可能提高模型性能的方法。
def kaiming_init_weights(m):if isinstance (m, (nn.Linear)): init.kaiming_normal_(m.weight)model SparseMoELanguageModel()
model.apply(kaiming_init_weights)本文作者使用 mlflow 跟踪并记录重要指标和训练超参数。 #Using MLFlow
m model.to(device)
# print the number of parameters in the model
print(sum(p.numel() for p in m.parameters())/1e6, M parameters)# create a PyTorch optimizer
optimizer torch.optim.AdamW(model.parameters(), lrlearning_rate)
#mlflow.set_experiment(makeMoE)
with mlflow.start_run():#If you use mlflow.autolog() this will be automatically logged. I chose to explicitly log here for completenessparams {batch_size: batch_size , block_size : block_size, max_iters: max_iters, eval_interval: eval_interval,learning_rate: learning_rate, device: device, eval_iters: eval_iters, dropout : dropout, num_experts: num_experts, top_k: top_k }mlflow.log_params(params)for iter in range(max_iters):# every once in a while evaluate the loss on train and val setsif iter % eval_interval 0 or iter max_iters - 1:losses estimate_loss()print(fstep {iter}: train loss {losses[train]:.4f}, val loss {losses[val]:.4f})metrics {train_loss: losses[train], val_loss: losses[val]}mlflow.log_metrics(metrics, stepiter)# sample a batch of dataxb, yb get_batch(train)# evaluate the losslogits, loss model(xb, yb)optimizer.zero_grad(set_to_noneTrue)loss.backward()optimizer.step() 8.996545 M parameters
step 0: train loss 5.3223, val loss 5.3166
step 100: train loss 2.7351, val loss 2.7429
step 200: train loss 2.5125, val loss 2.5233
.
.
.step 4999: train loss 1.5712, val loss 1.7508记录训练和验证损失可以很好地指示训练的进展情况。该图显示可能应该在 4500 次时停止当验证损失稍微增加时 接下来可以使用这个模型逐字符自回归地生成文本。
# generate from the model. Not great. Not too bad either
context torch.zeros((1, 1), dtypetorch.long, devicedevice)
print(decode(m.generate(context, max_new_tokens2000)[0].tolist()))DUKE VINCENVENTIO:
If it ever fecond he town sue kigh now,
That thou woldst is steen t.SIMNA:
Angent her; no, my a born Yorthort,
Romeoos soun and lawf to your sawe with ch a woft ttastly defy,
To declay the soul art; and meart smad.CORPIOLLANUS:
Which I cannot shall do from by born und ot cold warrike,
What king we best anone wraves going of heard and good
Thus playvage; you have wold the grace.
...本文参考内容
在实施过程中作者大量参考了以下出版物
混合专家模型https://arxiv.org/pdf/2401.04088.pdf超大型神经网络稀疏门控混合专家层https://arxiv.org/pdf/1701.06538.pdf来自 Andrej Karpathy 的原始 makemore 实现https://github.com/karpathy/makemore
还可以尝试以下几种方法来提高模型性能
提高混合专家模块的效率尝试不同的神经网络初始化策略从字符级到子词级的分词对专家数量和 k 的取值每个 token 激活的专家数量进行贝叶斯超参数搜索。这可以归类为神经架构搜索。优化专家能力。 第一个100%开源的MoE大模型7B的参数1B的推理成本
机器之心 2024年09月05日 10:14 辽宁 机器之心报道 机器之心编辑部 训练代码、中间 checkpoint、训练日志和训练数据都已经开源。 尽管大语言模型 (LM) 在各种任务上取得了重大进展但在训练和推理方面性能和成本之间仍然需要权衡。
对于许多学者和开发人员来说高性能的 LM 是无法访问的因为它们的构建和部署成本过高。改善成本 - 性能的一种方法是使用稀疏激活混合专家 (MoE)。MoE 在每一层都有几个专家每次只激活其中的一个子集参见图 2。这使得 MoE 比具有相似参数量的密集模型更有效因为密集模型为每个输入激活所有参数。 出于这个原因行业前沿模型包括 Gemini-1.5、 GPT-4 等在内的模型都使用了 MoE。
然而大多数 MoE 模型都是闭源的虽然有些模型公开发布了模型权重但有关训练数据、代码等的信息却很有限甚至有些研究没有提供这些信息。由于缺乏开放资源和对研究细节的深入探索在 MoE 领域无法构建具有成本效益的开源模型从而接近闭源前沿模型的能力。
为了解决这些问题来自艾伦人工智能研究院、 Contextual AI 等机构的研究者引入了 OLMoE 这是一个完全开源的混合专家语言模型在类似大小的模型中具有 SOTA 性能。 论文地址https://arxiv.org/pdf/2409.02060论文标题OLMoE: Open Mixture-of-Experts Language Models
特别的该研究使用 5.1 万亿个 token 预训练了 OLMoE-1B-7B 模型该模型总共拥有 69 亿参数其中每个输入 token 只激活 13 亿参数。
结果是与使用具有约 1B 参数的密集模型例如 OLMo 1B 或 TinyLlama 1B 实现了类似的推理成本只是需要更多的 GPU 内存来存储约 7B 的总参数。实验表明MoE 的训练速度比具有等效激活参数的密集 LM 快 2 倍左右。
如图 1 所示OLMoE-1B-7B 显著优于所有开源 1B 模型并且与推理成本和内存存储明显更高的密集模型相比表现出了竞争力。 通过指令和偏好调优该研究还创建了 OLMoE-1B-7B-INSTRUCT它在常见基准 MMLU、GSM8k、HumanEval 等上超越了各种更大的指令模型包括 Llama2-13B-Chat 、OLMo-7B-Instruct (0724) 和 DeepSeekMoE-16B。
受控实验强调了 MoE见表 1和一般 LM 的关键设计选择。结果表明使 MoE 性能卓越的一个关键设计决策是使用细粒度路由和粒度专家granular experts在每一层使用 64 个小专家其中 8 个被激活。
此外路由算法的选择也很重要该研究发现无丢弃dropless基于 token 的路由优于基于专家的路由。最后该研究分析了 OLMoE-1B-7B 中的路由行为发现路由在预训练的早期就饱和了专家很少被共同激活并且专家表现出领域和词汇的专业化。 最后作者希望这个完全开源的 MoE 能够促进更多研究和分析从而提高对这些模型的理解。训练代码、中间检查点每 5000 step 、训练日志和训练数据都已经开源。
论文作者 Niklas Muennighoff 表示OLMoE 是第一个 100% 开源的混合专家 LLM。 预训练与自适应
预训练架构
OLMoE 是由 N_L 个 transformer 层组成的语言模型仅包含解码器。对于 OLMo 这样的密集模型原本模型中单一的前馈网络被 N_E 个小型前馈网络专家组成的混合专家网络所替代对于每个输入 token x只有 k 个专家会被选中并被激活负责处理这个输入。 其中路由器r是一个经过训练的线性层将输入的原始数据映射到被选中的 k 个专家上。对路由器的输出应用 softmax 函数计算 N_E 个专家的路由概率。然后每个被指定的专家 E_i 处理输入 x其输出乘以其各自的路由概率。再将所有选定的 Top-k 专家的结果相加构成模型单个层的 MoE 模块输出。
MoE 模型的训练往往涉及对一个已经存在的大型密集模型转换成一个稀疏模型也就是所谓的「稀疏升级」。这个过程中需要改变模型的训练目标比如调整 auxiliary load balancing 以及路由器的损失函数。具体的方法如下表所示 在这项研究中论文作者使用了总计 69 亿参数中的 13 亿活跃参数每层有 64 个专家其中有 8 个被激活。他们使用了一种名为「无丢弃 token」的路由方法对于每个输入 token路由器网络将分配 8 个专家来处理它。
论文作者引入了两个辅助损失函数负载平衡损失 L L B \mathcal{L}_{LB} LLB和路由器 z 损失 L R Z \mathcal{L}_{RZ} LRZ来训练 OLMoE-1B-7B。他们给这两个损失函数分别设定了权重α 和 β然后把它们和模型的主要学习目标交叉熵损失 L C E \mathcal{L}_{CE} LCE )结合起来最终计算的损失函数为 预训练数据
训练数据方面论文作者使用了来自两个不同来源的数据集DCLM 和 Dolma 1.7。这些数据集包括了多种类型的数据比如网络爬取的数据、编程问题解答、数学问题解答和学术论文等。他们将这些数据混合起来创建了一个名为 OLMOE-MIX 的新数据集。
下表中展示了预训练数据的组成 对于数据的处理论文作者使用了过滤器去除了包含太多重复 token 的内容、GitHub 上星标少于 2 的项目以及某些词出现频率过高的文档。他们将在每轮训练开始前随机混洗数据总计超过 5 万亿个 token。在「退火」阶段最后 100B 个 token他们首先重新混洗整个数据集然后按照此前 OLMo 论文中的方法将学习率线性衰减到 0。
自适应
论文作者从指令调优和偏好调优两方面基于之前的开放模型构造了 OLMoE-1B-7B-INSTRUCT。在指令调优集中他们增加了更多的代码和数学数据以提高模型在这些领域的性能。
GPT-4 和 Llama 3 在预训练阶段使用了像 GSM8k 或 MATH 这样的数学数据集的样本。按照这个思路论文作者还添加了「No Robots」和「Daring Anteater」的一个子集。这些数据集不仅质量高还更多样这是拓展模型适应性的两个关键因素。
下表展示了 OLMoE-1B-7B-INSTRUCT 所使用的数据 实验
该研究的评估程序由三部分组成预训练期间、预训练之后和自适应之后。
预训练期间如图 3 所示该研究在预训练期间使用当前最佳 OLMo 模型在常用下游任务上对 OLMoE-1B-7B 的性能进行了基准测试。 研究团队发现在所有任务中OLMoE-1B-7B 比密集 OLMo 模型以更少的计算量 (FLOP) 获得了更好的性能。尽管 OLMoE-1B-7B 使用了不到一半的 FLOP 进行训练并且仅使用 1B 个激活参数但 OLMoE-1B-7B 在训练结束时可与 OLMo-7B 媲美甚至优于 OLMo-7B。
预训练之后在表 4 中该研究在常见的下游任务上对 OLMoE-1B-7B 进行基准测试。
研究发现 OLMoE-1B-7B 在使用少于 2B 个激活参数的模型中表现最好使其成为许多 LM 用例中最经济的选择。
如果预算较大Qwen1.5-3B-14B 具有更强的性能但其激活参数和总参数比 OLMoE-1B-7B 多一倍以上。
研究发现尽管每条前向传播所需的计算量减少了约 6-7 倍但 OLMoE-1B-7B 的性能优于一些具有 7B 参数的密集 LM例如 Llama2-7B 但不如其他 LM例如 Llama3.1-8B 。上图 1 比较了 OLMoE-1B-7B 和其他 LM 的 MMLU 性能和激活参数表明 OLMoE-1B-7B 是其成本范围内最先进的。 自适应之后在表 5 中该研究对 OLMoE-1B-7B 的指令 (SFT) 和偏好 (DPO) 调优进行了基准测试。SFT 在所有测量任务上都改进了本文的模型。
DPO 在大多数任务上都有帮助尤其是 AlpacaEval这与先前研究的结果一致。DPO 模型称之为 OLMoE-1B-7B-INSTRUCT在所有基准测试模型中具有最高平均值。 强的离谱一份MoE的可视化指南
Grootendorst DASOU 2024年10月13日 17:59 北京 编辑AI椰青
整理https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-mixture-of-experts
在查看最新发布的 LLMs 时你可能会在标题中看到“MoE”这个词。那么这个“MoE”到底代表什么为什么现在有这么多 LLM 都在使用它呢
在本视觉指南中我们将通过 50 多个图示来详细探讨这一重要组件专家混合模型MoE 本指南将围绕 MoE 的两个主要组成部分——专家Experts和路由器Router)——在典型的基于 LLM 架构中的应用展开讨论。
目录
第一部分什么是专家混合模型
第二部分专家的角色
密集层Dense Layers稀疏层Sparse Layers专家能学到什么专家的架构
第三部分路由机制
路由器Router专家的选择路由的复杂性
第四部分负载平衡与优化
KeepTopK 策略Token 选择策略辅助损失函数专家容量使用 Switch Transformer 简化 MoE切换层Switching Layer容量因子Capacity Factor简化的辅助损失函数
第五部分视觉模型中的专家混合模型
Vision-MoE从稀疏 MoE 到软 MoE
第六部分带有 Mixtral 8x7B 的激活与稀疏参数
什么是专家混合模型
专家混合模型MoE是一种通过使用多个不同的子模型或“专家”来提升 LLM 质量的技术。
MoE 的两个主要组成部分为
专家Experts每个前馈神经网络层FFNN现在都有一组可以选择的“专家”。这些“专家”通常本身也是前馈神经网络FFNN。
路由器或门控网络Router 或 Gate Network决定哪些 token 被发送到哪些专家。
在具有 MoE 的 LLM 的每一层中我们可以发现一些相对专业化的专家 请注意这些“专家”并不像人类领域中的“心理学”或“生物学”专家那样在特定学科上表现出高度专业化。实际上它们更多是从词汇层
次上学习句法信息 更具体地说它们擅长于在特定上下文中处理特定的 tokens。
路由器或门控网络负责为每个输入选择最合适的专家 每个专家并不是一个完整的 LLM而是 LLM 架构中一个子模型部分。
专家
要理解专家的含义及其工作方式首先需要了解 MoE 用来替代的内容密集层Dense Layers。
密集层Dense Layers
专家混合模型MoE最初是从大型语言模型LLMs中相对基础的功能开始的即前馈神经网络Feedforward Neural Network, FFNN。
请记住在标准的Decode-only Transformer 架构中FFNN 通常应用于层归一化Layer Normalization之后 FFNN 允许模型利用注意力机制生成的上下文信息并进一步转化这些信息从而捕捉数据中更复杂的关系。
然而FFNN 的规模增长速度很快。为了学习这些复杂的关系它通常需要对接收到的输入进行扩展 稀疏层Sparse Layers
在传统的 Transformer 中FFNN 被称为密集模型Dense Model因为它的所有参数包括权重和偏置项都会被激活。所有参数都被用于计算输出没有任何部分被遗弃。
如果我们仔细观察密集模型可以发现输入在某种程度上激活了所有参数 相比之下稀疏模型Sparse Models仅激活总参数中的一部分这与专家混合模型MoE)密切相关。
为了说明这一点我们可以将密集模型分解为多个部分称为专家并重新训练它。随后在某一时刻只激活部分专家 其核心思想是每个专家在训练过程中学习不同的信息。而在推理时仅使用与当前任务最相关的特定专家。
当面对一个问题时我们可以选择最适合该任务的专家 专家学习了什么
正如我们之前所见专家学习的信息比整个领域的信息更加细粒度。因此有时将它们称为“专家”可能会产生误导。 然而在解码器模型中的专家似乎并没有表现出相同类型的专业化。这并不意味着所有专家都是等同的。
在 Mixtral 8x7B 论文中每个 token 都被其选择的第一个专家进行了着色这是一个很好的例子。 https://arxiv.org/pdf/2401.04088 这一可视化结果也表明专家往往更关注句法而不是某个特定领域的内容。
因此尽管解码器中的专家似乎没有特定的专业化但它们确实在某些类型的 tokens 上表现得更加一致。
专家的架构
虽然将专家可视化为密集模型的隐藏层被切分成若干部分很有帮助但实际上它们通常本身就是完整的 FFNN。 由于大多数 LLM 具有多个解码器块Decoder Blocks因此给定的文本会在生成之前通过多个专家 不同 tokens 被选中的专家可能各不相同这导致了不同的“路径”被选择 如果我们更新解码器块的可视化现在它将包含多个 FFNN每个 FFNN 对应一个“专家” 解码器块现在包含多个可以在推理时使用的 FFNN即“专家”)。
路由机制The Routing Mechanism
现在我们有了一组专家模型如何知道该使用哪些专家呢
在专家之前会加入一个路由器也称为门控网络它会被训练来选择每个 token 应选择的专家。
路由器The Router
路由器或门控网络本身也是一个 FFNN它根据特定的输入选择专家。路由器会输出概率值并利用这些概率来选择最匹配的专家 专家层返回被选定专家的输出并乘以门控值选择概率)。
路由器和专家其中仅选择少部分共同构成了 MoE 层 给定的 MoE 层有两种类型稀疏专家混合模型Sparse Mixture of Experts和密集专家混合模型Dense Mixture of Experts)。
两者都使用路由器来选择专家但稀疏 MoE 只选择少数几个专家而密集 MoE 则选择全部专家但可能会以不同的分布进行选择。 例如面对一组 tokensMoE 会将这些 tokens 分布到所有专家而稀疏 MoE 则只会选择少数几个专家。
在目前的 LLM 中当你看到“MoE”时它通常指的是稀疏 MoE因为稀疏 MoE 允许使用部分专家从而减少计算开销这对于 LLM 来说是一个重要的特性。
专家的选择
门控网络可以说是 MoE 中最重要的组件因为它不仅决定了推理时要选择哪些专家还决定了训练时的选择。
最基本的形式是我们将输入x与路由器的权重矩阵W相乘 然后我们对输出应用 SoftMax 操作为每个专家创建一个概率分布 G(x) 路由器利用这个概率分布来为给定的输入选择最匹配的专家。
最后我们将每个路由器的输出与各自选择的专家输出相乘并将结果相加 让我们将所有内容整合起来探索输入如何在路由器和专家中流动 路由的复杂性
然而这个简单的功能往往会导致路由器总是选择相同的专家因为某些专家可能比其他专家学习得更快 这不仅会导致专家选择的不均匀分布还会导致某些专家几乎没有被训练过。这会在训练和推理过程中引发问题。
因此我们希望在训练和推理期间各个专家的使用具有同等的重要性这就是所谓的负载平衡。某种程度上这是为了防止模型在同一组专家上过拟合。
负载平衡与优化
为了平衡专家的重要性我们需要重点关注路由器因为它是决定某一时刻选择哪些专家的主要组件。
KeepTopK
一种对路由器进行负载平衡的方法是使用一个简单的扩展策略称为 KeepTopK。通过引入可训练的高斯噪声我们可以防止总是选择相同的专家 https://arxiv.org/pdf/1701.06538 然后除希望激活的前 k 个专家例如 2 个以外的所有专家权重都将被设为 -∞ 将这些权重设为 -∞ 时SoftMax 操作后的输出概率将变为 0 尽管存在许多有前景的替代方案许多 LLM 仍然使用 KeepTopK 策略。需要注意的是KeepTopK 也可以不使用额外的噪声。
Token 选择策略
KeepTopK 策略会将每个 token 路由到若干选定的专家。这种方法被称为 Token 选择策略Token Choice它允许一个给定的 token 被路由到一个专家top-1 路由 或者被路由到多个专家top-k 路由 这种策略的主要优点在于它可以对各个专家的贡献进行加权并将其整合起来。
辅助损失Auxiliary Loss
为了在训练期间实现专家的均匀分布网络的常规损失中加入了辅助损失也称为负载平衡损失。
辅助损失增加了一个约束强制专家在训练过程中具有相同的重要性。
辅助损失的第一个组成部分是对整个批次中每个专家的路由值进行求和 这为我们提供了每个专家的重要性得分即在不考虑输入的情况下给定专家被选中的概率。
我们可以使用这些重要性得分计算变异系数Coefficient of Variation, CV它表示各个专家的重要性得分之间的差异程度。 例如如果重要性得分之间的差异较大那么 CV 值就会较高 相反如果所有专家的得分都相似则 CV 值较低这是我们期望的情况 通过使用这个 CV 得分我们可以在训练过程中更新辅助损失使其尽可能降低 CV 得分从而使每个专家具有相同的重要性 最后辅助损失将作为一个独立的损失项参与训练优化。
专家容量Expert Capacity
专家的不平衡不仅体现在被选中的专家上还体现在分配给这些专家的 token 分布上。
例如如果输入 token 被不成比例地分配到某些专家上而不是平均分配这可能导致某些专家的训练不足 这里不仅要考虑使用了哪些专家还需要关注这些专家被使用的频率。
解决这个问题的方法是限制每个专家能够处理的 token 数量即专家容量Expert Capacity。当一个专家达到其容量时多余的 token 将被分配到下一个专家 如果两个专家都达到了其容量token 将不会被任何专家处理而是直接传递到下一层。这种情况被称为 token 溢出token
overflow。 使用 Switch Transformer 简化 MoE
第一个解决 MoE 训练不稳定性问题如负载平衡的基于 Transformer 的 MoE 模型是 Switch Transformer。它通过简化架构和训练过程提高了训练稳定性。
切换层Switching Layer
Switch Transformer 是一个 T5 模型编码器-解码器结构它将传统的 FFNN 层替换为切换层Switching Layer。切换层是一个稀疏的 MoE 层它为每个 token 选择单个专家top-1 路由。 路由器在选择专家时并没有使用特殊的方法只是对输入与专家权重相乘的结果取 softmax与之前的方法相同。 这种架构top-1 路由假设每个 token 只需要一个专家来学习如何进行路由。这与我们之前讨论的 top-k 路由将 token 分配给多个专家)有所不同。
容量因子Capacity Factor
容量因子是一个重要参数它决定了每个专家可以处理的 token 数量。Switch Transformer 通过引入直接影响专家容量的容量因子扩展了这一概念。 专家容量的组成部分非常简单 如果我们增加容量因子则每个专家能够处理更多的 token。 然而如果容量因子过大就会浪费计算资源。相反如果容量因子过小模型性能会因为 token 溢出而下降。
辅助损失Auxiliary Loss
为了进一步防止 token 被丢弃Switch Transformer 引入了简化版的辅助损失。
在简化版的辅助损失中不再计算变异系数而是将分配的 token 数量与每个专家的路由概率进行加权比较 由于目标是希望将 token 在 N 个专家中均匀分配因此我们希望向量 P 和 f 的值为 1/N。
α 是一个超参数用于在训练过程中微调此损失的重要性。值过高会影响主要损失函数而值过低则无法有效进行负载平衡。
视觉模型中的专家混合模型
MoE 并不仅限于语言模型。视觉模型如 ViT使用 Transformer 架构因此也有潜力使用 MoE。
快速回顾一下ViTVision Transformer是一种将图像分割为若干块并将其作为 tokens 处理的架构。 这些图像块或 tokens会被投射到嵌入向量加上额外的位置嵌入向量中然后输入到常规编码器中 当这些图像块进入编码器时它们会像 tokens 一样被处理这使得这种架构非常适合 MoE。
Vision-MoE
Vision-MoEV-MoE是图像模型中第一个实现 MoE 的例子之一。它将 ViT 中的密集 FFNN 层替换为稀疏 MoE。 这种改进使得 ViT 模型通常比语言模型小)能够通过增加专家的数量来大幅扩展。
为了降低硬件限制每个专家都设置了一个较小的预定义容量因为图像通常包含大量的图像块。然而低容量往往会导致图像块被丢弃类似于 token 溢出。 为了保持容量较低网络会为每个图像块分配重要性得分并优先处理这些得分较高的图像块从而避免溢出图像块的丢失。这种方法被称为批量优先路由Batch Priority Routing。 因此即使 token 数量减少我们仍然能够看到重要的图像块被成功路由。 优先路由使得在处理较少的图像块时仍能聚焦于最重要的图像块。
从稀疏 MoE 到软 MoE
在 V-MoE 中优先评分机制能够区分出重要和不重要的图像块。然而图像块被分配给每个专家后未被处理的图像块中的信息就会丢失。
软 MoESoft-MoE旨在通过混合图像块从离散的图像块token分配转变为软分配。
第一步我们将输入 x图像块嵌入与一个可学习矩阵 Φ 相乘。这将生成路由信息它告诉我们某个 token 与某个专家的相关程度。 然后对路由信息矩阵进行 softmax 操作在列上从而更新每个图像块的嵌入向量。 更新后的图像块嵌入本质上是所有图像块嵌入的加权平均。 从视觉上看这就像是所有图像块被混合。这些组合后的图像块被发送到每个专家。生成输出后它们再次与路由矩阵相乘。 路由矩阵在 token 层面影响输入并在专家层面影响输出。
因此我们获得了“软”图像块/token这些 token 被处理而不是离散输入。
Mixtral 8x7B 的激活与稀疏参数对比
MoE 的一个重要特点是其计算需求。由于在同一时刻只会使用部分专家我们可以拥有比实际使用的更多的参数。
尽管给定的 MoE 拥有更多的参数稀疏参数但由于我们在推理时只使用部分专家活跃参数因此激活的参数较少。 换句话说我们仍然需要将整个模型包括所有专家加载到设备中稀疏参数但在实际运行推理时我们只需要使用部分参数活跃参数)。MoE 模型需要更多的显存来加载所有专家但推理时运行速度更快。
让我们以 Mixtral 8x7B 为例来探讨稀疏参数与活跃参数的数量差异。 在此例中我们可以看到每个专家的参数量为 5.6B而不是 7B尽管一共有 8 个专家。 我们需要加载 8x5.6B46.7B的参数加上所有共享参数但推理时只需要使用 2x5.6B12.8B)的参数。
结论
以上就是我们对专家混合模型MoE的探索之旅希望这篇文章能帮助你更好地理解这一有趣技术的潜力。如今几乎所有的模型架构中都有 MoE 变体这也预示着它可能会长期存在下去。 大模型混合专家模型MoE概述
关注大模型的 AI大模型前沿 2024年12月13日 10:06 福建
随着 GPT-4、DeepSeekMoE 等模型的发布中均涉及到了混合专家模型MoEMixture of Experts的话题MoE 模型已经成为开放 AI 社区的热门话题。2023年6月美国知名骇客George Hotz在接受采访时透露GPT-4 由 8 个 220B 的专家模型组成。假如把 8 个专家模型比喻为比GPT-3还大的脑袋那GPT-4就是一个八个头的超级大怪兽。 GPT-4MoE比GPT-3Transformer和GPT-3.5RLHF强大一个数量级的关键可能就是来源于MoE架构。之前的GPT大模型增大参数的方法是在一个GPT模型上堆层数现在变成了堆模型数。将来大语言模型的研究新方向可能就不是增大单一模型的向量维度和层数了而是增大整体架构的模型数了。GPT-4引入MoE似乎是个必然因为无论是算力、数据、稳定性万亿级参数的单个大模型训练很困难而且推理成本也会居高不下跑万亿个参数的计算才能算出一个token的速度和成本比较不可观。所以将若干个模型堆成一个MoE大模型似乎是个必然趋势。
那么究竟什么是MoE大模型MoE大模型具备哪些优势呢
MoE全称为Mixed Expert Models混合专家模型简单理解就是将多个专家模型混合起来形成一个新的模型。在理解MOE之前有两个思想前提可以帮助我们更容易地理解MOE架构。一是在现实生活中如果有一个包括了多个领域知识的复杂问题我们该使用什么样的方法来解决呢最简单的办法就是先拆分任务到各领域然后把各个领域的专家集合到一起来攻克这个任务最后再汇总结论。这个思想可以追溯到集成学习MoE和集成学习的思想异曲同工都是集成了多个模型的方法区别在于集成学习不需要将任务分解为子任务。集成学习是通过训练多个基学习器来解决同一问题并且将它们的预测结果简单组合例如投票或平均。而MOE是把大问题先做拆分再逐个解决小问题再汇总结论。二是模型规模是提升模型性能的关键因素之一。在有限的计算资源下用更少的训练步数训练一个更大的模型往往比用更多的步数训练一个较小的模型效果更佳。 MoE正是基于上述的理念它由多个专业化的子模型即“专家”组合而成每一个“专家”都有其擅长的领域。而决定哪个“专家”参与解答特定问题的是一个称为“门控网络”的机制。技术上常说的门控机制可能会先想到LSTM的门控机制但是这里的门控机制和LSTM里的门控不一样。LSTM的门是为了控制信息流动这里的门就更像我们日常中提到的门选择进门或是不进门是一个控制是否使用某个专家模型的概率分布值。 MoE基于Transformer架构主要由两部分组成
稀疏 MoE 层MoE层代替了传统 Transformer 模型中的前馈网络 (FFN) 层。MoE 层包含若干“专家”模型每个专家本身是一个独立的神经网络。在实际应用中这些专家通常是前馈网络 (FFN)但它们也可以是更复杂的网络结构。
门控网络或路由: 这个部分用于决定哪些 token 被发送到哪个专家。例如在上图中“More”这个 token 可能被发送到第二个专家而“Parameters”这个 token 被发送到第一个专家。同时一个 token 也可以被发送到多个专家。token 的路由方式是 MoE 使用中的一个关键点因为路由器由学习的参数组成并且与网络的其他部分一同进行预训练。
MoE 的一个显著优势是它们能够在远少于 Dense 模型所需的计算资源下进行有效的预训练。这意味着在相同的计算预算条件下您可以显著扩大模型或数据集的规模。特别是在预训练阶段与稠密模型相比混合专家模型通常能够更快地达到相同的质量水平。例如Google的Switch Transformer模型大小是T5-XXL的15倍在相同计算资源下Switch Transformer模型在达到固定困惑度 PPL 时比T5-XXL模型快4倍。
国内的团队DeepSeek 开源了国内首个 MoE 大模型 DeepSeekMoE。
DeepSeekMoE 2B可接近2B Dense仅用了17.5%计算量。
DeepSeekMoE 16B性能比肩 LLaMA2 7B 的同时仅用了40%计算量。
DeepSeekMoE 145B 优于Google 的MoE大模型GShard而且仅用 28.5%计算量即可匹配 67B Dense 模型的性能。
此外MoE大模型的优点还有
训练速度更快效果更好。相同参数推理成本低。扩展能力强允许模型在保持计算成本不变的情况下增加参数数量这使得它能够扩展到非常大的模型规模如万亿参数模型。多任务学习能力MoE在多任务学习中具备很好的性能。
MoE结合大模型属于老树发新芽MOE大模型的崛起是因为大模型的发展已经到了一个瓶颈期包括大模型的“幻觉”问题、逻辑理解能力、数学推理能力等想要解决这些问题就不得不继续增加模型的复杂度。随着应用场景的复杂化和细分化垂直领域应用更加碎片化想要一个模型既能回答通识问题又能解决专业领域问题尤其在多模态大模型的发展浪潮之下每个数据集可能完全不同有来自文本的数据、图像的数据、语音的数据等数据特征可能非常不同MoE是一种性价比更高的选择。国内大模型已经开始朝着MoE方向大步前进在2024年估计会有越来越多大模型选择MoE架构。
后文将介绍MoE的主要原理以此来理解MoE大模型优势产生的原因。
一、Adaptive mixtures of local experts
Adaptive mixtures of local experts这是大多数MoE论文都引用的最早的一篇文章发表于1991年。论文介绍了一种新的监督学习过程——由多个独立网络组成一个系统每个网络单独处理训练集合的子集。因为对于多任务学习如果是常见的多层网络学习各层之间的网络通常会有强烈的干扰效应这会导致学习过程变慢和泛化能力差。为了解决这个问题论文提出了使用多个模型即专家expert去学习使用一个门控网络gating network来决定每个数据应该被哪个模型去训练的权重这样就可以减轻不同类型样本之间的干扰。
对于一个样本 c第 i 个 expert 的输出为 o i c \mathbf{o}_{i}^{c} oic期望的输出向量为 dc那么损失函数为 E c ∥ d c − ∑ i p i c o i c ∥ 2 E^c \left\|\mathbf{d}^c - \sum_{i} p_{i}^{c} \mathbf{o}_{i}^{c}\right\|^2 Ec∥dc−∑ipicoic∥2
其中 p i c p_{i}^{c} pic 是门控网络分配给第 i 个expert的权重。
但是作者提到这个损失函数可能会导致专家网络之间的强烈耦合因为是所有专家网络的权重加总来共同计算损失的一个专家权重的变化会影响到其他专家网络的loss。这种耦合可能会导致多个专家网络被用于处理每条样本而不是专注于它们各自擅长的子任务。为了解决这个问题论文重新定义了损失函数以鼓励专家网络之间的相互竞争。 就是先让不同的expert单独计算loss然后再加权求和得到总体的loss。这意味着每个expert在处理特定样本的目标是独立于其他expert的。如果门控网络和expert都使用这个新的loss进行梯度下降训练系统会倾向于将某类特定样本分配给特定的expert。因为当一个expert在给定样本上的的loss小于所有expert的平均loss时它对该样本的门控score会增加当它的表现不如平均loss时它的门控score会减少。这种机制鼓励expert之间的竞争而不是合作从而提高了学习效率和泛化能力。
二、Sparsely-Gated MoE
在 2010 至 2015 年间两个独立的研究领域为混合专家模型 (MoE) 的后续发展做出了显著贡献
组件专家Eigen、Ranzato 和 Ilya 探索了将 MoE 作为更深层网络的一个组件。这种方法允许将 MoE 嵌入到多层网络中的某一层使得模型既大又高效。
条件计算Conditional Computation传统的神经网络通过每一层处理所有输入数据。Yoshua Bengio 等研究人员开始探索基于输入 token 动态激活或停用网络组件的方法。 2017 年Shazeer 等人将上述概念应用于 137B 的 LSTM 。通过引入稀疏性这项工作在保持极高规模的同时实现了快速的推理速度。在牺牲极少的计算效率的情况下把模型规模提升1000多倍。这项工作被发表在论文Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer中和 1991 年Adaptive mixtures of local experts的工作对比这里的 Sparsely-Gated MoE 主要有两个区别
Sparsely-Gated不是所有expert都会起作用而是极少数的expert会被使用来进行推理。这种稀疏性也使得我们可以使用海量的experts来把模型多倍扩大。token-level相比于sample-level即不同的样本使用不同的专家采用token-level即一个句子中不同的token使用不同的专家。
如上图所示每个token都会接一层MoE Layer每个MoE layer中包含了多个experts还有一个Gating Network会根据当前 token选择少数几个expert来进行计算。
2.1 门控网络Gating Network
门控网络Gating Network的设计和实现是Sparsely-Gated MoE 层的核心组成部分。门控网络负责为每个输入 token 选择一个稀疏的专家组合。一般门控网络最简单的就是概率分布形式根据概率去做出相应的选择。那么在深度学习中最简单的网络就是设定门控网络为一个简单的前馈神经网络最后接一层softmax作为每个专家的使用权重。
设 G(x) 和 Ei(x) 分别是门控网络和第 i 个 expert 的输出那么对于在当前的输入x输出就是所有 experts 的加权和 那么一个典型的门控网络就是一个带有 softmax 函数的简单的网络。这个网络将学习将输入发送给哪个expert 在这种设置下所有expert都会对输入进行运算再通过门控网络的输出进行加权求和如果experts的数量太大就会导致计算量非常大。但是如果有一种方法能使某些专家模型的门控网络的输出为0就没有必要对这个专家进行相应的计算就可以节省计算资源。其中效果比较好的包括带噪声的 TopK 门控 (Noisy Top-K Gating)。这种门控方法引入了一些可调整的噪声然后保留前 k 个值。具体来说
添加一些噪声 选择保留前 k 个值 应用 Softmax 函数 对于非TopK的部分由于值是负无穷这样在经过Softmax之后就会变成 0不会被选中。噪声可以使得不同expert的负载更加均衡。在具体实验中作者使用的 K2~4。
在MoE模型中加入噪声的原因主要有以下几点
提高模型的鲁棒性和泛化能力。当模型在训练或推理阶段遇到不确定或嘈杂的数据时鲁棒性较强的模型更能保持稳定的性能。减少过拟合的风险。负载均衡在MoE模型中通过加入噪声可以实现不同专家之间的负载均衡。
2.2 辅助损失——均衡专家利用负载率Balancing Expert Utilization
要理解专家负载均衡可以从batch size的角度出发。通常来讲较大的 batch size 推理性能更好但是由于样本在MoE层激活专家时需要并行MoE层中每个专家的实际batch size会减少。举个例子假设当前 batch 有10个token其中5个token被路由到了某个专家网络而其他5个token被路由到了其它5个不同的专家网络这就会导致各专家网络获得的 batch size 不均匀一个是另外的5倍就会导致算力利用率不足的问题。
并且如果我们把所有的token都发送给少数几个头部专家训练效率将会变得低下。因为在训练中门控网络会收敛至仅选择少数几个头部专家。这种情况会不断自我强化因为受青睐的专家训练得更快它的效果一直在被优化因此选择它们的频率也会更高。那么此时这几个少数专家就会过载每次需要计算大量token其他的专家模型就得不到训练而被闲置就会导致模型失衡性能下降。因此MoE模型中需要控制专家均衡。上文加入噪声机制可能就会导致门控在做出选择时跳过几个最优专家模型去选择其余的专家模型从而更好地协同工作。
此外为了缓解这种情况可以添加辅助损失鼓励给予所有专家同等的重视。论文提出了一种软约束方法。作者定义了专家相对于一批训练样本的重要性 Importance(X) 即该专家在这批样本中门控值的总和。然后定义了一个额外的损失函数 LImportance(X) 这个损失函数被添加到模型的整体损失函数中。这个损失函数等于重要性值集合的CVcoefficient of variation平方变异系数可以度量一组数据的离散程度乘以一个手动调整的缩放因子 wImportance 。这个额外的损失鼓励所有专家具有相等的重要性具体计算公式如下所示 这种损失确保所有专家能收到数量大致相等的训练样本。
三、GShard
GShard是谷歌2021年在论文GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding中提出的使用GShard实现MoE跨设备分片的方法是第一个将MoE的思想拓展到Transformer上的工作。它具体的做法是把Transformer的encoder和decoder中每隔一个的FFN层替换成使用Top-2门控的MoE层。 上图展示了编码器部分的结构。这种架构对于大规模计算非常有效当扩展到多个设备时MoE 层在不同设备间共享而其他所有层则在每个设备上复制。GShard MoE层中的专家网络experts被分布在不同的设备上。每个专家网络负责处理一部分输入数据并且每个token根据门控机制的输出被分配到一个或两个专家网络中。这样整个 MoE 层的计算被分散到了多个设备上每个设备负责处理一部分计算任务。
实现 MoE 跨设备分片的关键技术是模型并行化model parallelism和数据并行化data parallelism的结合。在模型并行化中模型的不同部分这里指MoE层的专家网络被分配到不同的设备上。在数据并行化中输入数据被分割成多个部分每个部分被分配给不同的设备进行处理。由于专家被分配到不同设备可以并行计算大大提升了模型的计算效率这也解释了为什么 MoE 可以实现更大模型参数、更低训练成本。
为了保持负载平衡和训练效率GShard 的作者除了引入上节 Sparsely-Gated MoE 中的辅助 loss 外还引入了一些关键变化
随机路由在 Top-2 设置中GShard 始终选择排名最高的专家但第二个专家是根据其权重比例随机选择的。专家容量设定一个阈值定义一个专家最多能处理多少 token。如果专家容量达到上限token 就会溢出并通过残差连接传递到下一层或被完全丢弃。专家容量是 MoE 中最重要的概念之一。为什么需要专家容量呢因为我们无法提前知道多少 token 会分配给每个专家因此需要一个固定的容量因子来防止专家过载。其计算公式为 专家容量希望将 batch 中的总 token 数平均分配给所有专家。然后为了应对 token 特征分布不均的情况会通过一个容量因子来扩展每个专家的容量。容量因子是一个大于 1.0 的数它的作用是为每个专家提供额外的缓冲空间以容纳可能超出平均分配的 token。这样在数据分布不均的情况下即使某些专家接收到的 token 数量超过了平均值也能够处理这些额外的 token而不会因为容量不足而导致计算跳过。但是如果容量因子设置得过高会导致计算资源和内存的浪费因为模型可能永远不会用到这额外的资源。经过研究容量因子 在1至1.25下表现出色。
注意在推理过程中有些计算过程是共享的例如自注意力 (self-attention) 机制它适用于所有 token。这就解释了为什么Mixtrial 8×7B不是56B而是47B。同理如果采用 Top-2 门控模型会使用 14B 的参数。但是由于自注意力操作 (专家间共享) 的存在实际上模型运行时使用的参数数量是 12B。
四、Switch Transformers
2022年Google提出Switch Transformers对MoE大模型的复杂性、通信成本以及训练微调过程的不稳定性进行了优化。Switch Transformers是一个1.6万亿参数的MoE拥有2048个专家可以使用transformers库运行。
Switch Transformers 简化了 MoE 路由算法设计了直观的改进模型降低了通信和计算成本。Switch Transformers 的训练方法减轻了不稳定性并且首次展示了用较低精度bfloat16格式训练大型稀疏模型的可能性。 上图模型参数随着专家数量的增加而增加训练时保持了相同的计算成本loss逐渐降低。这表明模型在保持计算效率的同时能够利用更多的参数来提高性能。 上图比较了使用相同计算资源下Switch Transformer 和 T5-Base 的 PPL。可以看到 Switch Transformer 模型在保持相同计算资源的情况下相对于 T5-Base 有显著的提升而且专家数越多模型参数越多、模型更稀疏效果越好。
4.1 Switch Transformer主要优化
Swith Transformer 在论文中提到其设计指导原则是——尽可能地把 Transformer 模型的参数量做大即加专家数量堆叠专家模型
和其他 MoE 模型的一个显著不同就是Switch Transformer 的门控网络每次只路由到 1 个 expert也就是每次只选取 Top1 的专家而其他的模型都是至少 2 个。这样从 MoE layer 的计算效率上讲是最高的同时也降低了 MoE 中的通信成本
4.2 改进预训练和Fine-Tuning技术
1. 改进预训练——精度选择
作者尝试了混合精度改进预训练较低的精度可以减少处理器间的通信成本、计算成本以及存储 tensor 的内存。然而在最初的实验中当专家和门控网络都使用 bfloat16 精度训练时出现了不稳定的训练现象。这种不稳定性主要是由门控计算引起的因为门控涉及指数函数等操作这些操作对精度要求较高。因此为了保持计算的稳定性和精确性保持更高的精度是重要的。为了减轻不稳定性门控过程使用了全精度。
下表显示了混合精度训练的效果将路由器输入转换为 float32同时保持其他部分的精度为 bfloat16。这种策略允许模型在几乎与 bfloat16 精度相同的训练速度下实现与 float32 训练相当的稳定性。 2. 改进预训练——更小的参数初始化
在深度学习中适当的权重初始化对于模型的成功训练至关重要。作者观察到在 Switch Transformer 模型中这一点尤其明显。
为了提高模型的稳定性作者建议减少默认的 Transformer 初始化规模。在 Transformer 模型中权重矩阵通常是从一个正态分布开始。作者建议将这个初始化正态分布标准差的超参数 s 从默认值1.0 减少 10 倍到 s 0.1模型效果和稳定性得到了提升。 3. 改进微调——Fine-Tuning 过程正则化
为了解决 Fine-Tuning 过程中的过拟合问题作者提出了增加 dropout的策略特别是在专家层expert layers中称之为“expert dropout”即在 Fine-Tuning 时只在专家层增加 dropout 率。 通过这种 expert dropout 策略有效地减少了过拟合的风险同时保持了模型在下游任务上的性能。这种正则化方法对于处理具有大量参数的稀疏模型特别有用因为它可以帮助模型更好地泛化到未见过的数据。
五、ST-MOE——用 Router z-loss 稳定模型训练
在论文 ST-MOE: Designing Stable and Transferable Sparse Expert Models 中作者提出了一种新的辅助损失函数称为 Router z-loss用于提高稀疏模型的训练稳定性同时保持或稍微提高模型质量。这个损失函数是针对稀疏专家模型中的路由器router部分设计的路由器负责将输入的 token 路由到最合适的专家expert层。这个损失函数的目标是鼓励路由器产生较小的logits 值因为较大的 logits 值在 softmax 激活函数中会导致较大的梯度这可能会引起训练不稳定。
Router z-loss 的定义如下 其中 B 是 batch 中的 token 数量 N 是专家的数量 x是门控的 logits值。这个损失函数通过惩罚较大的 logits 值来工作因为这些值在 softmax 函数中会导致较大的梯度。通过这种方式Router z-loss 有助于减少训练过程中的不稳定性。
稳定梯度更新在训练深度神经网络时梯度的更新可能会因为logits值过大而导致不稳定造成梯度消失或是爆炸激活函数的使用问题涉及指数计算导致梯度消失。适当减小logits的幅度可以使梯度更新更加平稳从而提高训练的稳定性。提高泛化能力过大的logits值可能会导致模型对训练集中的某些特征过度敏感通过对logits进行适当的惩罚可以使模型更加谨慎地做出预测避免对训练数据过拟合。
此外ST-MOE 的作者还提到了两个MoE模型的现象
专家如何学习
encoder层的专家倾向于专注于特定类型的 token 或浅层概念。例如某些专家可能专门处理标点符号而其他专家则专注于专有名词等。decoder 中的专家通常具有较低的专业化程度。此外研究者们还对这一模型进行了多语言训练发现模型并不会按照人们期望的结构进行训练尽管人们可能会预期每个专家处理一种特定语言但实际上并非如此。由于 token 路由和负载均衡的机制没有任何专家被特定配置以专门处理某一特定语言。
专家的数量对模型有何影响
增加更多专家可以提升效率但这些优势随着专家数量的增加而递减 (尤其是当专家数量达到 256 或 512 之后更为明显)比较符合常见的边际效用递减现象。并且这种特性在小规模模型下也同样适用即便是每层仅包含 2、4 或 8 个专家也会存在这种情况。所以在设计MoEt模型时要考虑这个算力和模型质量的均衡专家数量不是越多越好。
六、微调MoE模型
稠密模型和稀疏模型在过拟合的动态表现上存在显著差异。稀疏模型更易于出现过拟合现象因此在处理这些模型时尝试更强的内部正则化措施是有益的比如使用更高比例的 dropout。例如我们可以为稠密层设定一个较低的 dropout 率而为稀疏层设置一个更高的 dropout 率以此来优化模型性能。
另一种可行的 Fine-Tuning 策略是尝试冻结所有非专家层的权重。然而实验发现这会导致性能大幅下降我们可以尝试相反的方法仅冻结 MoE 层的参数。实验结果显示这种方法几乎与更新所有参数的效果相当同时可以加速 Fine-Tuning 过程并降低显存需求。
在Fine-Tuning时还需要考虑的一个问题是它们有需要特殊设置的超参数例如稀疏模型往往更适合使用较小的batch size和较高的学习率这样可以获得更好的训练效果。 论文 MoEs Meets Instruction Tuning2023 年 7 月还做了以下几个实验
非MoE模型微调MoE模型微调非MoE模型指令微调MoE模型指令微调
结论MoE模型指令微调 非MoE模型指令微调 非MoE模型微调 MoE模型微调并且子任务类型越多MoE模型指令微调的效果越好
七、MoE 模型后续研究方向
MoE模型现在的几个值得探索的方向
1、将 MoE蒸馏成稠密模型。Switch Transformers 的作者做了一些初步的蒸馏实验。通过将 MoE 蒸馏成一个稠密模型能够保留 30-40% 的稀疏性增益。
2、探索专家聚合技术。该技术对专家的权重进行了合并从而减少了参数量。
3、对MoE执行极致量化。QMoE2023年10月通过将MoE量化至每参数不足 1 比特从而将1.6T的Switch Transformer模型的内存使用量从3.2TB压缩至仅需160GB。
八、开源 MoE 模型
国内的MoE大模型DeepSeek团队开源的DeepSeekMoE模型、代码、论文均已同步发布。
模型下载https://huggingface.co/deepseek-ai微调代码https://github.com/deepseek-ai/DeepSeek-MoE技术报告https://github.com/deepseek-ai/DeepSeek-MoE/blob/main/DeepSeekMoE.pdf
此外还有一些国外的开源 MoE 模型开源了训练代码。
Megablockshttps://github.com/stanford-futuredata/megablocksFairseqhttps://github.com/facebookresearch/fairseq/tree/main/examples/moe_lmditaOpenMoEhttps://github.com/XueFuzhao/OpenMoE
下面是开源了模型但是没有开源代码
Switch Transformers基于T5的MoE专家数量共2048最大模型有1.6万亿个参数。NLLB MoE (Meta)NLLB 翻译模型的一个 MoE 变体。OpenMoE社区对基于 Llama 的模型的 MoE 尝试。Mixtral 8x7B一个性能超越了 Llama 2 70B 的高质量 MoE并且具有更快的推理速度。此外还发布了一个指令微调的版本。 via: 手把手教你从零开始实现一个稀疏混合专家架构语言模型MoE https://mp.weixin.qq.com/s/fV-dnbT5g8j_B3RXgRsuZQ 第一个100%开源的MoE大模型7B的参数1B的推理成本 https://mp.weixin.qq.com/s/FvsYm5HxH4f9Km4Aqrso_Q 强的离谱一份MoE的可视化指南 https://mp.weixin.qq.com/s/d2DNDH2lHaaiw9RcNj0iiw 大模型混合专家模型MoE概述 https://mp.weixin.qq.com/s/X_EUAAjaKMpwQeatwW6Z5A
