Sparsely Activated Mixture-of-Experts are Robust Multi-Task Learners

TL;DR

The paper tackles interference in multi-task learning by introducing sparsely activated Mixture-of-Experts with task-aware gating (MT-TaG), enabling task-conditioned routing that preserves the original compute budget. By replacing FFN blocks with MoE layers and learning per-task gates, MT-TaG achieves superior transfer to low-resource tasks, better generalization to unseen related tasks, and robustness to unrelated tasks. Empirical results on GLUE and extended task mixtures show MT-TaG outperforms dense baselines and MT-Switch across multiple encoder sizes and task counts, demonstrating scalable, robust multi-task adaptation. The approach offers a principled balance between model capacity and shared representations, with implications for efficient deployment of large, multi-task NLP systems.

Abstract

Traditional multi-task learning (MTL) methods use dense networks that use the same set of shared weights across several different tasks. This often creates interference where two or more tasks compete to pull model parameters in different directions. In this work, we study whether sparsely activated Mixture-of-Experts (MoE) improve multi-task learning by specializing some weights for learning shared representations and using the others for learning task-specific information. To this end, we devise task-aware gating functions to route examples from different tasks to specialized experts which share subsets of network weights conditioned on the task. This results in a sparsely activated multi-task model with a large number of parameters, but with the same computational cost as that of a dense model. We demonstrate such sparse networks to improve multi-task learning along three key dimensions: (i) transfer to low-resource tasks from related tasks in the training mixture; (ii) sample-efficient generalization to tasks not seen during training by making use of task-aware routing from seen related tasks; (iii) robustness to the addition of unrelated tasks by avoiding catastrophic forgetting of existing tasks.

Figures (1)

• Figure 1: Sparse MoE layer with $3$ Experts, $2$ Tasks, and $top$-$1$ expert routing with (a) Shared Gating, and (b) Task-aware Gating. $x_{s,1}$ and $x_{s,2}$ are tokens from Task $1$ and $2$ respectively. They share the same gate $G$ in sub-figure (a), and routed to respective task-specific gates in sub-figure (b). For simplicity, we only show the pathway for $x_{s,2}$ with a solid line, and show the gating behavior for $x_{s,1}$ with a dashed red line