Modeling Task Relationships in Multi-task Learning with Multi-gate Mixture-of-Experts/zh
| Research Paper | |
|---|---|
| Authors | Jiaqi Ma; Zhe Zhao; Xinyang Yi; Jilin Chen; Lichan Hong; Ed H. Chi |
| Year | 2018 |
| Venue | Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining (KDD '18) |
| Topic area | Machine Learning |
| Difficulty | Research |
| Source | View paper |
| Download PDF | |
Modeling Task Relationships in Multi-task Learning with Multi-gate Mixture-of-Experts 是 Jiaqi Ma、Zhe Zhao、Xinyang Yi、Jilin Chen、Lichan Hong 和 Ed H. Chi 于 2018 年发表的论文,刊于 Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining(KDD '18)。该论文提出了 Multi-gate Mixture-of-Experts(MMoE) 架构:由一组共享的专家子网络与每个任务各自独立的 softmax 门控网络组合而成,使一个多任务神经网络能够灵活建模任务之间的关系,而无需显著增加参数量。MMoE 已成为工业界多任务学习的基础技术之一,被广泛部署于大规模推荐系统中(包括 YouTube),并被认为可同时提升离线 AUC 以及线上互动与满意度指标。
概述
神经网络多任务学习的目标是训练单一模型,同时预测若干相关目标——例如同时预测用户是否会观看一个被推荐的内容,以及之后是否会喜欢它。在任务之间共享表示可望带来更高的样本效率和正则化效果,但在实践中,占主导地位的 Shared-Bottom 架构的预测质量对任务相关性非常敏感:当多个任务在共享主干中争夺容量时,梯度会相互冲突,联合训练可能反而不如单任务模型。早期的改进方法——如 L2-Constrained 网络、Cross-Stitch 网络以及张量分解多任务网络——以软约束取代硬性共享,但通常会增加大量任务特定参数,并失去推理时的效率优势,而这种效率正是生产环境采用多任务模型的原因之一。
MMoE 将单一的共享主干替换为一组 feed-forward 专家网络,并为每个任务配备各自的门控网络。门控是输入的线性 softmax 函数,逐样本地输出在专家上的混合权重;每个任务因此使用同一组专家、但以各自由输入决定的混合方式来消费它们。当任务相似时,各门控会收敛到重叠的专家上并从共享表示中获益;当任务冲突时,门控学到将不同任务路由到互不相交的专家上,模型由此回到接近独立模型的行为——而所有这些变化都不会触及专家池本身。
作者在三种现实程度递增的场景中验证 MMoE:一个任务相关性可控的合成回归基准、UCI Census-income 二分类基准,以及一个在数百亿用户反馈事件上训练的 Google 内容推荐系统。在所有三种场景中,MMoE 都能匹敌或超越此前基于软参数共享的基线方法,同时保留 Shared-Bottom 模型轻量的计算特性。
主要贡献
- 面向多任务神经网络的 Multi-gate Mixture-of-Experts(MMoE) 架构:由共享的专家网络池构成,每个任务通过各自基于输入的 softmax 门控对其进行组合。
- 基于正弦回归的任务相关性受控合成研究,以权重向量的余弦相似度作为可调代理量,刻画标签 Pearson 相关性,从而隔离地考察任务相关性下降时多任务模型的行为变化。
- 可训练性分析,表明 MMoE 不仅在平均损失上优于 Shared-Bottom 和 One-gate MoE(OMoE)基线,而且对随机初始化的方差显著更小——也就是说更难陷入劣质局部极小。
- 在 UCI Census-income 数据集上的基准结果,在相同参数预算下,MMoE 与 L2-Constrained、Cross-Stitch 以及 Tensor-Factorization 多任务基线持平或更优。
- 来自 Google 推荐系统的生产级证据:MMoE 在工程上线的 Shared-Bottom 模型基础上提升了互动 AUC 和离线 R²,并在线上互动与满意度指标上取得了具有统计显著性的提升,且未明显增加服务成本。
方法
记 $ K $ 为任务数。标准的 Shared-Bottom 多任务模型由一个共享的编码器 $ f $ 和每个任务各自的塔(tower)$ h_k $ 组成:
- $ y_k = h_k(f(x)). $
MMoE 将单一编码器 $ f $ 替换为一组 $ n $ 个专家网络 $ f_1, \ldots, f_n $,并为每个任务引入一个 softmax 门控网络 $ g^k $:
- $ y_k = h_k\!\left(\sum_{i=1}^{n} g^k(x)_i\, f_i(x)\right),\qquad g^k(x) = \mathrm{softmax}(W_{g_k}\, x), $
其中 $ W_{g_k} \in \mathbb{R}^{n \times d} $ 是每个任务各自可训练的矩阵。每个专家是一个使用 ReLU 激活的 feed-forward MLP;门控被有意设计得很轻量,以便相对于宽度相当的 Shared-Bottom 模型,新增的参数开销可以忽略。论文还引入了 One-gate MoE(OMoE) 基线——所有任务共用同一个门控——以分离任务级门控带来的贡献与 MoE 结构本身的贡献。
在合成研究中,论文用余弦相似度受控为 $ p $ 的两个权重向量 $ w_1, w_2 $ 生成两个回归标签:
- $ w_1 = c\, u_1,\qquad w_2 = c\!\left(p\, u_1 + \sqrt{1 - p^2}\, u_2\right), $
其中 $ u_1 \perp u_2 $,标签由 $ w_k^T x $ 的正弦函数的非线性混合再加上高斯噪声生成。余弦相似度 $ p $ 作为标签经验 Pearson 相关性的可控代理,提供了一个清晰的坐标轴,沿其变化即可调节任务相关性。
在 Census-income 基准上,作者基于人口统计特征构造两组任务对(收入 vs. 婚姻状况;受教育程度 vs. 婚姻状况)。在生产推荐系统中,两个二分类任务——一个与互动相关的信号、一个与满意度相关的信号——在数百亿条用户反馈事件上联合训练,所有基线均通过高斯过程超参数搜索调优,并约束每层最多 2048 个隐藏单元。
结果
在合成基准上,MMoE 大幅缩小了高相关性与低相关性区间之间的差距,远超 OMoE 和 Shared-Bottom,并在每个相关性水平上经 200 次独立运行的平均损失上均占优。OMoE 基线因缺少任务级门控,当任务相关性下降时性能急剧下滑——这证实了任务级门控才是该设计的核心要素。可训练性直方图进一步显示,Shared-Bottom 出现了长长的劣质局部极小尾巴,而 MMoE 的结果集中在最优可达损失附近。
在 UCI Census-income 上,MMoE 在两组(收入/婚姻状况和受教育程度/婚姻状况)主任务上的平均 AUC 均最高,小幅领先 L2-Constrained 和 Cross-Stitch,并大幅超越 Tensor-Factorization——后者在低任务相关性下表现崩溃。在辅助的婚姻状况任务上,单任务模型仍保持小幅领先,因为它针对该任务做了超参数微调,而多任务模型仅针对主任务调参。
在 Google 内容推荐系统中,MMoE 在每个训练检查点(2M、4M、6M 步)上都取得了最高的互动 AUC 和 R²。L2-Constrained 和 Cross-Stitch 甚至落后于 Shared-Bottom,因为它们的参数量大致翻倍,导致约束不足。线上 A/B 测试显示,MMoE 相对生产中的 Shared-Bottom 模型,互动指标提升 +0.25%,满意度指标提升 +2.65%,均在 95% 置信水平下显著——并且关键在于没有明显的服务成本增加,因为专家共享保留了 Shared-Bottom 的效率优势。
影响
MMoE 成为工业界最广泛采用的多任务架构之一,在大规模推荐、排序和广告场景中尤为常见。该设计影响了后续的带门控的多任务学习工作,包括 Customized Gate Control(CGC)和 Progressive Layered Extraction(PLE),并启发了更广义的 Sparsely-gated mixture-of-experts 与条件计算架构家族——它们使参数量增长而每样本 FLOPs 不变。其核心洞见——即在共享的专家池之上,使用任务级、依赖输入的门控,可以将任务冲突与共享表示的收益解耦——已延续到 transformer 时代的 MoE 设计中,被现代大语言模型广泛采用。该论文也常被引用作为门控机制改善非凸深度网络可训练性的证据,与门控循环单元方面的同类发现相互呼应。
参见
参考文献
- Jiaqi Ma, Zhe Zhao, Xinyang Yi, Jilin Chen, Lichan Hong, and Ed H. Chi. 2018. Modeling Task Relationships in Multi-task Learning with Multi-gate Mixture-of-Experts. 见 Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining(KDD '18)。ACM, New York, NY, USA, 1930–1939. https://doi.org/10.1145/3219819.3220007
- Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. 1991. Adaptive mixtures of local experts. Neural Computation 3, 1, 79–87.
- Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. 2017. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv:1701.06538.
- Rich Caruana. 1998. Multitask learning. 见 Learning to learn. Springer, 95–133.
- Ishan Misra, Abhinav Shrivastava, Abhinav Gupta, and Martial Hebert. 2016. Cross-stitch networks for multi-task learning. 见 CVPR. 3994–4003.
- Long Duong, Trevor Cohn, Steven Bird, and Paul Cook. 2015. Low Resource Dependency Parsing: Cross-lingual Parameter Sharing in a Neural Network Parser. 见 ACL (2). 845–850.
- Yongxin Yang and Timothy Hospedales. 2016. Deep multi-task representation learning: A tensor factorisation approach. arXiv:1605.06391.
- Paul Covington, Jay Adams, and Emre Sargin. 2016. Deep neural networks for YouTube recommendations. 见 RecSys. ACM, 191–198.