MMoE学习笔记利用门控专家网络高效建模多任务关系一、背景与动机在现代推荐系统、广告点击率CTR预估、自然语言处理等场景中往往需要同时预测多个相关任务如“点击率”和“转化率”。传统做法是训练多个独立模型但这样忽略了任务之间的关联信息导致模型泛化能力不足。而共享底层网络Shared-Bottom虽能捕捉共性却容易因任务冲突如“点击”和“购买”的目标不一致导致“跷跷板”现象一个任务性能提升另一个任务下降。MMoEMulti-gate Mixture-of-Experts多门控专家混合模型由Google在2019年提出核心思想是通过多个“专家网络”和“门控机制”动态分配任务所需的表示空间。每个专家网络学习不同的数据模式门控网络根据输入样本为不同任务选择不同的专家组合从而实现任务间的灵活共享与独立学习。### 二、MMoE核心原理MMoE的架构包含三个关键组件1.专家网络Experts多个并行的子网络通常为全连接层每个专家负责学习输入数据的某种隐式模式。2.门控网络Gates每个任务对应一个独立的门控网络输出一个概率分布softmax用于加权融合各专家的输出。3.任务塔Towers每个任务使用融合后的专家输出作为特征进行最终的预测。数学上假设有 ( k ) 个专家输入 ( x ) 经过每个专家得到 ( f_i(x) )对于任务 ( j )门控网络输出权重 ( g_j(x) \in \mathbb{R}^k )则任务 ( j ) 的融合表示为[y_j \text{Tower}j\left( \sum{i1}^k g_j(x)i \cdot f_i(x) \right)]其中 ( g_j(x) \text{softmax}(W{gj} x b_{gj}) )。### 三、实战代码从零实现MMoE模型下面用PyTorch实现一个简单的MMoE模型用于多任务回归问题如同时预测物品的评分和销量。我们将使用合成数据演示训练过程。pythonimport torchimport torch.nn as nnimport torch.optim as optimfrom torch.utils.data import DataLoader, TensorDatasetimport numpy as np# 定义专家网络一个简单的全连接层class Expert(nn.Module): def __init__(self, input_dim, hidden_dim): super(Expert, self).__init__() self.fc nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim) ) def forward(self, x): return self.fc(x)# 定义门控网络每个任务一个独立的门控输出softmax权重class Gate(nn.Module): def __init__(self, input_dim, num_experts): super(Gate, self).__init__() self.gate nn.Sequential( nn.Linear(input_dim, num_experts), nn.Softmax(dim1) ) def forward(self, x): return self.gate(x)# 定义MMoE整体模型class MMoE(nn.Module): def __init__(self, input_dim, num_experts, expert_hidden_dim, num_tasks, task_hidden_dim): super(MMoE, self).__init__() self.num_experts num_experts self.num_tasks num_tasks # 创建多个专家网络 self.experts nn.ModuleList([ Expert(input_dim, expert_hidden_dim) for _ in range(num_experts) ]) # 创建每个任务对应的门控网络 self.gates nn.ModuleList([ Gate(input_dim, num_experts) for _ in range(num_tasks) ]) # 创建每个任务的任务塔输出层 self.towers nn.ModuleList([ nn.Sequential( nn.Linear(expert_hidden_dim, task_hidden_dim), nn.ReLU(), nn.Linear(task_hidden_dim, 1) # 假设每个任务输出一个标量 ) for _ in range(num_tasks) ]) def forward(self, x): # 所有专家输出[batch_size, num_experts, expert_hidden_dim] expert_outputs [expert(x) for expert in self.experts] # list of [batch, hidden] expert_outputs torch.stack(expert_outputs, dim1) # [batch, num_experts, hidden] task_outputs [] for i in range(self.num_tasks): # 门控权重[batch, num_experts] gate_weights self.gates[i](x) # [batch, num_experts] # 对专家输出进行加权求和[batch, hidden] weighted_expert torch.bmm(gate_weights.unsqueeze(1), expert_outputs).squeeze(1) # 通过任务塔得到最终预测 tower_out self.towers[i](weighted_expert) task_outputs.append(tower_out) return task_outputs # list of [batch, 1]# 生成合成数据假设有2个任务输入维度10torch.manual_seed(42)num_samples 1000X torch.randn(num_samples, 10)y1 X[:, 0] * 2 X[:, 1] * 0.5 torch.randn(num_samples) * 0.1 # 任务1线性组合y2 X[:, 2] ** 2 X[:, 3] * 3 torch.randn(num_samples) * 0.1 # 任务2非线性组合# 转换为DataLoaderdataset TensorDataset(X, torch.stack([y1, y2], dim1))dataloader DataLoader(dataset, batch_size32, shuffleTrue)# 初始化模型、损失函数、优化器model MMoE(input_dim10, num_experts4, expert_hidden_dim16, num_tasks2, task_hidden_dim8)criterion nn.MSELoss()optimizer optim.Adam(model.parameters(), lr0.01)# 训练循环num_epochs 50for epoch in range(num_epochs): total_loss 0.0 for batch_X, batch_y in dataloader: optimizer.zero_grad() outputs model(batch_X) # outputs: [output_task1, output_task2] # 计算两个任务的损失并求和可调整权重 loss1 criterion(outputs[0], batch_y[:, 0:1]) loss2 criterion(outputs[1], batch_y[:, 1:2]) loss loss1 loss2 loss.backward() optimizer.step() total_loss loss.item() if (epoch1) % 10 0: print(fEpoch {epoch1}, Loss: {total_loss/len(dataloader):.4f})代码说明- 专家网络使用两层全连接激活函数ReLU输出固定维度expert_hidden_dim。- 门控网络通过Softmax生成权重确保每个任务的专家权重之和为1。- 训练时两个任务损失直接相加实际应用中可根据任务重要性调整权重。### 四、进阶实战带正则化的MMoE与对比实验为了验证MMoE的效果我们将其与共享底层模型Shared-Bottom进行对比并加入Dropout正则化。pythonimport torch.nn.functional as F# 定义Shared-Bottom模型作为对比基线class SharedBottom(nn.Module): def __init__(self, input_dim, shared_hidden_dim, num_tasks, task_hidden_dim): super(SharedBottom, self).__init__() self.shared nn.Sequential( nn.Linear(input_dim, shared_hidden_dim), nn.ReLU(), nn.Dropout(0.2), nn.Linear(shared_hidden_dim, shared_hidden_dim) ) self.towers nn.ModuleList([ nn.Sequential( nn.Linear(shared_hidden_dim, task_hidden_dim), nn.ReLU(), nn.Dropout(0.2), nn.Linear(task_hidden_dim, 1) ) for _ in range(num_tasks) ]) def forward(self, x): shared_out self.shared(x) return [tower(shared_out) for tower in self.towers]# 定义带Dropout的MMoE增强泛化class MMoEWithDropout(MMoE): def __init__(self, input_dim, num_experts, expert_hidden_dim, num_tasks, task_hidden_dim, dropout0.2): super(MMoEWithDropout, self).__init__(input_dim, num_experts, expert_hidden_dim, num_tasks, task_hidden_dim) # 为每个专家添加dropout重写专家网络 self.experts nn.ModuleList([ nn.Sequential( nn.Linear(input_dim, expert_hidden_dim), nn.ReLU(), nn.Dropout(dropout), nn.Linear(expert_hidden_dim, expert_hidden_dim) ) for _ in range(num_experts) ]) # 为任务塔添加dropout self.towers nn.ModuleList([ nn.Sequential( nn.Linear(expert_hidden_dim, task_hidden_dim), nn.ReLU(), nn.Dropout(dropout), nn.Linear(task_hidden_dim, 1) ) for _ in range(num_tasks) ])# 训练对比实验def train_model(model, dataloader, epochs30): optimizer optim.Adam(model.parameters(), lr0.01) for epoch in range(epochs): for batch_X, batch_y in dataloader: optimizer.zero_grad() outputs model(batch_X) loss sum(criterion(out, batch_y[:, i:i1]) for i, out in enumerate(outputs)) loss.backward() optimizer.step() # 返回最终平均损失 final_loss 0.0 with torch.no_grad(): for batch_X, batch_y in dataloader: outputs model(batch_X) loss sum(criterion(out, batch_y[:, i:i1]) for i, out in enumerate(outputs)) final_loss loss.item() return final_loss / len(dataloader)# 实例化模型并训练shared_model SharedBottom(input_dim10, shared_hidden_dim16, num_tasks2, task_hidden_dim8)mmoe_model MMoEWithDropout(input_dim10, num_experts4, expert_hidden_dim16, num_tasks2, task_hidden_dim8)print(训练Shared-Bottom模型...)shared_loss train_model(shared_model, dataloader, epochs30)print(fShared-Bottom最终损失: {shared_loss:.4f})print(\n训练MMoE模型...)mmoe_loss train_model(mmoe_model, dataloader, epochs30)print(fMMoE最终损失: {mmoe_loss:.4f})# 观察门控权重分布可视化第一个任务的第一个batchsample_batch, _ next(iter(dataloader))with torch.no_grad(): gate_weights_task0 mmoe_model.gates[0](sample_batch[:5]) # 取前5个样本 print(\n第一个任务的门控权重前5个样本) print(gate_weights_task0.numpy())代码说明- 对比实验显示MMoE通常能比Shared-Bottom获得更低的联合损失尤其在任务冲突时。- 通过打印门控权重可以观察模型是否自动为不同样本分配不同的专家组合如某些样本倾向使用专家1和3另一些倾向专家2和4。### 五、关键设计细节与经验1.专家数量选择通常2-8个专家过多易过拟合过少则无法捕获多样性。可通过交叉验证调整。2.门控网络初始化建议门控权重初始均匀分布softmax输出接近1/k避免初始训练时门控偏向某个专家。3.任务权重平衡若任务重要性不同可在损失函数中加入权重如loss w1*loss1 w2*loss2。4.特征交互对于高维稀疏特征如推荐系统可在专家网络前嵌入层处理专家网络内部可增加BatchNorm。### 六、总结MMoE通过引入多个专家网络和独立门控机制实现了对多任务关系的灵活建模。相比传统共享底层模型它具有以下优势-任务特异性每个任务通过门控选择最相关的专家组合减少任务冲突。-参数共享效率专家网络在所有任务间共享但门控网络允许任务动态调节共享程度。-可解释性门控权重可反映不同样本对专家模式的偏好辅助理解数据分布。实际应用中MMoE已被广泛应用于推荐系统如同时预测点击、收藏、购买、广告竞价预估CTR和CVR以及多目标排序。需要注意的是当任务数量过多时门控网络数量线性增长可能带来训练开销此时可考虑PLEProgressive Layered Extraction等改进版本。通过本文的代码实战相信读者能快速上手MMoE并根据自身业务需求进行调整。多任务学习的核心在于“求同存异”——让模型既共享底层知识又尊重每个任务的独特性这正是MMoE的精髓所在。