Variational Multi-Task Learning with Gumbel-Softmax Priors

1 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 [ 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 . 2 Methodology . 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 . 2.1 Variational Multi-Task 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 ( 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 . 2.2 Learning Task Relatedness via Gumbel-Softmax Priors . 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 .

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.