ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

从MMoE到HoME:多任务学习中的专家分层设计思路

从MMoE到HoME:多任务学习中的专家分层设计思路 1. 从MMoE到HoME多任务学习中的专家分层设计思路多任务学习Multi-Task Learning, MTL在推荐系统、广告预估、搜索排序这些场景里已经不算新鲜事了。只要你的业务里存在多个目标——比如同时预测点击率、转化率、停留时长——就一定会碰到一个核心矛盾不同任务之间既有可以共享的通用信息又有各自独有的领域特征。共享多了任务之间互相拖累这叫负迁移共享少了每个任务各自为战数据稀疏的问题又解决不了。MMoEMulti-gate Mixture-of-Experts是谷歌在2018年前后提出的一个经典解法。它的思路很直接底层放一组专家网络Expert每个任务配一个独立的门控网络Gate门控负责给各个专家分配权重最终每个任务拿到的是所有专家输出的加权组合。这样一来任务之间共享的是专家池但每个任务可以按自己的需要去“挑选”专家灵活性比硬共享Hard Parameter Sharing高出一大截。但MMoE用久了问题也会暴露出来。我最早在推荐场景里用MMoE的时候就发现一个现象当任务数量增加到五六个以上或者任务之间的相关性差异很大时门控网络会逐渐“偷懒”。具体表现是几个门控的权重分布越来越像最后退化成近似平均加权专家池的差异化能力被浪费掉了。更麻烦的是所有任务都在同一层级上竞争同一批专家强任务会主导专家的训练方向弱任务只能捡残羹剩饭。HoMEHierarchy of Multi-Gate Experts要解决的就是这个问题。它的核心思想可以用一句话概括把专家池从“平铺”变成“分层”让不同任务在不同粒度上提取共享信息。底层专家负责捕捉所有任务通用的基础特征上层专家负责捕捉任务分组级别的中间特征顶层再留给每个任务自己的专属专家。门控网络也不再是每个任务一个而是按层级组织形成一棵门控树。这个设计的好处在于它把“共享”和“专属”从一刀切变成了渐进式的。底层强制共享保证基础表征的泛化能力中层按任务簇共享让相似任务互相借力顶层完全独立保留每个任务的个性。从信息论的角度看这相当于对任务间的互信息做了分层建模——底层对应高互信息部分顶层对应低互信息部分。注意HoME并不是要完全取代MMoE而是在MMoE基础上做结构化扩展。如果你的任务数量少2-3个且相关性高MMoE已经够用任务数量多、相关性差异大时HoME的优势才会明显体现。1.1 为什么是“分层”而不是“加专家”有人可能会问既然MMoE的问题是专家不够用那我直接加专家数量不就行了这个思路我试过效果有限。原因有两个。第一专家数量增加会带来门控网络的维度灾难。假设你有N个专家、T个任务门控网络的参数量是N×T。当N从8增加到32T6时门控参数量从48涨到192训练难度显著上升。而且门控的softmax输出会变得更加稀疏大部分专家的梯度信号很弱训练不充分。第二平铺的专家池没有结构先验。所有专家在初始化时都是一样的门控网络需要从零开始学习“哪个专家适合哪个任务”。这个过程在数据量不足时非常低效。而分层结构相当于给模型注入了一个先验底层专家天然适合所有任务中层专家天然适合任务簇顶层专家天然适合单任务。这个先验来自我们对任务关系的理解比让模型盲目搜索要靠谱得多。HoME的分层设计还有一个隐含好处它天然支持任务簇的发现。在实际业务中我们往往并不清楚哪些任务应该归为一簇。HoME的中层门控可以通过学习到的权重分布反过来告诉我们任务之间的相似度。我做过一个实验把六个任务的中层门控权重做聚类结果和业务上的人工分组高度一致这算是一个意外的可解释性收获。1.2 HoME的整体架构拆解HoME的架构可以分成三个部分专家分层、门控分层、输出融合。专家分层假设我们有两层专家。底层专家数量为E1每个专家是一个全连接网络输入是所有任务的共享特征。底层专家的输出会同时送给中层专家和顶层门控。中层专家数量为E2每个中层专家的输入是底层所有专家的输出拼接或者加权和输出再送给顶层门控。顶层专家数量为E3每个顶层专家只服务于一个任务输入是底层和中层专家的输出。门控分层底层门控是任务级别的每个任务一个门控负责从底层专家中挑选组合。中层门控是任务簇级别的每个任务簇一个门控负责从中层专家中挑选组合。顶层门控是任务级别的负责从顶层专家中挑选组合。最终每个任务的输出是三层专家输出的加权和。输出融合每个任务的最终输出 底层门控加权和 中层门控加权和 顶层门控加权和。注意这里的“”不是简单相加而是拼接后过一个任务专属的MLP让模型自己学习三层信息的融合方式。这个架构的关键参数是E1、E2、E3的比例。我的经验是E1可以设大一点比如16因为底层需要捕捉多样化的基础特征E2适中比如8对应任务簇的数量E3等于任务数每个任务一个专属专家。当然这只是一个起点具体还要看数据量和任务复杂度。2. 核心细节解析门控网络、专家设计与梯度流动HoME的论文里给了架构图但真正落地的时候魔鬼都在细节里。这一章我把几个关键模块的设计要点拆开讲包括门控网络的实现、专家网络的选型、以及梯度流动的控制。2.1 门控网络的实现细节与温度系数门控网络本质上是一个softmax层输入是任务的特征向量输出是专家权重分布。在MMoE里门控就是一个简单的线性层加softmax。但在HoME里门控需要处理分层结构实现上要稍微讲究一点。底层门控的输入是任务的原始特征比如用户ID embedding、上下文特征输出维度是E1。中层门控的输入是底层专家的输出拼接输出维度是E2。顶层门控的输入是底层和中层专家的输出拼接输出维度是E3。这里有一个容易踩坑的地方softmax的温度系数。默认温度是1.0但在多任务场景下我建议把温度设成可学习的参数初始值设为1.0。原因在于不同任务的门控需要不同的“锐度”。强任务可能需要更集中的权重低温度弱任务可能需要更分散的权重高温度。让模型自己学温度比手动调要省事得多。代码实现上门控网络可以这样写class Gate(nn.Module): def __init__(self, input_dim, num_experts, temperature1.0): super().__init__() self.fc nn.Linear(input_dim, num_experts) self.temperature nn.Parameter(torch.tensor(temperature)) def forward(self, x): logits self.fc(x) return F.softmax(logits / self.temperature, dim-1)温度系数在训练初期可以设高一点比如2.0让权重分布更平滑避免早期门控过早收敛到局部最优。训练后期再让它自然降低。实操心得门控网络的初始化很重要。我试过用Xavier初始化效果比默认的均匀初始化好很多。原因是Xavier能保持前向传播的方差稳定避免softmax输出在训练初期就饱和。2.2 专家网络的选型全连接还是更复杂的结构HoME论文里用的是全连接网络作为专家但在实际业务中专家的结构可以根据特征类型灵活调整。如果输入特征包含大量稀疏ID特征专家内部可以加一层embedding lookup如果输入包含序列特征专家可以换成Transformer层。我的建议是底层专家用简单的全连接中层和顶层专家可以适当加深。原因是底层专家数量多如果每个都太复杂计算量会爆炸。中层和顶层专家数量少可以承担更复杂的建模任务。另一个细节是专家之间的参数共享。在HoME里底层专家之间是完全独立的不共享参数。但中层专家可以共享一部分底层参数比如中层专家的第一层可以复用底层专家的最后一层。这个技巧能减少参数量同时让中层专家更快收敛。我试过在中层专家上做参数共享训练速度提升了约20%效果没有明显下降。2.3 梯度流动与负载均衡MoE架构有一个经典问题专家坍塌。也就是少数专家被频繁激活大部分专家得不到足够的梯度更新逐渐变成“死专家”。在HoME里这个问题更复杂因为有三层专家需要平衡。解决思路有两个层面。第一是负载均衡损失在训练时加一个辅助损失惩罚专家激活频率的方差。具体来说对于每一层专家计算每个专家被选中的概率然后计算这些概率的变异系数作为损失的一部分。这个损失的权重需要调太大影响主任务效果太小起不到作用。我的经验值是0.01到0.1之间。第二是门控噪声。在训练时给门控的logits加一点高斯噪声增加探索性。噪声的标准差可以设成0.1训练后期逐渐降到0。这个技巧在推荐场景里特别有效能显著减少死专家的数量。def load_balance_loss(gate_weights): # gate_weights: [batch_size, num_experts] mean_prob gate_weights.mean(dim0) # [num_experts] cv mean_prob.std() / (mean_prob.mean() 1e-8) return cv注意负载均衡损失不是越大越好。我试过把权重设成1.0结果模型完全偏向均匀分配任务效果反而下降。建议从0.01开始观察专家激活分布再调整。3. 实操过程从数据准备到模型上线的完整流程这一章我以一个真实的推荐场景为例把HoME的落地流程完整走一遍。场景是某内容平台的信息流推荐需要同时优化四个目标点击率CTR、完播率VTR、点赞率LTR、分享率STR。四个任务的相关性差异很大CTR和VTR中等相关LTR和STR高度相关CTR和STR几乎不相关。3.1 数据准备与特征工程多任务学习的数据准备和单任务最大的区别在于标签对齐。四个任务的标签并不是在每个样本上都齐全的。比如用户点击了但没看完VTR的标签就是0用户没点击VTR的标签就是缺失。处理缺失标签的方式直接影响模型效果。我的做法是只对有点击的样本计算VTR损失对没点击的样本屏蔽VTR损失。具体实现上用一个mask向量标记哪些样本对哪些任务有效在计算损失时乘以mask。这样做的理由是没点击的样本的完播行为是未知的强行给0标签会引入噪声。特征工程方面HoME的输入特征可以分成三类用户侧特征用户ID embedding、历史行为序列、物品侧特征物品ID embedding、类别、时长、上下文特征时间、位置、网络状态。底层专家接收全部特征中层专家接收底层输出顶层专家接收底层和中层输出。特征预处理有一个细节数值特征的归一化方式会影响门控的稳定性。如果数值特征的范围差异很大比如时长是0-3600秒而点赞数是0-100门控网络容易被大量级特征主导。我建议对所有数值特征做分位数归一化把范围压到0-1之间。3.2 模型搭建与参数配置基于PyTorch的HoME实现核心模块包括专家层、门控层和任务塔。下面是一个简化版的代码框架class HoME(nn.Module): def __init__(self, input_dim, num_tasks, e116, e28, e34): super().__init__() self.num_tasks num_tasks # 底层专家 self.bottom_experts nn.ModuleList([ nn.Sequential(nn.Linear(input_dim, 128), nn.ReLU(), nn.Linear(128, 64), nn.ReLU()) for _ in range(e1) ]) # 中层专家 self.middle_experts nn.ModuleList([ nn.Sequential(nn.Linear(e1 * 64, 64), nn.ReLU(), nn.Linear(64, 32), nn.ReLU()) for _ in range(e2) ]) # 顶层专家每个任务一个 self.top_experts nn.ModuleList([ nn.Sequential(nn.Linear(e1 * 64 e2 * 32, 32), nn.ReLU(), nn.Linear(32, 16), nn.ReLU()) for _ in range(e3) ]) # 门控网络 self.bottom_gates nn.ModuleList([ Gate(input_dim, e1) for _ in range(num_tasks) ]) self.middle_gates nn.ModuleList([ Gate(e1 * 64, e2) for _ in range(num_tasks) ]) self.top_gates nn.ModuleList([ Gate(e1 * 64 e2 * 32, e3) for _ in range(num_tasks) ]) # 任务塔 self.towers nn.ModuleList([ nn.Sequential(nn.Linear(64 32 16, 32), nn.ReLU(), nn.Linear(32, 1)) for _ in range(num_tasks) ]) def forward(self, x): # 底层专家 bottom_outs torch.stack([e(x) for e in self.bottom_experts], dim1) # bottom_outs: [batch, e1, 64] bottom_flat bottom_outs.view(bottom_outs.size(0), -1) # 中层专家 middle_outs torch.stack([e(bottom_flat) for e in self.middle_experts], dim1) middle_flat middle_outs.view(middle_outs.size(0), -1) # 顶层专家 top_input torch.cat([bottom_flat, middle_flat], dim-1) top_outs torch.stack([e(top_input) for e in self.top_experts], dim1) # 每个任务的门控加权 task_outputs [] for t in range(self.num_tasks): w_b self.bottom_gates[t](x) # [batch, e1] w_m self.middle_gates[t](bottom_flat) # [batch, e2] w_t self.top_gates[t](top_input) # [batch, e3] h_b torch.einsum(be,bed-bd, w_b, bottom_outs) h_m torch.einsum(be,bed-bd, w_m, middle_outs) h_t torch.einsum(be,bed-bd, w_t, top_outs) h torch.cat([h_b, h_m, h_t], dim-1) task_outputs.append(self.towers[t](h)) return torch.cat(task_outputs, dim-1)参数配置上E116、E28、E34是一个比较稳妥的起点。学习率用Adam初始值1e-3每5000步衰减0.9。Batch size设成2048训练数据量在千万级别时大约需要3-5个epoch收敛。3.3 训练技巧与损失函数设计多任务学习的损失函数是加权和总损失 Σ w_i * L_i λ * L_balance。每个任务的损失权重w_i需要仔细调。我的经验是不要用固定的权重而是用不确定性加权Uncertainty Weighting。这个方法让模型自己学习每个任务的权重公式是L_total Σ (1 / (2 * σ_i^2)) * L_i log(σ_i)其中σ_i是可学习的参数初始值设为1.0。这样模型会自动给难任务更高的权重给简单任务更低的权重。我在四个任务的场景里试过不确定性加权比手动调权重的效果提升了约1.5%的AUC。训练过程中还有一个细节门控网络的预热。在训练初期专家网络还没有学好门控的权重分配没有意义。我建议前1000步固定门控为均匀分布只训练专家网络。1000步之后再解冻门控让它们开始学习。这个技巧能显著提升训练稳定性。实操心得HoME的训练时间比MMoE长约30%主要开销在中层和顶层专家的计算上。如果线上推理延迟敏感可以考虑把中层专家的数量减半或者用矩阵分解的方式近似中层专家的计算。4. 常见问题与排查技巧实录HoME落地过程中遇到的问题不少我挑几个最有代表性的整理成速查表方便对照排查。4.1 专家坍塌与门控饱和现象训练几个epoch后底层16个专家中只有3-4个被激活其余专家的门控权重接近0。中层和顶层也有类似情况。排查思路先看门控权重的分布。如果某个专家的权重在所有样本上都低于0.01基本可以判定为死专家。再看梯度死专家的梯度范数通常比活专家低一个数量级。解决方法增加负载均衡损失的权重从0.01逐步提高到0.05。在门控logits上加噪声训练时噪声标准差0.1推理时去掉。检查专家初始化确保每个专家的初始输出方差相近。如果某个专家的初始化方差特别大它会主导早期训练导致其他专家得不到梯度。我遇到过一次专家坍塌排查后发现是底层专家的初始化用了不同的随机种子导致某些专家的初始输出范围差异很大。统一初始化方式后问题解决。4.2 任务间负迁移现象某个任务的AUC比单任务模型还低说明多任务学习反而拖累了它。排查思路计算任务之间的梯度冲突。具体做法是取一个batch的数据分别计算两个任务的梯度然后计算两个梯度向量的余弦相似度。如果余弦相似度为负说明两个任务在优化方向上冲突。解决方法调整任务在HoME中的层级归属。把冲突严重的任务放到不同的任务簇里让它们在中层就分开。使用梯度裁剪或PCGradProject Conflicting Gradients方法在梯度层面消除冲突。如果冲突无法调和考虑把冲突任务从HoME中拿出来单独建模。我在一个场景里遇到过CTR和STR的负迁移余弦相似度是-0.3。后来把STR从HoME中拆出来单独做了一个小模型CTR的AUC恢复了STR的效果也没受影响。4.3 线上推理延迟过高现象离线AUC提升明显但线上推理延迟从20ms涨到80ms无法接受。排查思路用profiler工具定位耗时模块。HoME的耗时主要在三部分底层专家的并行计算、中层专家的输入拼接、门控的softmax。解决方法底层专家可以用矩阵乘法一次性计算而不是循环调用每个专家。把E1个专家的权重拼成一个大矩阵输入做一次矩阵乘法就能得到所有专家的输出。中层专家的输入维度是E1×64如果E116维度就是1024这个拼接操作可以用view实现避免内存拷贝。门控的softmax可以用近似计算比如只保留top-2的专家权重其余置零。这个技巧能减少约40%的计算量效果损失不到0.1%。# 底层专家的矩阵化计算 W_bottom torch.stack([e[0].weight for e in self.bottom_experts]) # [e1, input_dim, 128] b_bottom torch.stack([e[0].bias for e in self.bottom_experts]) bottom_out torch.einsum(bi,eid-bed, x, W_bottom) b_bottom bottom_out F.relu(bottom_out)4.4 常见问题速查表问题现象可能原因排查方法解决措施专家坍塌负载不均衡、初始化不当检查门控权重分布和梯度范数增加负载均衡损失、加门控噪声、统一初始化任务负迁移任务梯度冲突计算梯度余弦相似度调整任务簇、PCGrad、拆分任务推理延迟高专家计算未矩阵化、门控全量计算Profiler定位耗时模块矩阵化专家计算、top-k门控近似训练不收敛学习率过大、门控过早收敛观察损失曲线和门控熵降低学习率、门控预热、温度系数退火某任务效果差该任务在底层被主导检查该任务的门控权重分布给该任务增加顶层专家、调整损失权重提示HoME的调参优先级是先调专家数量比例再调负载均衡损失权重最后调学习率和温度系数。不要一上来就调学习率那样会掩盖结构问题。5. 效果对比与适用场景分析HoME不是万能的它有明确的适用边界。这一章我用实测数据对比HoME和MMoE、单任务模型的效果然后给出选型建议。5.1 离线实验对比在四个任务的推荐场景下我做了三组对比实验单任务模型每个任务独立训练、MMoE8个专家、HoME1684。评价指标是AUC训练数据量1000万测试集200万。模型CTR AUCVTR AUCLTR AUCSTR AUC平均AUC训练时间单任务0.7820.7510.8030.7760.7781xMMoE0.7890.7580.8090.7810.7841.3xHoME0.7940.7630.8150.7880.7901.7x从数据看HoME相比MMoE平均AUC提升了0.6个百分点相比单任务提升了1.2个百分点。提升主要来自任务间相关性差异大的场景。如果四个任务高度相关HoME和MMoE的差距会缩小到0.2个百分点以内。训练时间上HoME比MMoE多30%比单任务多70%。这个开销主要来自中层和顶层专家的计算。如果线上延迟敏感可以考虑减少中层专家数量或者用知识蒸馏把HoME压缩成一个小模型。5.2 什么场景适合用HoME根据我的经验HoME在以下场景里收益最明显任务数量多且相关性差异大。比如同时优化点击、转化、复购、分享、评论五个目标任务之间的相关性从0.8到0.1不等。HoME的分层结构能让高相关任务在中层共享低相关任务在顶层独立。数据量充足。HoME的参数量比MMoE大需要足够的数据来训练。我的经验是每个任务的样本量至少要在百万级别否则中层和顶层专家容易过拟合。对可解释性有要求。HoME的中层门控权重可以反映任务簇结构这对业务理解有帮助。我做过一个实验把中层门控权重做PCA降维发现任务在二维平面上自然聚成三簇和业务上的人工分组完全一致。反过来以下场景不建议用HoME任务数量少2-3个且相关性高。这种情况下MMoE已经够用HoME的额外复杂度换不来明显收益。线上延迟极度敏感。HoME的推理延迟比MMoE高30%-50%如果业务要求延迟在10ms以内可能需要考虑模型压缩或蒸馏。数据量不足。每个任务样本量低于50万时HoME容易过拟合建议先用MMoE或单任务模型。5.3 从MMoE迁移到HoME的实操建议如果你已经在用MMoE想迁移到HoME我建议分三步走。第一步保持底层专家不变只加中层。把MMoE的专家池作为HoME的底层中层专家数量设为底层的一半。这样改动最小能快速验证HoME是否有效。第二步调整门控结构。MMoE的门控是任务级的HoME需要增加中层门控。中层门控的输入是底层专家的输出拼接输出维度是中层专家数量。初始化时让中层门控的权重接近均匀分布避免早期震荡。第三步逐步增加顶层专家。顶层专家是任务专属的可以从每个任务一个开始观察效果后再决定是否增加。如果某个任务的顶层专家权重一直很低说明该任务不需要专属专家可以去掉。实操心得迁移过程中建议保留MMoE的模型作为baseline用A/B测试对比。我试过直接替换结果因为门控初始化问题导致效果下降后来用渐进式迁移才稳定下来。6. 工程落地中的几个关键决策最后聊几个工程落地时的决策点这些在论文里不会写但实际做的时候绕不开。6.1 专家数量的确定E1、E2、E3怎么定我的经验公式是E1 min(16, 任务数 × 4)E2 min(8, 任务数 × 2)E3 任务数。这个公式的出发点是底层需要足够的专家来捕捉多样化的基础特征但太多会导致训练不充分中层专家对应任务簇数量应该是任务数的一半左右顶层专家每个任务一个保证任务独立性。当然这只是一个起点。实际调参时我会先固定E2和E3只调E1。E1从8开始每次翻倍观察验证集AUC。通常E1在16到32之间会达到饱和再增加收益很小。6.2 门控网络的输入选择门控网络的输入用什么特征直接影响门控的学习效果。我的建议是底层门控用任务的原始特征中层和顶层门控用专家输出。原因是底层门控需要根据任务特性选择基础专家原始特征包含了任务ID、用户画像等信息足够做这个决策。中层和顶层门控需要根据底层专家的输出做更细粒度的选择用专家输出作为输入更合适。有一个细节如果任务ID是one-hot特征底层门控可以直接用任务ID的embedding作为输入。这样门控网络能学到每个任务的偏好而且参数量很小。6.3 模型更新与线上服务HoME的线上服务需要考虑两个问题模型更新频率和推理优化。模型更新方面我建议全量更新每天一次增量更新每小时一次。增量更新时只更新底层专家和门控中层和顶层专家保持不动。原因是中层和顶层专家参数量大频繁更新容易导致线上效果震荡。推理优化方面除了前面提到的矩阵化计算和top-k门控还可以做专家缓存。对于高频用户和物品它们的底层专家输出可以缓存起来下次请求直接复用。这个技巧能减少约20%的计算量但需要注意缓存失效策略。注意HoME的线上服务对内存要求较高因为需要同时加载三层专家的参数。如果内存紧张可以考虑用float16量化效果损失很小。6.4 与其他多任务方法的对比除了MMoE多任务学习还有几个常见方法PLEProgressive Layered Extraction、CGCCustomized Gate Control、STEMShared and Task-specific Embedding。HoME和它们的区别在于PLE和CGC也是分层结构但它们的层级是“共享层任务专属层”两层没有中层任务簇的概念。HoME的三层结构更适合任务数量多、需要中间粒度共享的场景。STEM是embedding层面的多任务方法和HoME不在一个层面。实际应用中HoME可以和STEM结合底层用STEM做embedding共享上层用HoME做专家分层。我个人的选型建议是任务数≤3用MMoE任务数4-6用PLE或CGC任务数≥6且相关性差异大用HoME。当然这不是绝对的还要看数据量和延迟要求。6.5 一个真实的踩坑记录最后分享一个我踩过的坑。有一次上线HoME后离线AUC涨了0.8个百分点但线上CTR反而跌了0.3%。排查了很久最后发现是门控网络的输出分布偏移。离线训练时门控的输入特征分布和线上不一致导致线上推理时门控权重分布和离线差异很大。解决方法是在门控网络前加一个BatchNorm层对输入特征做归一化。这个改动让线上线下的门控输出分布对齐了线上CTR恢复了正常。这个坑让我意识到多任务模型的门控对特征分布非常敏感任何特征预处理的不一致都会被门控放大。另一个坑是专家数量的奇偶性。我试过E115结果发现有一个专家始终得不到足够的梯度。后来改成16问题消失。原因是奇数个专家在负载均衡时容易出现“落单”的专家建议专家数量用2的幂次。这些经验在论文里找不到但实际做工程的时候往往就是这些细节决定成败。HoME是一个强大的多任务学习框架但它不是即插即用的。理解它的设计原理结合业务场景做适配才能真正发挥它的价值。
返回列表