Multi-Task Learning via Task Multi-Clustering
Abstract
Multi-task learning has the potential to facilitate learning of shared representations between tasks, leading to better task performance.
Some sets of tasks are related, and can share many features that are useful latent representations for these tasks.
Other sets of tasks are less related, possibly sharing some features, but also competing on the representational resources of shared parameters.
We propose to discover how to share parameters between related tasks and split parameters between conﬂicting tasks, by learning a multi-clustering of the tasks.
We present a mixture-of-experts model, where each cluster is an expert that extracts a feature vector from the input, and each task belongs to a set of clusters whose experts it can mix.
In experiments on the CIFAR-100 MTL domain, multi-clustering outperforms a model that mixes all experts in accuracy and computation time.
The results suggest that the performance of our method is robust to regularization that increases the model’s sparsity when sufﬁcient data is available, and can beneﬁt from sparser models as data becomes scarcer.
1. Introduction
Multi-task learning (MTL) is the process of jointly learning to perform multiple tasks (Ruder, 2017).
When different tasks are related to each other, we might expect joint learning to provide beneﬁts such as achieving better performance at the tasks, requiring less data, and improving generalization and robustness.
The reasons why MTL might lead to such beneﬁts go beyond the mere availability of more training data or supervision signal: MTL can facilitate structured representation, i e., have some parameters shared between related tasks.
In this work, we focus on tasks that are related in the sense that they share featurizations of the input — i.e., latent representations — from which each task’s output can be computed.
A standard hypothesis is that learning models with shared parameters has the potential to lead to better performance.
However, when models for unrelated tasks are forced to share parameters, they compete on the same representational resources, which may hurt their performance.
Ideally, we would learn models jointly for tasks that share useful features, and separately for tasks that do not; but this principle meets two obstacles.
First, a pair of tasks may share some useful features but not others, so that a simple binary decision — whether to learn the two tasks jointly or separately — may be too coarse.
Second, it is often unclear how one could obtain prior information about the degree to which different tasks are related or competing in the structural sense.
In this work, we propose to overcome these obstacles by discovering a multi-clustering of the tasks, i.e., an assignment of each task to multiple clusters of tasks to which it is partially related.
Each cluster corresponds to an “expert” that extracts a vector of features from the input.
A model for the task is then obtained by a mixture of the experts to which the task is clustered.
Our gradient-based method jointly learns the experts, the multi-clustering, and the mixtures.
We present experimental results on the CIFAR-100 MTL domain suggesting that our method has both computational and performance beneﬁts over existing methods with the same neural architecture and number of experts.
The results further provide evidence that regularizing for model sparsity, which decreases the number of experts used by tasks, can improve performance when data is somewhat scarce.
Experimenting with a range of regularization coefﬁcients
when data is more abundant indicates that performance is largely insensitive to small changes in that hyper-parameter, and that our method of selecting experts during evaluation time outperforms simpler baselines.
In summary, this paper contributes: (1) a method for multitask learning via task multi-clustering using discrete gating; (2) an evaluation of the computational and performance beneﬁts of this method; and (3) evidence that the method’s performance is robust to variations in the regularization hyper-parameter in the high-data regime and can beneﬁt from sparsity regularization in the lower-data regime.
2. Related Work
Extensive research has been done in multi-task learning, seeking to improve task performance (Misra et al., 2016;
Pinto & Gupta, 2017) and reduce the requirement of computational (Shazeer et al., 2017) and memory resources (Mallya & Lazebnik, 2018).
As a result, many creative models and algorithms have been proposed, including ones that cluster tasks in a top-down fashion (Lu et al., 2017), learn relationships between tasks using matrix priors (Long et al., 2017), and more (Ruder, 2017).
Our proposed model uses the mixture-of-experts (MoE) principle for multi-task learning.
Two related models are the cross-stitch network (Misra et al., 2016), where outputs from all experts are mixed together by a task-speciﬁc linear combination, and the routing network (Rosenbaum et al., 2017), where one expert per gating layer is selected, depending on the input and the task.
Routing and cross-stitch networks can be viewed as two extremes of a spectrum, with the former making gating decisions that are hard (yes / no) and sparse (exactly one expert), and the latter performing soft and dense “gating”.
Our method enjoys the best of both approaches, by learning how many experts to select for each task, as well as task-speciﬁc linear combinations that mix the selected experts.
This paper only studies networks with a single gating layer, but our method can be straightforwardly extended to multiple gating layers.
A similar approach is taken in Shazeer et al.
(2017) in the single-task setting, where the gating layer selects a predetermined number of experts.
In contrast, we learn a taskspeciﬁc number of experts in the multi-task setting, which requires a different mechanism for expert selection (Section 3).
During training time, our gating is a variant of dropout (Srivastava et al., 2014) in which entire experts are dropped out rather than single neurons, and the dropout probabilities are learned and task-speciﬁc rather than ﬁxed.
Another MoE model for the single-task setting is DeepMoE (Wang et al., 2018), which uses ReLU activation to encourage sparsity in input-dependent gating weights.
In contrast, we make discrete gating choices using a probabilistic model that is conditioned only on the task.
After training, we ﬁx the set of experts that contribute to each task, facilitating efﬁcient model deployment, and improving the model’s interpretability (Section 4.3).
3. Task Multi-Clustering
A task is a tuple x X , Y , P, L y , with X and Y , respectively, input and output domains, P a distribution over X ˆ Y , and L : Y ˆ Y Ñ R a loss function.
We say that a model m : X Ñ Y achieves good performance if it has low expected loss E p x,y q„ P r L p m p x q , y qs .
In multi-task learning (MTL), we assume some distribution Q over a set of tasks
T , and a dataset of tuples p t, x, y q , such that t „ Q and p x, y q | t „ P t .
In our experiments, Q is the uniform distribution over N tasks.
We propose a mixture-of-experts (MoE) model for MTL in which the tasks in T are clustered into M experts.
The binary gating variables b i p t q P t 0 , 1 u are functions of the task representation t indicating whether task t is clustered into expert i .
Intuitively, we would like a subset of the tasks
T i “ t t | b i p t q“ 1 u to be clustered into expert i if these tasks, and only these tasks, share useful features f i p x q of the input x .
The expert should learn to extract such features.
Given the expert featurizations f i p x q and the multiclustering b i p t q , we compute for each task t a task-speciﬁc representation z p t, x
q as a mixture of the experts into which it is clustered
[Equation]
where w i p t
q are the mixture weights.
We learn a task head g mapping from the latent representation z p t, x q to the output m p t
[Equation]
p
[Equation]
Our training objective is to minimize the expected loss between the predicted output and the true output min b
[Equation]
We note that, during both training and evaluation time, only the experts of the mixture (1) for which b i p t q “ 1 need to be called.
Sparsity of the multi-clustering can therefore yield signiﬁcant computational efﬁciency.
To perform gradient-based optimization of the objective (2), we can choose differentiable parametrizations of all w i ’s, f i ’s, and g as neural networks.
Unfortunately, the multi-clustering indicators b i are inherently discrete.
To relax this modeling assumption, we introduce the probability p
i
p t q that expert i will be used for task t , and have b i p t q „ Bernoulli p p i
p t qq .
This makes z p t, x ; b q and m p t, x ; b q random variables, requiring special care in the optimization.
Because the discrete gating variables are not differentiable, in order to optimize the gating distribution, we can either relax it into a reparametrizable distribution using the Gumbel– softmax trick (Jang et al., 2016), or use the score-function trick to sample from the same gating distribution that we optimize.
Sampling from a relaxed gating distribution generates soft gating values
b i p t q P p 0 , 1 q , and does not enjoy the computational beneﬁts of sparsity during training time.
We therefore choose the score-function trick
[Equation]
where p i p b i ; t q is p i p t
q if b i “ 1 , else 1 ´ p i p t q .
To estimate the gradient (3) in each optimizer step, we sample a single gating combination b for each element of the mini-batch.
We deﬁne the density of each task as the expected number of experts used in its computation, k p t q “ ř i p i p t q , and refer to its complement M ´ k p t q as the sparsity.
We desire sparser representations for faster training and evaluation, as long as task performance does not degrade much.
To our optimization objective, we add the expected task density, in our case the mean 1 N ř Nt “ 1 k p t q , as a regularization term, namely L 1 regularization on p .
The coefﬁcient of this term, λ , can be viewed as the relative marginal cost of computing one expert and incurring a unit of loss.
By tuning λ , we can control the trade-off between computation and loss.
During evaluation time, given an input x for task t , we wish to keep the same task density k
p t q as in training time.
Moreover, efﬁcient model deployment and computation may require a ﬁxed, deterministic multi-clustering b .
We therefore round k p t q to the nearest integer, and pick the set of k p t q most likely experts
[Equation]
Other approaches to choosing the ﬁxed multi-clustering b from the learned distribution p could be rounding the probabilities to 1 or 0, or sampling from them.
However, rounding would induce a different expert density at evaluation time than in training time, which may create a covariate shift for the downstream network, and sampling from the gating distribution may not select the most likely experts, thus degrading performance.
As a result of our gating discretization technique (4), the model computation time is O p k q , where regularization induces selection of only k experts.
This is generally much more efﬁcient than mixing all M experts in O p M q computation time.
As the diversity of the set of tasks T increases, we must also increase the total number of experts M , but not the task density k , contributing to the advantage of the multi-clustering method over naive MoE.
4. Experiments
We evaluate our method on the CIFAR-100 MTL dataset, in which each task is to classify an image into one ofclasses, with cross-entropy loss between the true class and the predicted class distribution.
There are N “ 20 tasks, corresponding to 20 non-overlapping super-classes (ﬁsh, ﬂowers, etc.).
Each image in the dataset is involved in exactly one task and has exactly one true label for that task.
The deep network architecture of each of the experts f i p x q
[Figure]
is four sequential convolutional layers, each with 32 ﬁlters,
ReLU activation, and 2 ˆ 2 max pooling.
Each task head g p t, z q is a fully connected layer that takes the mixture of experts z p t, x q and outputs classiﬁcation logits.
Two N ˆ M matrices represent w i p t q and the logits of p i p t q , with each element of w initialized by a standard Gaussian, and p initialized uniformly as 0 logits.
We optimize the mean task loss (2) using Adam with an initial learning rate of 10 ´ 3 .
We present three experiments: (1) a comparison between different gating techniques in terms of accuracy, (2) a comparison between different ways to select experts during evaluation time; and (3) a qualitative evaluation of the multiclustering similarity between tasks.
4.1. Comparison of Gating Techniques
To compare with cross-stitch networks (Misra et al., 2016), we generalize their method to have more than 2 tasks andexperts by representing the mixture weights as an N ˆ M matrix.
This is equivalent to setting b ” 1 in our model, and compares sparse gating to dense, ungated MoE.
The results are summarized in Figure 1.
With sufﬁcient data, multi-clustering outperforms ungated MoE with the same number of parameters.
While MoE is largely insensitive to the number of experts in a wide range (20–100), multi-clustering beneﬁts from sparser gating as the amount of data decreases, presumably due to a regularization effect.
Our results consistently outperform those reported in
Rosenbaum et al.
(2017): routing networks achieve 60% accuracy, whereas our model with λ “ 0 . 1 , using the same architecture, achieves 62.7% accuracy.
4.2. Expert Selection during Evaluation Time
We turn to the question of how task-speciﬁc subsets of experts should be selected during evaluation time, based on the learned distribution p .
We wish to ﬁx the multi-
[Figure]
clustering after training to deploy a ﬁxed model for each task with only the selected experts.
This has the potential to make the model more computationally efﬁcient and interpretable.
Prior works predeﬁned the number of experts per task as a hyper-parameter k .
In Shazeer et al.
(2017), the k experts with the largest (noisy) weights are selected.
In contrast, we choose the task-speciﬁc k p
t
q to be the expected number of experts selected during training time for task t .
We speculate that this choice of k p t q reduces the covariate shift in the distribution of z p t, x q experienced by the task head g p t, z q , thus improving its performance.
Given k p t q , we select the maximum-posterior subset of experts for the task, which is the k p t q most likely experts.
We compare this method of selecting the ﬁnal multiclustering to choosing k to be mean expected number of experts over all tasks.
The results, summarized in Figure 2, suggest that choosing task-speciﬁc k p t q outperforms choosing the mean k .
We speculate that the performance degradation is more pronounced in sparser multi-clustering because this increases the risk of having too few features for tasks that need more features than the average.
Figure 2 also compares the performance of ungated MoE as a function of the total number of experts M .
Performance is largely insensitive to the number of experts, if there are sufﬁciently many, but forcing all M experts to be mixed does degrade performance.
This suggests that tasks do beneﬁt from our discrete multi-clustering formulation.
4.3. Task Similarity
Using hard gating enables interpretable task relatedness involving how similarly tasks share experts.
In our evaluation of this property, after adding sparsity regularization with coefﬁcient λ “ 0 . 1 , we ﬁnd that some tasks have a strong preference for some experts, as depicted in Figure 3.
The left maps show the probabilities of the M experts being used for the N tasks.
The right maps display the normalized correlation in gating probabilities between pairs of tasks (rows in the left maps).
With the larger dataset (bottom), there is a clear correlation in gating probabilities between tasks, such as tasks 3 and 6, which are “food containers” and “household furniture”, respectively.
With the smaller dataset (top), such correlations between tasks are unclear.
5. Conclusion
In multi-task domains, tasks are often partially related, in the sense that they beneﬁt from sharing some but not all of their latent representation with other tasks.
We presented a method for discovering a multi-clustering of tasks into feature-extracting experts, such that each task uses a mixture of the experts it selected.
Through these experts, the task shares each portion of the parameters in its model with a different cluster of partially related tasks.
The discovered structure of the multi-clustering facilitates more efﬁcient learning, in terms of both task performance and computational efﬁciency.
The model’s sparsity allows evaluating only a subset of the experts per task, both during training and evaluation time.
By ﬁxing the multi-clustering after training, we enable deployment of smaller models per task, which only involve the most useful experts.
Finally, with sufﬁcient data, the model’s performance is robust to signiﬁcantly increasing the sparsity through regularization.
The single multi-clustering gating layer presented in this work can be extended to models with multiple such layers.
This extension may require techniques that reduce the considerable variance of the score-function gradient estimator (3).
Our method should further be evaluated on a variety of domains and architectures.
Acknowledgements
This work was performed at the Real-Time Intelligent Secure Exe- cution (RISE) Lab and Berkeley AI Research (BAIR).
In addition to NSF CISE Expeditions Award CCF-1730628, this research is supported by gifts from Alibaba, Amazon Web Services, Ant Financial, Arm, CapitalOne, Ericsson, Facebook, Google, Huawei, Intel, Microsoft, Nvidia, Scotiabank, Splunk and VMware.