Modeling Task Relationships in Multi-task Learning with Multi-gate Mixture-of-Experts

By Jiaqi W. Ma, Zhe Zhao, Xinyang Yi, Jilin Chen, Lichan Hong, Ed H. Chi. 2018

1,663 citations recorded by the source corpus.

Neural-based multi-task learning has been successfully used in many real-world large-scale applications such as recommendation systems. For example, in movie recommendations, beyond providing users movies which they tend to purchase and watch, the system might also optimize for users liking the movies afterwards. With multi-task learning, we aim to build a single model that learns these multiple goals and tasks simultaneously. However, the prediction quality of commonly used multi-task models is often sensitive to the relationships between tasks. It is therefore important to study the modeling tradeoffs between task-specific objectives and inter-task relationships. In this work, we propose a novel multi-task learning approach, Multi-gate Mixture-of-Experts (MMoE), which explicitly learns to model task relationships from data. We adapt the Mixture-of-Experts (MoE) structure to multi-task learning by sharing the expert submodels across all tasks, while also having a gating network trained to optimize each task. To validate our approach on data with different levels of task relatedness, we first apply it to a synthetic dataset where we control the task relatedness. We show that the proposed approach performs better than baseline methods when the tasks are less related. We also show that the MMoE structure results in an additional trainability benefit, depending on different levels of randomness in the training data and model initialization. Furthermore, we demonstrate the performance improvements by MMoE on real tasks including a binary classification benchmark, and a large-scale content recommendation system at Google.

Subfields: Computer Vision and Pattern Recognition, Information Systems, Management Science and Operations Research, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computer Vision and Pattern Recognition

In Information Systems

In Management Science and Operations Research

In Optimization & Training

Citations in the corpus

Massively Multilingual Neural Machine Translation in the Wild: Findings and Challenges (2019, background): “Oftentimes these metrics vary in their effectiveness across languages; WMT shared tasks (Ma et al., 2018b) report that the speciﬁc language, dataset, and system signiﬁcantly affect which metric has the strongest correlation with human ratings.”
Pathways: Asynchronous Distributed Dataflow for ML (2022, background): “Graph neural network training (Jia et al., 2020), neural architecture search (Pham et al., 2018), and multi-modal multi-task learning systems (Ma et al., 2018; Lepikhin et al., 2020; Zhao et al., 2019) are examples of inherently heterogeneous and dynamic tasks that do not ﬁt naturally in the SPMD…”
A Review of Sparse Expert Models in Deep Learning (2022, background): “Ma et al. (2018) introduces a multi-gate routing algorithm where each task gets it own unique gating function.”
DSelect-k: Differentiable Selection in the Mixture of Experts with Applications to Multi-Task Learning (2021, neutral): “In the rest of this section, we present experiments on the following real MTL datasets: MovieLens, Multi-MNIST, Multi-Fashion MNIST, and on a real-world, large-scale recommender system.”
Dynamic Neural Networks: A Survey (2021, background): “Others Time series forecasting [248], [249], [250], Link prediction [251], Recommendation system [77], [252], [253], [254] Graph classification [121], Document classification [156], [255], [256], [257], Stereo confidence estimation [258]”
Deep Mixture of Experts via Shallow Embedding (2018, neutral): “, video captioning [26], multi-task learning [20], etc.”
Pareto Multi-Task Learning (2019, methodology): “Similar to the setting in [35], we set the income, education level, and marital status as three binary targets to be predicted: • Task 1: whether the person’s income exceeds $50K/year.”
Understanding and Improving Knowledge Distillation (2020, methodology): “Inspired by (Ma et al., 2018), we synthesize a classification dataset with K classes and C super-classes, such that each super-class has K/C classes, and each class will be assigned with a carefully generated basis vector, so that we could control the class correlations within the same super-class.”
Multi-Task Reinforcement Learning with Soft Modularization (2020, methodology): “Our method is also related to works on mixture of experts [14, 7, 13, 36, 2, 20].”
Designing Effective Sparse Expert Models (2022, methodology): “Ma et al. (2018) recommended using a separate gating or router network for each task, an idea that may soon be revisited for Transformer architectures.”
Recommending what video to watch next: a multitask ranking system (2019, methodology): “Our implementation of the expert networks is identical to multilayer perceptrons with ReLU activations [30].”
Improving Video-Text Retrieval by Multi-Stream Corpus Alignment and Dual Softmax Loss (2021, background): “A gate module (Ma et al. 2018) is added to integrate entity and action representations with fusion ones, which improves fusion expert’s performance.”
Progressive Layered Extraction (PLE): A Novel Multi-Task Learning (MTL) Model for Personalized Recommendations (2020, methodology): “To evaluate the performance of PLE, we conduct extensive experiments on real-world industrial recommendation dataset and major available public datasets including census-income [5], synthetic data [13] and Ali-CCP 1.”
HyperGrid: Efficient Multi-Task Transformers with Grid-wise Decomposable Hyper Projections (2020, methodology): “has also been a go-to strategy for continual learning [Yoon et al., 2019]. Outside of the NLP domain, ﬂexible parameter sharing approaches are also dominant strategies for learning multi-task models [Ma et al., 2018, 2019]. The key novelty behind our work lies in the decomposable and factorized form”
Beyond Distillation: Task-level Mixture-of-Experts for Efficient Inference (2021, background): “This approach has seen success in a variety of domains such as classification, recommender systems and NLP ((Ma et al., 2019, 2018; Clark et al., 2019; Collobert and Weston, 2008; Ruder et al., 2019; Tan et al., 2019)).”
Small Towers Make Big Differences (2020, methodology): “…multi-task deep learning has gained popularity through its success in a wide range of applications, including natural language processing (Collobert and Weston, 2008), computer vision (Girshick, 2015; Ren et al., 2015), and online recommendation systems (Bansal et al., 2016; Ma et al., 2018, 2019).”
A Survey of Graph Neural Networks for Recommender Systems: Challenges, Methods, and Directions (2021, background): “Recently, some works of deep learning-based recommendation, such as MMOE [120], PLE [165], and so on, studied how to deploy multiple separate models, each of which is designed towards one objective, and then optimize them together.”
Gumbel-Matrix Routing for Flexible Multi-task Learning (2019, methodology): “A.1 GENERATION We follow a process for synthetic data generation based on (Ma et al., 2018), which allows for explicit control of the relatedness between tasks.”
ST-MoE: Designing Stable and Transferable Sparse Expert Models (2022, methodology): “Ma et al. (2018) recommended using a separate gating or router network for each task, an idea that may soon be revisited for Transformer architectures.”
Multi-Task Learning with Deep Neural Networks: A Survey (2020, background): “(Zhang et al., 2014; Dai et al., 2016; Zhao et al., 2018; Liu et al., 2019; Ma et al., 2018) propose architectures which are variations on the shared trunk idea.”