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