Variational Multi-Task Learning with Gumbel-Softmax Priors
Authors: Jiayi Shen
Xiantong Zhen
Marcel Worring
Ling Shao
Venue: NeurIPS 2021 Poster (2021)
Keywords: Multi-task learning
Variational inference
Gumbel-softmax

Abstract: Multi-task learning aims to explore task relatedness to improve individual tasks, which is of particular significance in the challenging scenario that only limited data is available for each task. To tackle this challenge, we propose variational multi-task learning (VMTL), a general probabilistic inference framework for learning multiple related tasks. We cast multi-task learning as a variational Bayesian inference problem, in which task relatedness is explored in a unified manner by specifying priors. To incorporate shared knowledge into each task, we design the prior of a task to be a learnable mixture of the variational posteriors of other related tasks, which is learned by the Gumbel-Softmax technique. In contrast to previous methods, our VMTL can exploit task relatedness for both representations and classifiers in a principled way by jointly inferring their posteriors. This enables individual tasks to fully leverage inductive biases provided by related tasks, therefore improving the overall performance of all tasks. Experimental results demonstrate that the proposed VMTL is able to effectively tackle a variety of challenging multi-task learning settings with limited training data for both classification and regression. Our method consistently surpasses previous methods, including strong Bayesian approaches, and achieves state-of-the-art performance on five benchmark datasets.

Variational Multi-Task Learning with Gumbel-Softmax Priors
Jiayi Shen', Xiantong Zhen Marcel Worring's Ling Shao?
AIM Lab, University of Amsterdam, Netherlands
IInception Institute of Artificial Intelligence, Abu Dhabi, UAE
Abstract
Multi-task learning aims explore task relatedness to improve individual tasks, which of particular significance the challenging scenario that only limited data is available for each task. To tackle this challenge, we propose variational multi-task learning (VMTL), a general probabilistic inference framework for learning multiple related tasks. We cast multi-task learning as variational Bayesian inference problem, in which task relatedness is explored in a unified manner by specifying priors. To incorporate shared knowledge into each task, we design the prior of task to be a learnable mixture of the variational posteriors of other related tasks, which is learned by the Gumbel-Softmax technique. In contrast previous methods, our VMTL can exploit task relatedness for both representations and classifiers in a principled way by jointly inferring their posteriors. This enables individual tasks to fully leverage inductive biases provided by related tasks, therefore improving the overall performance of all tasks. Experimental results demonstrate that the proposed VMTL is able to effectively tackle variety of challenging multi-task learning settings with limited training data for both classification and regression. Our method consistently surpasses previous methods, including strong Bayesian approaches, and achieves state-of-ther oo-thh--rttttttt performance on five benchmark datasets.
Introduction
Multi-task learning 8 exploits knowledge shared among related tasks to improve their overall performance. This is especially significant when only limited data is available for each task. The central goal in multi-task learning is to explore task relatedness to improve each individual task [5],]] which is non-trivial since the underlying relationships among tasks can be complicated and highly nonlinear [52]. Early works deal with this by learning shared features, designing regularizers imposed
on parameters 33
24
23
21 or exploring priors over parameters 20 3 53 54] Recently,
multi-task deep neural networks have been developed, learning shared representations in the feature layers while keeping classifier layers independent
[4 learn small number of related tasks based on one shared input. However, lots of training data is usually available. Few-shot learning [12] addresses multi-task learning with limited data, but usually relies on large number of related tasks. Further, in general, most existing works explore either representations or classifiers for shared knowledge, leaving exploring them jointly an open problem.
31. Some methods
[32
In this work, we tackle the following challenging multi-tast task learning setting [33] each task contains very limited training data and only handful of related tasks are available to gain shared knowledge from. In addition, instead of sharing the same input, each of the multiple tasks has its own input space from different domain and these tasks are related by sharing the same target space, e.g., the same class labels in classification tasks. In this setting, it is difficult to learn proper model for each task independently without overfitting 33 52 is therefore crucial to leverage the inductive bias provided by related tasks learned simultaneously.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

To address this, we propose variational multi-task learning (VMTL), a novel variational Bayesian inference approach that can explore task relatedness in unified way. Specifically, as shown in Fig. [ we develop probabilistic latent variable model by treating both the feature repre- sentation and the classifier as latent variables. To explore task relatedness, we place condi- tional priors over the latent variables, which are dependent on data from other related tasks. By doing so, multi-task learning is cast as the joint variational inference feature representa- tions and classifiers for all tasks in single unified framework. The probabilistic framework enables the model to better capture the uncertainty caused by limited data in each task [13]
To leverage the knowledge provided by related tasks, we propose to specify the priors each task mixture of the variational posteriors of other tasks. In particular, the mixing weights constructed with the Humbel-Softmnn tech- nique [22] and jointly learned with the probabilis- tic modeling parameters by back-proye pockppropagation. The obtained Gumbel-Softmax priors enable the model effectively explore different correlation patterns among tasks and, more importantly, pro- vide unified way to explore shared knowledge jointly for representations and classifiers. We val- idate the effectiveness of the proposed VMTL by extensive evaluation on five multi-task learning benchmarks with limited data for both classifi- cation and regression. The results demonstrate the benefit of variational Bayesian inference for modeling multi-task learning. VMTL consis- tently achieves state-of-the ofrtthe-rrt performance and
T
DErr
Wt
N
Xt,n
Zt,
Yin
Fig. 1: A graphical illustration of the proposed model for multi-tast learning, (Xt,n= Vt..) is the n-th sample from task t. w, and Zt.x are the intro- duced latent variables, i.e., classifier/regressor and representation, respectively, whose priors are conditioned the data from other tasks, DI.TY-
surpasses previous methods all tasks.
2
Methodology
Consider a set of related tasks, each of which is classification or regression problem. Tasks in this paper share the same label space for classification or the same target space for regression but have different distributions. Each task has its own training data D, {x,,n Vt.n J where Ne the number of training samples. We consider the challenging setting where each task has limited labeled data, which makes it difficult to learn proper model for each task independently 333 52
In contrast to most previous works, we will explore the knowledge shared among tasks for learning both the classifirssreegressoss and representations of each task. Note that in this section we derive methodology mainly using terminologies related to classification tasks, but also applicable to regression tasks.
2.1 Variational Multi-Tast Learning
To better capture uncertainty caused by limited data, we explore multi-task learning in a probabilistic framework. We define the conditional predictive distribution: p()'e/X,, Di:rv)) with respect to the
current task D,: }/- = {yr,n) ) and X; ) where we use Di:rye denote the data from all tasks excluding 'D. of task f. From a probabilistic perspective, jointly learning multiple tasks amounts maximizing the conditional predictive log-likelihood follows:
1 T Σ log.('',,,
1
Σ
log p(yr.a Xy,,, DLTAA).
We condition the prediction for samples each task on the data of other tasks, D Di,T to leverage the shared knowledge from related tasks.
The probabilistic multi-task learning formalism in ① provides general framework to explore task re- latedness by conditioning on other tasks. More importantly. it enables the model to incorporate shared knowledge from related tasks into the learning of both the classifier and the feature representation in unified way by specifying their priors.

---

We introduce the latent variables w, and zt to represent the classifier and the latent representation of the sample Xt,n, respectively. The joint distribution for the task can be factorized as
No. II n=1
p()', Z,, we Wr, DI.TV))
p(y,, Ze,ns we) Pear))((we/PPnnd
(2)
where we use Z, = n J= collectively represent the set of latent representations of samples in task t. Here, we specify conditional priors dependent on other tasks to explore task relatedness. The probabilistic latent variables capture uncertainties different levels: Zt,n captures the uncertainty of each sample, while w, works at the categorical level. A graphical illustration of the probabilistic latent variable model shown in Fig.
Solving the model for multi-tase task learning involves inferring the true joint posterior p(2,, W. / (Dirr) over all latent variables 2, and wa, which is generally intractable. We introduce variational joint distribution q( 2 wr(D,) for the current task to approximate the true posterior. Also under the conditional independence assumption, the joint variational posterior distribution be factorized with respect to classifiers and latent representations as follows:
N₂
((,,,,, = go ( (D)) II (Ze. Xt,n),
m=1
where (Wr)Dr) and qo xt.n) ) are variational posterior distributions for classifiers and latent representations respectively, and 9 and ó are the associated parameters. For computational feasibility, they are defined as fully factorized Gaussians, as is common practice [2
To obtain good approximation, the variational posterior should be close to the true posterior. This is usually achieved by minimizing the Kullback-Leibler (KL) divergence between them:
DKL [g(,, w. (P.) welPi.r)].
By applying Bayes' rule to the true posterior, we derive the evidence lower-bound (ELBO) with the probabilistic classifier and representation for the conditional predictive log-likelithe lillliikeoodd
log () DI:T() 2 T
Rewerge -ag Re ~2 (log Doog((yr,n /ze,a wi)l
T
(5)
DD. (go (Zt,al Xt,nx) (xt.n))lp((za)) [p( Ne,n; Dirr DRI. [go (D:::D:))ll) (.r)D
The objective function provides general probabilistic inference framework for multi-task learning. As priors naturally serve as regularizations in Bayesian inference, they offer unified way of sharing information across multiple tasks for improving both classifiers and representations. The detailed derivation is given in the supplementary materials.
In what follows, will describe the specification of the prior distributions over the classifier and the representation, as well as their variational posterior distributions.
Learning Task Relatedness via Gumbel-Softmax Priors
The proposed variational multi-tass task inference framework enables related tasks to provide supportive information for the current task through the conditional priors of the latent variables. offers a unified way to incorporate shared knowledge from related tasks into the inference of representations and classifiers for individual tasks.
Classifiers To leverage the shared knowledge, we propose to specify the prior of the classifier for the current task using the variational posteriors over classifiers of other tasks:
p(w (1) (DITT))
Σ onego (w) (w (m-)) D:),
(6)
where y) denotes the 7-tth iteration. practice, avoid entanglement among tasks, we define the prior of task t in the current iteration to be combination of the variational posteriors [40] of the remaining

---

tasks from the last iteration. Particularly, our designed prior resembles the empirical Bayesian prior in that the prior of each task is specified by the observed data of other tasks, which provides a principled way to explore the shared knowledge among multiple tasks.
20 3 46 44
During training, we aim to have each task learn more shared knowledge from its most relevant tasks, while maximally reducing the interference from irrelevant tasks. To this end, we adopt the Gumbel-Softma technique learn the mixing weights for all related tasks and define the mixing weights as follows:
exp((log Tu gt))TT)
Σ Lift exp (log Tez + g()/TT))
Here Oni is the mixing weight that indicates the relatedness between tasks / and i. mii is the learnable parameter in the Gumbel-Softmax formulation, which denotes the probability of two tasks transferring positive knowledge. 9ti is sampled from a Gumbel distribution, using inverse transform sampling by drawing u ~ Uniform( 0, 1) and computing By using the Gunbel-Softmax technique, our model effectively handles possible negative transfer between tasks. The Gumbel-Soffmaal technique encourages the model to reduce the interference from the less relevant tasks by minimizing the corresponding mixing weights. The more negative the effects of interference between pairwise tasks are, the smaller the mixing weight is likely to be.
- -log(- log())). is the softmax temperature.
With respect to the inference of the variational posteriors, we define them as fully factorized Gaussians for each class independently:
qo (w = II go (we,e Dr.) = II N N diag(o?,.))
(8)
where Pt... and of can be directly learned with back-propagation [5]
As an alternative to direct inference, classifiers can also be inferred by the amortization technique 17. which will further reduce the computational cost. In particular, different classes share the inference network to generate the parameters of the specific classifier. In practice, the inference network takes the mean of the feature representations in each class input and returns the parameters 1t, and Tre for Wt.c. The amortized classifier inference enables the cost to be shared across classes, which reduces the overall cost. Therefore, it offers an effective way to handle scenarios with large number of object classes and can still produce competitive performance, as shown in our experiments, even in the presence of the amortization gap [10].
Representations In a similar way to (). we specify the prior over the latent representation Zt,n of sample Xt,n as a mixture of distributions conditioned on the data of every other task, as follows:
p( ziz Xt.n DI.Ty) = Σ Bugg i Xt... D.).
(n-1)
kipte
Here, (8 is the mixing weight which defined in similar way as 7 The conditional distribution (o(zi,n|Xt.n x Dx) on the right hand side of (9) indicates that we leverage the data D; from every other task i help the model infer the latent representation of xt,n The contribution of each other task is determined by the learnable mixing weight. In practice, the distribution qa (x,,, D:) is inferred
by an amortized network 15 28 To be more specific, the inference network takes an aggregated representation of Xt.. and D; as input and returns the parameters of distribution Jg. The aggregated representation is formulated as (T)) which is established by the cross attention mechanism [26). The sample Xt,n from the current task acts as a query, and Die plays the roles of the key and the value. Di. includes samples from the i-th task, which have the same label as Xt,a. This is formulated as:
Xt, DT.
fX D.) DotProduct(xx.. Di,es Di.)) == softmax
>i.e>
(10)
Vd
where f is the aggregation function, and Di. is matrix, with each row containing sample from class с of the i-th related task. is the dimension of the input feature. Since we are dealing with supervised learning in this work, class labels of training samples always available at training time. The intuition here is to find similar samples to help build the representation of the current sample.
The inference of the variational posterior qa (Zi,nl Xt... over latent representations is also achieved using the amortization technique [5] 28). The amortized inference network takes Xt,, as input and returns the statistics of its probabilistic latent representation.

---

2.3 Empirical Objective Function
By integrating multi-task learning with Gumbel-Softmax priors:
6 and (9) into 5.. we obtain the following empirical objective for variational
(0MTL(O, 3) T
log plylll/z w (())))]
+ DKL. [go xx,n))
Σ istt
BrigolZenn |x,, D.1]] + DKL. [40 (w./DD))l Σ angoo (w: (w.go(
(11)
where zt,n (z x) and w (m)) qo (wel (Dz). To sample from the variational posteriors, we
adopt the reparameterization trick 28 and M are the number of Monte Carlo samples for the
variational posteriors of latent representations and classifiers, respectively. In practice, L and M are set to 10, which yields good performance while being computationally efficient. We investigate the sensitivity of L and M in the supplementary materials. 0 represents the statistical parameters associated with the classifiers or the inference parameters for the amortized classifier; ø denotes parameters of the shared inference network for the latent representation.
We minimize this empirical objective function to optimize the model parameters jointly. The log-lik logllikelhhood term is implemented as the cross-entropy loss. In the practical implementation of the KL terms, we adopt the closed-form solution based on its upper bound as done in (35). e.g., DaL (g( wr)ll Eiet (w:g(w:)] < Σizz Dn; (((w))ll((w)) (w:)). Minimizing the KL terms encourages the model to leverage shared knowledge provided from related tasks at the instance level for representations and at the categorical level for classifiers.
At test time, we obtain the prediction for a test sample x, by calculating the following probability using Monte Carlo sampling:
I.
p(y)
Σ Σ
p(yle, we (m))
(12)
MT.
mm1
where we draw samples from posterions: x! (z(x))) and we (m) ~ go (w (D))
3
Related Works
Multi-Task Learning Settings There are mainly two typical multi-task learning settings: single- input multi-output (SIMO) and multi-input multi-output (MIMO), In the first setting different tasks are assumed to share the same input data during both training and testing [52]. Recently, this setting has been integrated into deep neural networks boost performance for multi-label prediction (38] 42
and various pixel-level prediction tasks, such as depth prediction, semantic segmentation and surface 50]. In the multi input multi-output setting, each task has its own input from a different domain and output B3 55 3 [T]. Tasks are related by sharing the same target space, e.g., class label space or regression target space. This setting, which the focus of our work, is made challenging by the distribution shift between different tasks.
normal prediction
25
Probabilistic Multi-Tane Learning This has been widely developed to explore shared priors for all tasks. The relationships among multiple tasks are investigated by designing priors over model I 3 under the Bayesian framework, or sharing the covariance structure of parameters [1I] In particular, Bakker et al. 3 adopted a Bayesian approach which introduces a joint learnable prior distribution that allows similar tasks to share model parameters and others to be more loosely connected. In contrast these works, our probabilistic modeling explores task relatedness by specifying priors under the variational Bayesian inference framework.
parameters [49
45
Regularization-basee Multi-Tass Learning This has been extensively investigated [52 with the aim of exploring different constraints, e.g., loc.1 and lp.g norms, imposed on model parameters or
23 21). In doing so, similar tasks are encouraged to share the same parameter patterns or the same subset of features. In particular, Long et al. [33] exploited
features as regularizers 3C
24
5

---

the task relatedness underlying parameter tensors to regularize the maximum posteriori estimation. This method was also developed for the multi-input multi-output setting. In our work, we explore shared knowledge for feature representations and classifiers/regressos jointly in single framework.
Gumbel-Softman Technique This technique 2 was recently introduced into multi- learning for exploring task relationships. In some multi-task learning models, it is used for searching the optimal architecture for each task by assigning kernels in each convolutional layer 77, choosing the child node during the forward pass [19], or controlling whether the sub-module is skipped [43] These works were exclusively for the single-input multi-output setting. In our work, we also adopt the Gumbel-Softmax technique for multi-task learning, using it to learn different correlation patterns for individual tasks. In this way, each task is encouraged to find the most relevant tasks for leveraging the shared knowledge, while reducing interference by irrelevant tasks.
4 Experiments
We conduct extensive experiments evaluate the proposed VMTL on five benchmark datasets for both classification and regression tasks. We perform comprehensive comparisons with previous methods and ablation studies gain insights into the effectiveness of our methods, We report the major experimental results here and provide more results in the supplementary materials.
4.1 Datasets
Office-Home
47 contains images from four domains/tasks: Artistic (A), Clipart (C), Product (P) and Real-world (R). Each task contains images from 65 object categories collected under office and home settings. There are about 15, 500 images total.
Office-Caltech [6 contains the ten categories shared between Office-31 339 and Caltech-256 [8] One task uses data from C5l (C), and the other three tasks use data from Office-31, whose images were collected from three distinct domains/tasks, namely Amazon (A), Webcam (W) and DSLR (D). There are 151 samples per category per task, and 2, 533 images in total.
ImageCLEs 331. the benchmark for the ImageCLEF domain adaptation challenge, contains 12 common categories shared by four public datasets/tasks: Caltech-256 (C), ImageNet ILSVRC 2012 (I), Pascal VOC 2012 (P), and Bing (B). There are 2, 400 images in total.
DomainNet 36 is a large-scale dataset with approximately 0.6 million images distributed among 345 categories. It contains six distinet domains: Clipart (C), Infograph (I), Painting (P), Quickdraw (Q), Real (R) and Sketch (S). Due to the large number of object categories, this is an extremely challenging benchmark for multi-task learning with limited data. To the best of our knowledge, has therefore not been used for multi-task learning before.
Rotated MNIST 29 adopted for angle regression tasks. Each digit denotes one task and the tasks are related because they share the same rotation angle space. Each image is rotated by 0" through 90% in intervals of 10", where the rotation angle is the regression target.
4.2 Experimental Settings
Implementation Details We adopt the standard evaluation protocols (52) for the Office-Home, Office-Caltedi and ImageCLEF datasets, randomly selecting 5%. 10%, and 20% of samples from each task in the dataset as the training set, using the remaining samples as the test set [33]. Note that when we use 5%. 10% and 20% of labeled data for training. there are on average about 3, 6 and 12 samples per category per task, respectively. For the large-scale DomainNer dataset, we set the splits to 1%, 2% and 4%, which also results an average of 3, and 12 samples per category per task, respectively. For the regression dataset, Rotated MNIST, we set the splits to 0.1% and 0.2%, which results in an average of 6 and 12 samples per task per angle, respectively. For fair comparison, all methods in our paper share the same split documents.

---

Office
Office-Caltech
ImageCLEF
75-
99
85
Average Accuraly(N)
75-
92
10
STL
VSTL EHTL
VATTLA
VMT
45:
5%
10%
20%
10%
22%
50%
60-
5%
10%
20%
40%
60%
split
spiit
split
Fig. The performance terms of average accuracy with 95% confidence interval under different proportions of training data on Office- Home, Office-Caltechh and ImageCLEF. Best viewed in color.
Table I: Effectiveness of variational Bayesian representations (z) and classifiers (w).
Table 2: Performance comparison of our Gumbel-
Softmax priors with other alternative priors.
5%
10%
20%
Priors
5%
10%
20%
50.4
59.5
65.620.1
Mean
57.51000
64, 410.1
68. 5.0.1
$7.4200
64.4+s1
68.5
Learnable weighted
57.2102
64.2.002
68 9102
x
56,8
63.7
68.5
58.3=a.1
65.0+a0
69.2=0.2
Gumbel-Softtax
58.3
65.0
69.2.
Following the experimental settings in [33, adopt the VGGnet 41), remove the final classifier
layer and employ the remaining model to pre-extract the input feature x. From the input, we infer its latent representation z using amortized inference by multi-layer perceptrons (MLPs) (28] In experiments, the temperature of the schedule applied in [22]: start with high temperature and gradually anneal it small but
priors (6) and (9) is annealed using the same
non-zero value. For the KL-divergence in [T], we use the annealing scheme from 6
We adopt the Adam optimizer 2 with learning rate of le-4 for training. All the results are obtained based on a 95% confidence interval from five runs. The detailed network architectures are given the supplementary materials. The code will be available at http:://tthhbbccm//uuuuu
Baselines Since the multi-task learning setting with limited training data is relatively new, most previous methods designed under other settings are not directly applicable due to the distribution shift between different tasks. To enable comprchensive comparison, implement several strong baselines including both Bayesian and non-Bayesian approaches. The single-task learning (STL) is implemented by task-specific feature extractors and classifiers without sharing knowledge among tasks. The basic multi-task learning (BMTL) is a deterministic model, which has a shared feature extractor and task-specific classifiers. Meanwhile, STL and BMTL, which we name VSTL and VBMTL, by placing priors as uninformative standard Gaussian distributions (details be found in the supplementary materials). VSTL and VBMTI. serve as baselines to directly demonstrate the benefits of the probabilistic modeling of multi-task learning based variational Bayesian inference. We VMTL-AC represent our proposed method with amortized classifiers. We compare against four representative methods [5 (8), which are implemented by following the same experimental setup as our methods. Details of the implementation can be found in supplementary materials.
define variational Bayesian extensions of
4.3 Results
We demonstrate the effectiveness of our VMTL in handling very limited data, the benefits of variational Bayesian representations and classifiers, and the effectiveness of Gumbel-Softtax priors in learning task relationships. Detailed information can be found in the supplementary materials.
Effectiveness in Handling Limited Data Our methods are highly effective at handling scenarios with very limited data. To show this, we compare our VMTL and VMTL-AC with other baselines. including single-task baselines (STL and VSTL) and multi-task baselines (BMTL and VBMTL) on the Office-Home Office-Caltech and ImageCLEF datasets. We employ a large range of train-test splits, from 5% to 60%. As shown in Fig. 2 our proposed VMTL outperforms the baselines on all three datasets with limited data (less than 20%) by large margins. The proposed VMTL-AC achieves comparable performance to VMTL on most benchmarks, despite the amortization gap.
7

---

(W)
Ro (4)
(M)
(9)
Head nver Proeast Cipatt Mt
fead nor Proeact Ciiert AR
Food neer Prosuct Cleart Art
Frnd nedkanee Prosuer Cleanr
Fig. Tasks show different correlation patterns 04 and 3 for classifiers (w) and representations (2).
Similar tasks are more correlated, e.g., World is closer to Art than C1 Clipart.
We observe that the proposed methods not only show dominant performance in the setting with limited data but also produce comparable even better performance than baselines in the setting with large amounts of data, which demonstrates the effectiveness of our methods in handling challenging scenarios with limited training data. The proposed methods show less significant improvements under setting with larger amounts of training data than the setting with limited data. The reason could be that the knowledge transfer becomes less crucial since individual tasks suffer less from over-fitting caused by limited data. Compared with other multi-task learning methods, the proposed methods still show certain improvements under the setting with large amounts of training data.
Effectiveness of Variational Bayesian Approximation We investigate the effect of variational Bayesian inference for representations and classifiers separately. We conduct these experiments on the Office-Home dataset. The results with different train-test splits shown in Table T which indicates whether z and w are probabilistic or not, e.g., a x for w means a deterministic classifier. More detailed experimental results provided in the supplementary materials. As be seen, both variational Bayesian representations and classifiers can benefit performance. This benefit becomes more significant when training data very limited, which indicates the effectiveness of leveraging shared knowledge by conditioning priors on related tasks. In addition, in Fig. 2 VSTL demonstrates be very strong Bayesian model, which again demonstrates the benefits of variational Bayesian representations and classifiers, even without explicitly leveraging shared knowledge among tasks.
Effectiveness of Gumbel-Softtmx Priors The introduced Gumbel- -ofobel-offtax priors provide an effective way learn data-driven task relationships. To demonstrate this, we compare with two alternatives without the Gumbel-Softmax technique in Table including the mean and the learnable weighted posteriors of other tasks. The proposed Gumbel-Softmn
priors perform the best, consis- tently surpassing other alternatives. We would like to highlight that the advantage of GumbellSofftnn priors is even larger under very limited training data, e.g., 5%, which creates challenging scenario where it is crucial to leverage shared knowledge provided by related tasks. To further investigate the learned Gumbel-Softmax priors, we visualize the values the mixing weights a and for both Office-Home with the 5% split are shown in Fig. 3 We can
VMTL and VMTL-AC. The results see that the relatedness among tasks is relatively sparse, which could avoid interference by irrelevant tasks. It is worth noting that the mixing weight between tasks is higher when they more correlated. For instance, the task Real World is closer to Art than to Clipart.
Comparison with Other Methods The comparison results on the small-scale Office-Home, Office- Caltech, ImageCLEF datasets and large-scale DomainNet dataset shown in Tables 3 respectively. The average accuracy all tasks is used for overall performance measurement. The best results are marked in bold, while the second-best are marked by underlining.
5
and 6
Our proposed methods achieve the best performance on all four datasets with different train-test splits. On the most challenging setting with 5% training data, VMTL shows large performance advantage over the baselines. This demonstrates the effectiveness of VMTL in exploring relatedness to improve the performance of each task from limited data. VMTL-AC can also produce comparable performance and is better than most previous methods.
We would like to highlight that Bakker et al. [3, Bayesian method for multi-task learning optimized by an 33 use point estimation to find approximate values of model parameters. However, methods
maximization algorithm, produce very competitive performance. Long et al.

---

train spil: 0.1%
train-test split: 0.2%
DATL
112
.09
1955
04
09
.00
10p
last
Fig. 4: Performance comparison (normalized mean squared errors) for rotation angle regression.
Table 3: Performance comparison of different methods on the Office-Home dataset.
Methods
5%
A0g
10%
20%
Ang
SIL VSTL
30/7284
67.5
61.7333
09:2:35
50.00 8,03
18.4
07.5
58.3.00
54..334
50,6
73 1ub
Maxi
37.
33.3
69.0em
64.0ss
S1.1-41
52.0
43.s
76.2as
07.4
60
55.9
52.0eas
$1.3ass
73.8
Buiker
40/0.:1
355.00
698
52.8.11
$2.5.3
42.3
75.7155
095.c5
00.0.02
01.3
56.5
$3,7100
75.4.02
08.7.00
Lings al 11
47.8.aa
37.9um
70.4ess
53.
57.2
43.3mm
7677as
63.4mm
65.1exx
46.7em:
790
76.6
Kendall et 65
40:2:8
33.0m4
6M5532
60.7.31
5133311
49
49.100
39.7
134112
07.4153
57.2.02
59.5233
53.8
M0.100
73.6.04
06.802
Chan al 18
37.9
31.4m0
67.7
50.4xas
49.
47.1am
37.
30.3esc
66.aaa
35.3.m:
58.3waz
$3.5.a
79.5
73.
66.2.u
BNTIL.
37.6244
31.5.03
885533
0338833
90:4:11
11.0
410
09:3:
59.5.:
46.033
51.8:a5
73.9-04
VEMTL
38.6.85
318.ma
64.5
51.3
43.
76.
60.9.:
59.2.k
$3.3444
91.9m2
VMTT
39.2
475
53.5.a
Table 4: Performance comparison of different methods on the Office-Caltech dataset.
Mothods
Avg
A0g.
SIL VSTT.
N7:9188
96.4..5
97.7.00
K7.8:3
64.3.14
90.7.62
94.9
02 11.
U5.2 u5.2
867.0o
92.4.00
88.3
89.1asa
07.0ea3
81.4ps
59 19.0
93.1um
96.5am
50 00:0
64.5
91.1-m3
05.5eM
94.5em:
95.0sx
$7.7
43.
Rackem a al
03:2as
94.0es
94.7
$5.4
91 9188
94. 9cm
97. 97.6cm
06/6
900
05.0cm>
95.2pm
94-4es:
90.5ae
91.S
Long et al.
82.7
84.3183
97:1-62
39.2.08
934131
95.0.65
98.1
950012
91.3332
94.8.65
955.33
849111
992143
91.0.04
Kendall es
03.6.a
92.5aa0
95.0um
$3.9mm
21.2aas
94.0um
98.2au
03.
00:4as2
03.8.a2
03:3am:
993aea
91:2am
05.7.04
852.00
91.1
97.0ka
99.2.45
91:1:04
BMTL
93.1
VISMT
90.3
M0.3
97.1253
34.4
90/
94.2.01
97.lan
87.8
93.2.93
57.8p4
VMTL-ACC
93.5
05/4443
967743
90.0.03
93.9
94.9.a3
97.1
97.930
906-13
05.1.43
95.5.a
96.464-
95.6cc3
90.3-04
95.2
VMTT.
09.7
95.7
96.4
89 7p94
918
95.4gn
97 97.0gm
97 4g34 975
05.3
95.629
98.4
91.1gn
98.2
Table 5: Performance comparison of different methods on the ImageCLEF dataset.
Mathods
Aug.
STL VSTL
K5-4aan
71-4es
37.7
Det
62.6
S6. Duei
77.8
56 Bass
47 67.6
69:7cm
92.9psm 93.30X3
73.5ax
54.6
T6 2
87.01a
73.2144
60.5544
39.0.03
64.9
89.5.03
19.1
56.0133
48.02.1
70.8.45
$7.3143
72.7ue5
55.4.04
17.2
90.11s
76.515
72.8.03
54.9004
13.7
55.2
56.3.3
89.2
59.4.03
Kendali
952aan
86.
66.6e
50.4
14hest
91.9cm:
38 38.9mo
14.3am>
$0.4
16.9-a
98.3em
91.0ea:
75.6
56.9
79.2
Qun et 35
91.013
85.5144
65.4.03
502.04
74
90.7.04
18.1
75.6134
546.35
77.3.65
93.1.33
02.1145
74.41C7
55.8.08
12.9.0
BMTT.
88.3aa
73.3
6122c0
40.0es
65.
637aas
90.8-a
79.3mm
56.3e3
S1.
72.0am
92.9
86 86-3em
71.0ax
36.0maa
768800
VIMTI
75.0185
01:9.33
417711
67.
91.1.82
S017
000
5111::5
13.--*
997.33
87.116
567003
VMTI--AC
89.3am
81.5ass
72.3
39.5
757
35.9a
71.9.u
99.
77.6-42
92.9.3
88.0
75.6
60.5-02
91.3
71.6u
59.5pn4
Table 6: Performance comparison of different methods on the scaled DomainNer dataset.
Methods
18
Aag.
25
0
Asg.
STL VSTE.
1500
4.0.a4
19.7.03
7,5182
50.6.n2
9.6324
17.7at2
20.5
6 6:22
26.0.n
7,2182
$6.5
14.0
21 2138
189und
5.4553
23.5.m2
15.2
54.5
12.2
21.8a83
26.4291
8.9 8.9
30.9:n
20.2
60.7e0
17.8
27.5
Bakber al
20:2:n
5.8153
26.2.01
21.0s8:
5288n2
14.2a11
23.5133
23.4.a1
8.6281
32.5.m
21.5s33
58.7
19.2.s2
21 2153
Long at 130
19.1ad
K1sta
30.7.z
575m1
57.
15.1am
2.8r53
24.3ax5
11.0a52
36.--u
7.0202
60.gp:
18.0s13
26 3633
Kandill al 55
175gu
4.5r
34.4-0
19.4281
507mm
11.521
21.3at1
23.5
77mm
32.7-m
22.5aa2
6022m
18. 18.5
Gin al.
20
642m
272
19.6es:
5277mm
13.5
23. 333
36.0g3
89pm
34.5-m
20.0
62.6
21 21/
20.0az1
BMTL.
1833g
52500
22.3_a1
15:0es:
53.3mm
11.5p11
21.0aaa
26-2337
8720
30
20.7
59.640
17.8
VBMTL.
181
S2aca
25.0m4
13,4
51.0p0
12.0
20.0
26.2435
89.0
30770s
17.54a2
59.0.4
17.5
26.8s22
VMELAC
193-ni
66463
28.1-04
12.1
52.3.m:
13,621
21.4aa
26.5aa2
97.01
306-02
14.6aa1
$5.2.0
18.3aa2
26.3120
VMEL
252-m1
884ca
30.1.01
13,2483
36.4003
17.5
25.2.41
31.1aa1
11.9
35.1-a1
17.1
61.6sez
21.9
29.8isa
consistently outperform Bakker et al. [3] and Long et al. (33). which shows the great effectiveness of variational Bayesian inference for multi-task learning. In addition, the methods (25) and (38] weigh multiple loss functions by considering the homoscedastic uncertainty of each task, which shows the benefit of modeling uncertainty, while they are also outperformed by our methods.
It also worth mentioning that VMTL-AC demonstrates computational advantages as well, with faster convergence than VMTL due to amortized learning. Besides, we find that VMTL-AC demonstrates good robustness against adversarial attacks. This could be due to the fact that amortized learning applies the mean feature representations to generate classifiers, which is more robust to attacks. Due to space limit, more experiments and detailed discussions about the amortized classifier are provided in the supplementary materials.
9

---

In addition, we apply our methods to regression task on Rotated MNIST for rotation angle estimation. The comparison results are shown in Fig. 4 The proposed VMTL and VMTL-AC outperform other methods with lower average NMSE. For some digits/tasks, e.g. 0, 6, 9, their NMSE are higher than other digits/tasks. The reason could be that these digits have some degree of rotational invariance, which makes it harder for them to regress their target angles.
5 Conclusion
In this paper, we develop variational multi-task learning (VMTL), a novel variational Bayesian inference approach for simultaneously learning multiple tasks. We propose a probabilistic latent variable model to explore task relatedness, where we introduce stochastic variables as the task-specific classifiers and latent representations in a single Bayesian framework. In doing so, multi-task learning is cast variational inference problem, which provides unified way to jointly explore knowledge shared among representations and classifiers by specifying priors. We further introduce Gumbel- Softmax priors, which offer an effective way to learn the task relatedness in data-driven manner for each task. We evaluate VMTL on five benchmark datasets of multi-task learning under both classification and regression. Results demonstrate the effectiveness of our variational multi-task learning and our method delivers state-of-the-art performance.
Acknowledgment
This work is financially supported by the Inception Institute of Artificial Intelligence, the University of Amsterdam and the allowance Top consortia for Knowledge and Innovation (TKIs) from the Netherlands Ministry of Economic Affairs and Climate Policy.
References
[1] C. Archambeau, S. Guo, and O. Zoeter. Sparse bayesian multi-task learning. In NIPS, volume 1, page 41,
2011.
[2] A. Argyriou, T. Evgeniou, and M. Pontil. Multi-task feature learning. In Advances in neural information processing systems, pages 41-48, 2007.
[3] B. Bakker and T. Heskes. Task clustering and gating for bayesian multitask learning. Journal of Machine Learning Research, 4(May):83 99, 2003.
[4] J. Baxter: A model of inductive bias learning. Journal of artificial intelligence revearch, 12:149-198, 2000.
[5] C. Blundell, J. Cornebise, K. Kavukcuoglu, and D. Wierstra. Weight uncertainty in neural network. In International Conference Machine Learning, pages 1613-1622. PMLR, 2015.
[6] S. R. Bowman, L. Vilnis, O. Vinyals, A. M. Dai, R. Jozefowicz, and S. Bengio. Generating sentences from continuous space. arXiv preprint wwX.v:1511.06331.
2015.
[7] F. J. Bragman, R. Tanno, S. Ourselin, D. C. Alexander, and J. Cardoso. Stochastic filter groups for multi-task cnns: Learning specialist and generalist convolution kernels. In Proceedings of the IEEE
International Conference on Computer Vision, pages 1385-1394 2019.
[8] R. Caruana. Multitask learning. Machine learning, 28(1):41 11:::11175 75,, 1997.
[9] Z. Chen, Badrinarayanan, C.-Y. Lee. and A. Rabinovich. Gradnorm: Gradient normalization for adaptive
loss balancing in deep multitask networks. In International Conference on Machine Learning, pages 794-803. PMLR, 2018.
[10] C. Cremer,
and D. Duvenaud, Inference suboptimality variational autoencoders. arXiv preprint vrXiv:1801.03555 2018.
[11] H. Daumé III. Bayesian multitask learning with latent hierarchies. arXiv preprint arXiv:0907.0783, 2009.
[12] C. Finn, Abbeel, and Levine. Model-agnostic meta-learning for fast adaptation of deep networks.
arXiv preprint (XXiv:7703.0444 2017.
[13] C. Finn, K. Xu, and S. Levine. Probabilistic model-agnostic meta-learning, In Proceedings of the 32nd
International Conference on Neural Information Processing Systems, pages 9537-9548, 2018.
[14] Y. Gao, Bai, 7. Jie, Ma, K. Jia, and W. Liu. Mtl-nas: Task-agnostic neural architecture search towards gencral-purpose multi-task learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11 11543-11552. 2020.
[15] S. Gershman and N. Goodman. Amortized inference in probabilistic reasoning. In Proceedings of the annual meeting of the cognitive science society, volume 36, 2014.
10

---

[16] B. Gong. Y. Shi, F. Sha. and K. Grauman. Geodesic flow kernel for unsupervised domain adaptation. In
2012 IEEE Conference on Computer Vision and Pattern Recognition, pages 2066-2073, IEEE, 2012.
[17] J. Gordon, J. Bronskill, M. Bauer, S. Nowozin, and R. Turner. Meta-leaming prohabilistic prrtellaminggrroohllliiitt inference for prediction. arXiv preprint arXiv:
2018.
[18] G. Griftin, A. Holub, and P. Perona. Caltech-256 object category dataset. 2007.
[19] P. Gua, C.-Y. Lee, and D. Ulbricht. Learning to branch for multi-task learning. In International Conference
Machine Learning, pages 3854-3863, PMLR, 2020.
[20] T. Heskes. Empirieal bayes for learning to learn. 2000.
[21] A. Jalali, S. Sanghavi, C. Ruan, and P. K. Ravikumar, A dirty model for multi-task learning. In Advances
in neural information processing systems, pages 964-972, 2010.
[22] E. Jang, S. Gu, and B. Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint 1rXiv:1611.00 144, 2016.
[23] P. K. Jawanpuria, M. Lapin, M. Hein, and B. Schiele. Efficient output kernel learning for multiple tasks. In Advances in neural information processing systems, pages 1189-119 1197, 2015.
[24] Z. Kang. K. Grauman, and F. Sha. Learning with whom share in pulli-task feature learning. In ICML,
volume 2. page 4, 2011.
[25] A. Kendall, Y. Gal, and R. Cipolla. Multi-task learning using uncertainty weigh losses for scene geometry and semantics. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages
7482-7491, 2018.
[26] H. Kim, A. Mnih, Schwarz, M. Garnelo, A. Eslami, D. Rosenbaum, O. Vinyals, and W. Teh. Attentive neural processes. arXiv preprint wxXiv/9901.07700 2019.
[27] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint Xri:::4422.600
2014.
[28] D. P. Kingma and M. Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.66 14, 2013.
[29] Y. LeCun, I.. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applicd to document recognition. Proceedings of the IEEE, 86(1 1):22)8-2222244 1998.
[30] H. Liu, M. Palatucci, and J. Zhang. Blockwise coordinate descent procedures for the multi task lasso, with applications to neural semantic basis discovery. In Proceedings of the 26th Annual International Conference on Machine Learning, pages 649-6566 656, 2009.
[31] L. Liu, Y. Li, Z. Kuang, J.-H. Xue, Y. Chen, W. Yang. Q. Liao, and W. Zhang. Towards impartial multi-task learning. In International Conference Learning Representations, 2021.
[32] S. Liu, E. Johns, and A. J. Davison. End-to-end multi-task learning with attention. In Proceedings of the IEEE Conference on Computer Vision and Pottern Recognition, pages 1871-1880, 2019.
[33] M. Long, Z. Cao, J. Wang, and S. Y. Philip. Learning multiple tasks with multilinear relationship networks. In Advances in neural information processing systems, pages 1594-1603, 2017.
[34] I. Misra, A. Shrivastava, A. Gupta, and M. Hebert. Cross-stitch networks for multi-task learning. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition, pages 4003,4003. 2016.
[35] E. Nalisnick and P. Smyth. Learning priors for invariance. In International Conference on Artificial Intelligence and Statistics, pages 366-375. PMLR, 2018.
[36] X. Peng. Q Bai, X. Xia, Z. Huang, K. Saenko, and B. Wang. Moment matching for multi-source domain adaptation. In Proceedings of the IEEE International Conference on Computer Vision, pages 1406-1415, 2019.
[37] T. K. Pong, P. Tseng, S. Ji, and J. Yc. Trace norm regularization: Reformulations, algorithms, and multi-task learning. SIAM Journal on Optimization, 2066):346--344 2010.
[38] W. Qian, B. Chen, and F. Gechter. Multi-task variational information bottleneck. arXiv preprint
arXiv:2007.00339, 2020.
[39] K. Saenko, B. Kulis, M. Fritz, and T. Darrell. Adapting visual category models new domains. In European conference on computer vision, pages 213-226. Springer, 2010.
[40] Y. Shi, N. Siddharth, B. Paige, and P. H. Torr. Variational mixture-of of----erttt----eeee autoencoders for mulli-mode deep generative models. Advances in Neural Information Processing Systems, 2019.
[41] K. Simonyan and A. Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint 1xXiv:1409.1566 2014.
[42] G. Strezoski, N. Noord, and M. Worring, Many task learning with task routing. In Proceedings of the IEEE International Conference on Computer Vision, pages 1375-1384, 2019.
[43] X. Sun. R. Panda. R. Feris, and K. Saenko. Adashare: Learning what to share for efficient deep multi-task learning. arXiv preprint arXiv: 1.143,,,,, 2019.
[44] H. Takahashi, T. Iwata, Y. Yamanaka, M. Yamada, and Yagi. Variational autoencoder with implicit optimal priors. In Proceedings of the AAAI Conference on Artificial Intelligence, pages 5066 (666,507,,20 2019.
11

---

[45] M. K. Titsias and M. Lázaro-Gredille Spike and slab variational inference for multi-task and multiple kernel learning. In Advances in neural information processing systems, pages 2339-2347,, 2011.
[46] J. Tomczak and M. Welling. Vae with a vampprior. In International Conference Artificial Intelligence and Statistics,, pages 1214-1223, 2018.
[47] H. Venkateswara. J. Eusebio. S. Chakraborty, and S. Panchanathan. Deep hashing network for unsupervised domain adaptation. (IEEE) Conference Computer Vision and Puttern Recognition (CVPR). 2017. [48] Y. Xue, X. Liao, L. Carin, and B. Krishnapuram. Multi-task learning for classification with dirichlet process priors. Journal of Machine Learning Research, 8(Jan):35 63, 2007.
[49] K. Yu, V. Tresp, and A. Schwaighofer. Learning gaussian processes from multiple tasks. Proceedings the 22nd international conference on Machine learning, pages 1012-1019, 2005.
[50] T. Yu, Kumar, A. Gupta, Levine, Hausman, and C. Finn. Gradient surgery for multi-tasi task learning.
arviv preprint arXiv:2001.06756 2020.
[51] Y. Zhang and Yang. Learning sparse task relations in multi-task learning. Proceedings of the AAAI Conference on Artificial Intelligence, 2017.
[52] Zhang and Q. Yang. A survey on multi-task learning. IEEE Transactions on Knowledge and Data Engineering, 2021.
[53] Y. Zhang and D. Yeung, Multi-task learning using generalized t process. In Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, pages 964-971, 2010.
[54] Y. Zhang and D.-Y. Yeung. A convex formulation for learning task relationships in multi-task learning. arXiv preprint 12X.v:1003.3536 2012.
[55] Y. Zhang, Y. Zhang, and W. Wang. Deep multi-task learning via generalized tensor trace norm, arXiv preprint arX:: 2020.
12


Reviews:
[{"review_id": "pE2CEzn7rhj", "reviewer": "Reviewer_k1co", "summary": "This paper proposes a variational multi-task learning method that uses Gumbel-Softmax priors for both the weight (w) of each task and the latent representation (z) of each data point. The Gumbel-Softmax trick is used for learning the relatedness between tasks. Amortized inference is also used for faster inference of the posteriors of the weight vectors and also for the latent representations. Experiments on different benchmark datasets show improvement over baselines especially in the low-data regime.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is in general well-organized and easy to follow. The notations are clear and consistent throughout. Introducing a learnable Gumbel-Softmax prior for learning from relevant tasks is original and interesting. I like the extensive ablation studies on handling limited amounts of data, using gumbel-softmax prior, and amortization. Also the paper considers multiple baselines (both Bayesian and non-Bayesian) for comparison. In particular, the proposed method is able to perform relatively well with limited data.\n\nMy major remaining concern about the paper is about the use of different Gumbel-Softmax weights for the weight (w) and latent representation (z). Given that the weight and latent representations are coupled, maybe it makes more sense to have them to be the same? \n\nI am also curious if there is any other suitable metric for comparing the uncertainty prediction of the Bayesian MTL methods. Currently it seems only accuracy or NSME is used as a metric for evaluation.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes. Maybe can also warn about possible negative transfer if tasks are not very related.", "review_text": "", "extra_scores": {"ethical_concerns": "N.A.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5 hours"}, "extra_text": {}}, {"review_id": "_eDciFZN4t", "reviewer": "Reviewer_Q1H6", "summary": "In this work, the authors propose to use Gumbel-Softmax as the prior of a task for multi-task learning under the variational Bayesian inference setup. Instead of exploiting task relatedness for representations only, the proposed method jointly infers posteriors for both representations and classifiers. State-of-the-art results are achieved in the multi-input multi-output setting with a small training data regime. ", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n* Thorough literature review in terms of deterministic and probabilistic approaches.\n* Leverage the advantages of probabilistic-based models and achieve good results on the small data regime.\n* Various ablation studies demonstrate the effectiveness of Gumbel-Softmax in MTL environments and effectiveness of the proposed variational Bayesian representation z and classifier w.\n* Although Gumbel-Softmax has been widely used in multi-task learning [7,19,42], to the best of my knowledge it is the first time being applied to the variational multi-task learning setup.\n\nWeaknesses:\n* It would be great to discuss the actual training and inference time of the proposed method compared to deterministic approaches to see the runtime impact of the additional sampling steps.\n* The benefit of a probabilistic-based approach on small data regimes is clearly demonstrated in the experiments, where 3-12 samples per category per task are provided. Can the authors shed some light on the scenario when training with large amounts of data? Is varitionable-based approaches still favorable in that setting?\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The main limitation I can think of is the performance of the proposed variational Bayesian approach in the large data regime as mentioned above. I was wondering if there is a way to combine the advantages of the proposed approach with deterministic approaches so that it can generalize well in more settings?", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "5DKtGEsSUXv", "reviewer": "Reviewer_xbXZ", "summary": "This work proposes a probabilistic method (variational inference) to exploit task relatedness and ignore harmful interference in multi-task learning paradigms. In particular, the authors focus on a unique MTL setting (training data is extremely limited; tasks have different input spaces; tasks have the same target space) to highlight the utility of their method. Experimental results on 5 datasets and compared to 6 related work is highly compelling to establishing the efficacy of the proposed method.", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This work is grounded in the probabilistic multi-task learning domain but presents a novel approach: using variational inference to model task relatedness. This method appears highly effective when training data is extremely limited as it augments each task with a strong inductive bias by learning correlation patterns among tasks (Figure 3). Moreover, these correlation patterns appear to be sparse (good) which the authors attribute to using a Gumbel-Softmax to determine task relatedness. \n\nThe submission is clearly written and cogent. The authors motivate their proposed method, outline the situation in which it is useful, offer intuition and explanation in defining methodology and comprehensively present compelling experimental results which adhere to their introduction. I believe the approach is technically sound; although, I did not carefully check for conflicting assumptions or unsupported jumps in the methodology. The extensive empirical evaluation significantly bolsters the author's proposed method.\n\nThat said, the work addresses a very limited use case of MTL, and I have not encountered a use case -- and have difficulty imagining a real world scenario -- where this technique would be used. The datasets used in the empirical evaluation are quite unusual: \"small\" datasets with an even smaller 0.1% -> 20% of the available training data being used for the method to be effective. Moreover, I would posit a transfer learning algorithm or meta learning algorithm would be significantly more suitable to this paradigm than a multi-task learning approach. Additionally, the method is quite involved and complex which reduces its utility and adoption. With regards to presenting novel research concepts, this work suggests their contribution is jointly exploring learned representations and classifiers; however, I would suggest the domain of soft-parameter sharing and even neural architecture search in MTL systems presently address this domain, but with significantly more utility on a variety of MTL applications. Integrating variational inference into multi-task learning is novel, but again, the extensibility of this contribution and application to MTL systems seems extremely limited.  \n\nIn summary, this is a very high quality and well-written work, but I struggle to see the significance of the proposed method to either real world systems or as a building block in future work. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I cannot recall this work containing meaningful discussion of its limitations (aside from the limited scope of its applicability) and societal impact is not addressed. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "3MbIfS3t8zA", "reviewer": "Reviewer_TXRx", "summary": "This paper proposes a unified few-shot multi-task learning framework which learns latent representations of data points and latent classifier weights through an amortized variational inference. Specifically, the proposed method regularizes the inference networks by minimizing KL divergence between the target variational distribution and the mixture of the non-target variational distribution. The mixing probability of the non-target variational distribution is sampled through Gumbel-Softmax distribution with learnable parameters. The intuition behind this regularization is that a model is able to learn from related tasks to improve the inference of latent representations and classifier weights (especially useful for the few-shot multi-task learning setting). The extensive experiments demonstrates that the proposed method is more effective than other methods for the few-shot multi-task learning.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- Originality\n\nThis paper tackles the uncertainty of few-shot multi-task learning through a unified probabilistic inference of representation and classifier weights, which is novel in the literature of few-shot multi-task learning. The difference from the related work is also well-explained in Section 3. \n\n- Quality\n\nThe technical quality of this paper is intuitive and well-demonstrated, however, there is some missing part for technical details, such as, KL regularization (see the below questions).  The experimental results are very impressive and coherent to the motivation and design purpose, for example, the proposed method is most useful for the most limited data setting.\n\n- Clarity\n\nThis paper is well-written, but there are equations that can lead misunderstandings.\n\n---\nI appreciate the authors to address my concern, and I will raise my score.\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The below is my and questions and suggestions to improve this paper.\n\n- Although the idea of using an aggregated (or mixture of) variational posterior as a prior distribution is new for multi-task learning, it is already introduced in the literature of improving Variational Autoencoder (https://arxiv.org/pdf/1809.05284.pdf). It would be better if any discussion of this paper is included in the paper.\n\n- This paper regularizes a model to use the data of not-target tasks by minimizing KL divergence between the variational distribution of target tasks and the mixture of the non-target distributions. I am not enthusiastic for the theory, but as far as I know any analytic solution of KL divergence between Gaussian distribution and a mixture of Gaussian distributions does not exists, and only bound exists. Can you elaborate how to compute KL divergence in your framework?\n\n- The function f is introduced in equation 10, however, its usage is not explained in relation to equation 9. Is this function q_\\phi in equation 9?\n\n- I didn't understand how to select the class c in equation 10. Is this class is the same class of x_{t, n}? if it is, how it work when we do not know the label of x_{t, n} as equation 9 and 10?\n\n- I have checked the implementation details of supplementary files to understand the above question. As I read, the cross attention f in equation 10 is used for the prior distribution of the latent representation, however, the cross attention layer appears only for the inference network \\theta of classifier weights (in Table B.3). This is very different as I read in the paper.\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers all found the approach to be “original and interesting,” with their main concerns addressed in the discussion. The technique looks to be very useful in the multi-task setting with limited data. During the discussion period, additional experiments showed that it doesn’t hurt in the large-data setting either, although the improvement is less pronounced. To me, this shows that the Bayesian approach seems to be doing the right thing: not getting in the way of single-task performance when enough data is available, but leveraging multiple tasks when data is scarce. One reviewer has concerns that the approach is too complex to be widely adopted, but given that code was made available as part of the submission, I am less concerned about this.
Decision: Accept (Poster)