Variational Multi-Task Learning with Gumbel-Softmax Priors
Year: 2022

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.

heading: 1 Introduction
text: 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 [2, 54], 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 [37, 24, 23, 21] or exploring priors over parameters [20, 3, 48, 53, 54]. Recently, multi-task deep neural networks have been developed, learning shared representations in the feature layers while keeping classifier layers independent [33, 38, 55, 31]. Some methods [32, 25, 9, 14] learn a 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 a 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.
In this work, we tackle the following challenging multi-task learning setting [33]: each task contains very limited training data and only a 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 a 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 a proper model for each task independently without overfitting [33, 52]. It is therefore crucial to leverage the inductive bias [4] provided by related tasks learned simultaneously.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
To address this, we propose variational multi-task learning (VMTL), a novel variational Bayesian inference approach that can explore task relatedness in a unified way. Specifically, as shown in Fig. 1, we develop a probabilistic latent variable model by treating both the feature representation and the classifier as latent variables. To explore task relatedness, we place conditional 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 of feature representations and classifiers for all tasks in a 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 of each task as a mixture of the variational posteriors of other tasks. In particular, the mixing weights are constructed with the Gumbel-Softmax technique [22] and jointly learned with the probabilistic modeling parameters by back-propagation. The obtained Gumbel-Softmax priors enable the model to effectively explore different correlation patterns among tasks and, more importantly, provide a unified way to explore shared knowledge jointly for representations and classifiers. We validate the effectiveness of the proposed VMTL by extensive evaluation on five multi-task learning benchmarks with limited data for both classification and regression. The results demonstrate the benefit of variational Bayesian inference for modeling multi-task learning. VMTL consistently achieves state-of-the-art performance and surpasses previous methods on all tasks.

heading: 2 Methodology
text: Consider a set of T related tasks, each of which is a 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 t has its own training data Dt = {xt,n,yt,n}Ntn=1, where Nt is the number of training samples. We consider the challenging setting where each task has limited labeled data, which makes it difficult to learn a proper model for each task independently [33, 52]. In contrast to most previous works, we will explore the knowledge shared among tasks for learning both the classifiers/regressors and representations of each task. Note that in this section we derive our methodology mainly using terminologies related to classification tasks, but it is also applicable to regression tasks.

heading: 2.1 Variational Multi-Task Learning
text: To better capture uncertainty caused by limited data, we explore multi-task learning in a probabilistic framework. We define the conditional predictive distribution: p(Yt|Xt,D1:T\t) with respect to the current task Dt: Yt = {yt,n}Ntn=1 and Xt = {xt,n} Nt n=1, where we use D1:T\t to denote the data from all tasks excluding Dt of task t. From a probabilistic perspective, jointly learning multiple tasks amounts to maximizing the conditional predictive log-likelihood as follows:
1
T T∑ t=1 log p(Yt|Xt,D1:T\t) = 1 T T∑ t=1 Nt∑ n=1 log p(yt,n|xt,n,D1:T\t). (1)
We condition the prediction for samples in each task on the data of other tasks, D1:T\t, to leverage the shared knowledge from related tasks.
The probabilistic multi-task learning formalism in (1) provides a general framework to explore task relatedness 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 a unified way by specifying their priors.
We introduce the latent variables wt and zt,n to represent the classifier and the latent representation of the sample xt,n, respectively. The joint distribution for the task t can be factorized as
p(Yt,Zt,wt|Xt,D1:T\t) = Nt∏ n=1 p(yt,n|zt,n,wt)p(zt,n|xt,n,D1:T\t)p(wt|D1:T\t), (2)
where we use Zt = {zt,n}Ntn=1 to 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 at different levels: zt,n captures the uncertainty of each sample, while wt works at the categorical level. A graphical illustration of the probabilistic latent variable model is shown in Fig. 1.
Solving the model for multi-task learning involves inferring the true joint posterior p(Zt,wt|D1:T ) over all latent variables Zt and wt, which is generally intractable. We introduce a variational joint distribution q(Zt,wt|Dt) for the current task to approximate the true posterior. Also under the conditional independence assumption, the joint variational posterior distribution can be factorized with respect to classifiers and latent representations as follows:
q(Zt,wt|Dt) = qθ(wt|Dt) Nt∏ n=1 qφ(zt,n|xt,n), (3)
where qθ(wt|Dt) and qφ(zt,n|xt,n) are variational posterior distributions for classifiers and latent representations respectively, and θ and φ are the associated parameters. For computational feasibility, they are defined as fully factorized Gaussians, as is common practice [28, 5].
To obtain a 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 [ q(Zt,wt|Dt)||p(Zt,wt|D1:T ) ] . (4)
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-likelihood in (1) :
1
T T∑ t=1 log p(Yt|Xt,D1:T\t) ≥ 1 T T∑ t=1 { Nt∑ n=1 { Ewt∼qθEzt,n∼qφ [log p(yt,n|zt,n,wt)]
−DKL[qφ(zt,n|xt,n)||p(zt,n|xt,n,D1:T\t)] } − DKL [ qθ(wt|Dt)||p(wt|D1:T\t) ]} .
(5)
The objective function provides a general probabilistic inference framework for multi-task learning. As priors naturally serve as regularizations in Bayesian inference, they offer a 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, we will describe the specification of the prior distributions over the classifier and the representation, as well as their variational posterior distributions.

heading: 2.2 Learning Task Relatedness via Gumbel-Softmax Priors
text: The proposed variational multi-task inference framework enables related tasks to provide supportive information for the current task through the conditional priors of the latent variables. It 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 (η) t |D1:T\t) := ∑ i 6=t αtiqθ(w (η−1) i |Di), (6)
where η denotes the η-th iteration. In practice, to avoid entanglement among tasks, we define the prior of task t in the current iteration to be a combination of the variational posteriors [40] of the remaining
tasks from the last iteration. Particularly, our designed prior resembles the empirical Bayesian prior [20, 3, 46, 44] 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.
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-Softmax technique learn the mixing weights for all related tasks and define the mixing weights as follows:
αti = exp((log πti + gti)/τ)∑ i6=t exp((log πti + gti)/τ)) . (7)
Here αti is the mixing weight that indicates the relatedness between tasks t and i. πti is the learnable parameter in the Gumbel-Softmax formulation, which denotes the probability of two tasks transferring positive knowledge. gti is sampled from a Gumbel distribution, using inverse transform sampling by drawing u ∼ Uniform(0, 1) and computing g = − log(− log(u)). τ is the softmax temperature. By using the Gumbel-Softmax technique, our model effectively handles possible negative transfer between tasks. The Gumbel-Softmax 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.
With respect to the inference of the variational posteriors, we define them as fully factorized Gaussians for each class independently:
qθ(wt|Dt) = C∏ c=1 qθ(wt,c|Dt,c) = C∏ c=1 N (µt,c, diag(σ2t,c)). (8)
where µt,c and σ 2 t,c 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 as input and returns the parameters µt,c and σt,c 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 a 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 (6), we specify the prior over the latent representation zt,n of a sample xt,n as a mixture of distributions conditioned on the data of every other task, as follows:
p(z (η) t,n|xt,n,D1:T\t) := ∑ i6=t βtiqφ(z (η−1) t,n |xt,n,Di). (9)
Here, β is the mixing weight which is defined in a similar way as in (7). The conditional distribution qφ(zt,n|xt,n,Di) on the right hand side of (9) indicates that we leverage the data Di from every other task i to 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 qφ(zt,n|xt,n,Di) is inferred by an amortized network [15, 28]. To be more specific, the inference network takes an aggregated representation of xt,n and Di as input and returns the parameters of distribution qφ. The aggregated representation is formulated as (10) which is established by the cross attention mechanism [26]. The sample xt,n from the current task acts as a query, and Di,c plays the roles of the key and the value. Di,c includes samples from the i-th task, which have the same label as xt,n. This is formulated as:
f(xt,n,Di) = DotProduct(xt,n, Di,c, Di,c) := softmax( xt,nD > i,c√ d )Di,c, (10)
where f is the aggregation function, and Di,c is a matrix, with each row containing a sample from class c of the i-th related task. d is the dimension of the input feature. Since we are dealing with supervised learning in this work, class labels of training samples are 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 qφ(zt,n|xt,n) over latent representations is also achieved using the amortization technique [15, 28]. The amortized inference network takes xt,n as input and returns the statistics of its probabilistic latent representation.

heading: 2.3 Empirical Objective Function
text: By integrating (6) and (9) into (5), we obtain the following empirical objective for variational multi-task learning with Gumbel-Softmax priors:
L̂VMTL(θ, φ, α, β) = 1
T T∑ t=1 { Nt∑ n=1 { 1 ML L∑ `=1 M∑ m=1 [ − log p(yt|z(`)t,n,w (m) t ) ] + DKL [ qφ(zt,n|xt,n)||
∑ i 6=t βtiqφ(zt,n|xt,n,Di) ]} + DKL [ qθ(wt|Dt)|| ∑ i6=t αtiqθ(wi|Di) ]} ,
(11) where z(`)t,n ∼ qφ(zt,n|xt,n) and w (m) t ∼ qθ(wt|Dt). To sample from the variational posteriors, we adopt the reparameterization trick [28]. L 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. θ 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-likelihood 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., DKL [ q(wt)|| ∑ i 6=t αtiq(wi) ] ≤ ∑ i 6=t αtiDKL [ q(wt)||q(wi) ] . 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 xt by calculating the following probability using Monte Carlo sampling:
p(yt|xt) ≈ 1
ML L∑ l=1 M∑ m=1 p(y|z(l)t ,w (m) t ), (12)
where we draw samples from posteriors: z(l)t ∼ qφ(z|x) and wt(m) ∼ qθ(w|Dt).

heading: 3 Related Works
text: Multi-Task Learning Settings There are mainly two typical multi-task learning settings: singleinput 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 to boost performance for multi-label prediction [38, 42] and various pixel-level prediction tasks, such as depth prediction, semantic segmentation and surface normal prediction [9, 14, 25, 32, 34, 50]. In the multi-input multi-output setting, each task has its own input from a different domain and output [33, 55, 3, 51]. Tasks are related by sharing the same target space, e.g., class label space or regression target space. This setting, which is the focus of our work, is made challenging by the distribution shift between different tasks.
Probabilistic Multi-Task 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 parameters [49, 45, 1, 3] under the Bayesian framework, or sharing the covariance structure of parameters [11]. 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 to these works, our probabilistic modeling explores task relatedness by specifying priors under the variational Bayesian inference framework.
Regularization-based Multi-Task Learning This has been extensively investigated [52], with the aim of exploring different constraints, e.g., l∞,1 and lp,q norms, imposed on model parameters or features as regularizers [2, 30, 37, 24, 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
the task relatedness underlying parameter tensors to regularize the maximum a 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/regressors jointly in a single framework.
Gumbel-Softmax Technique This technique [22] was recently introduced into multi-task 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 [7], 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.

heading: 4 Experiments
text: We conduct extensive experiments to evaluate the proposed VMTL on five benchmark datasets for both classification and regression tasks. We perform comprehensive comparisons with previous methods and ablation studies to gain insights into the effectiveness of our methods. We report the major experimental results here and provide more results in the supplementary materials.

heading: 4.1 Datasets
text: 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 in total.
Office-Caltech [16] contains the ten categories shared between Office-31 [39] and Caltech-256 [18]. One task uses data from Caltech-256 (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 8 ∼ 151 samples per category per task, and 2, 533 images in total. ImageCLEF [33], 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 distinct 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, it has therefore not been used for multi-task learning before.
Rotated MNIST [29] is 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.

heading: 4.2 Experimental Settings
text: Implementation Details We adopt the standard evaluation protocols [52] for the Office-Home, Office-Caltech 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 DomainNet dataset, we set the splits to 1%, 2% and 4%, which also results in an average of 3, 6 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 a fair comparison, all methods in our paper share the same split documents.
Following the experimental settings in [33], we 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 our experiments, the temperature of the Gumbel-Softmax priors (6) and (9) is annealed using the same schedule applied in [22]: we start with a high temperature and gradually anneal it to a small but non-zero value. For the KL-divergence in (11), we use the annealing scheme from [6].
We adopt the Adam optimizer [27] with a learning rate of 1e-4 for training. All the results are obtained based on a 95% confidence interval from five runs. The detailed network architectures are given in the supplementary materials. The code will be available at https://github.com/autumn9999/VMTL.git.
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 a comprehensive comparison, we 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, we define variational Bayesian extensions of STL and BMTL, which we name VSTL and VBMTL, by placing priors as uninformative standard Gaussian distributions (details can be found in the supplementary materials). VSTL and VBMTL serve as baselines to directly demonstrate the benefits of the probabilistic modeling of multi-task learning based on variational Bayesian inference. We use VMTL-AC to represent our proposed method with amortized classifiers. We compare against four representative methods [3, 25, 33, 38], which are implemented by following the same experimental setup as our methods. Details of the implementation can be found in supplementary materials.

heading: 4.3 Results
text: 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-Softmax 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.
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 the 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 are shown in Table 1, which indicates whether z and w are probabilistic or not, e.g., a × for w means a deterministic classifier. More detailed experimental results are provided in the supplementary materials. As can be seen, both variational Bayesian representations and classifiers can benefit performance. This benefit becomes more significant when training data is very limited, which indicates the effectiveness of leveraging shared knowledge by conditioning priors on related tasks. In addition, in Fig. 2, VSTL demonstrates to be a 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-Softmax Priors The introduced Gumbel-Softmax priors provide an effective way to learn data-driven task relationships. To demonstrate this, we compare with two alternatives without the Gumbel-Softmax technique in Table 2, including the mean and the learnable weighted posteriors of other tasks. The proposed Gumbel-Softmax priors perform the best, consistently surpassing other alternatives. We would like to highlight that the advantage of Gumbel-Softmax priors is even larger under very limited training data, e.g., 5%, which creates a 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 of the mixing weights α and β for both VMTL and VMTL-AC. The results on Office-Home with the 5% split are shown in Fig. 3. We can 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 are 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, OfficeCaltech, ImageCLEF datasets and large-scale DomainNet dataset are shown in Tables 3, 4, 5 and 6, respectively. The average accuracy of all tasks is used for overall performance measurement. The best results are marked in bold, while the second-best are marked by underlining.
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 a 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], a Bayesian method for multi-task learning optimized by an expectation-maximization algorithm, produce very competitive performance. Long et al. [33] use point estimation to find approximate values of model parameters. However, our methods
consistently outperform Bakker et al. [3] and Long et al. [33], which shows the great effectiveness of our 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 is 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.
In addition, we apply our methods to a 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.

heading: 5 Conclusion
text: 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 as a variational inference problem, which provides a unified way to jointly explore knowledge shared among representations and classifiers by specifying priors. We further introduce GumbelSoftmax priors, which offer an effective way to learn the task relatedness in a 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.

heading: Acknowledgment
text: 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.

Review summaries:
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.

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.

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. 

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.