Multi-Task Neural Processes

Neural processes have recently emerged as a class of powerful neural latent variable models that combine the strengths of neural networks and stochastic processes . As they can encode contextual data in the network ’ s function space , they offer a new way to model task relatedness in multi-task learning . To study its potential , we develop multi-task neural processes , a new variant of neural processes for multi-task learning . In particular , we propose to explore transferable knowledge from related tasks in the function space to provide inductive bias for improving each individual task . To do so , we derive the function priors in a hierarchical Bayesian inference framework , which enables each task to incorporate the shared knowledge provided by related tasks into its context of the prediction function . Our multi-task neural processes methodologically expand the scope of vanilla neural processes and provide a new way of exploring task relatedness in function spaces for multi-task learning . The proposed multi-task neural processes are capable of learning multiple tasks with limited labeled data and in the presence of domain shift . We perform extensive experimental evaluations on several benchmarks for the multi-task regression and classification tasks . The results demonstrate the effectiveness of multi-task neural processes in transferring useful knowledge among tasks for multi-task learning and superior performance in multi-task classification and brain image segmentation1 . 1 INTRODUCTION . As deep neural networks are black-box function approximations , it is difficult to introduce prior domain or expert knowledge into a prediction function ( Jakkala , 2021 ) . In contrast , Gaussian processes ( Rasmussen , 2003 ) explicitly define distributions over functions and perform inference over these functions given some training examples . This enables reliable and flexible decisionmaking . However , Gaussian processes can suffer from high computational complexity due to the manipulation of kernel matrices . Therefore , there has been continuous interest in bringing together neural networks and Gaussian processes ( Damianou & Lawrence , 2013 ; Wilson et al. , 2016 ; Garnelo et al. , 2018a ; Jakkala , 2021 ) into so-called neural processes . Neural processes ( Garnelo et al. , 2018b ) combine the computational efficiency of neural networks with the uncertainty quantification of stochastic processes . They are a class of neural latent variables model , which deploy a deep neural network to encode context observations into a latent stochastic variable to model prediction functions . Neural processes provide an elegant formalism to efficiently and effectively incorporate multiple datasets into learning distributions over functions . This formalism is also promising in multi-task learning to improve individual tasks by transferring useful contextual knowledge among related tasks . Their capability of estimating uncertainty over predictions also makes them well-suited for multi-task learning with limited data , where each task has only a few training samples . However , neural processes rely on the implicit assumption that the context and target sets are from the same distribution and can be aggregated by a simple average pooling operation ( Kim et al. , 2019 ; Volpp et al. , 2020 ) . This makes it non-trivial to directly apply neural processes to modeling multiple heterogeneous tasks from different domains , where the context data of different tasks are from distinctive distributions ( Long et al. , 2017 ) . 1 Our code is available in https : //anonymous.4open.science/r/Multi-Task-Neural-Processes/ . In this paper , we develop multi-task neural processes ( MTNPs ) , a methodological extension of neural processes for multi-task learning , which fills the theoretical gap of neural processes for multi-task learning . Particularly , we propose to explore task relatedness in the function space by specifying the function priors in a hierarchical Bayesian inference framework . The shared knowledge from related tasks is incorporated into the context of each individual task , which serves as the inductive bias for making predictions in this task . The hierarchical architecture allows us to design expressive datadependent priors . This enables the model to capture the complex task relationships in multitask learning . By leveraging hierarchical modeling , multi-task neural processes are capable of exploring shared knowledge among related tasks in a principled way by specifying the function prior . We validate the effectiveness of the proposed multi-task neural processes by extensive experiments in both multi-task classification and regression . The results demonstrate that multi-task neural processes can effectively capture task relatedness in the function space and consistently improve the performance of each individual task , especially in the limited data regime . 2 PRELIMINARIES : NEURAL PROCESSES . In this section , we briefly review vanilla neural processes ( Garnelo et al. , 2018b ) based on which we derive our multi-task neural processes . Let a data set be given composed of training samples and corresponding labels . In order to better reflect the desired model behaviour at test time ( Garnelo et al. , 2018b ) , the training data is split into the context set D = ( X , Y ) and a target set D∗ = ( X∗ , Y∗ ) , where X = { x1 , · · · , xn } is a subset of the training data , Y = { y1 , · · · , yn } the corresponding set of labels and n the size of the context set . To ensure both sets have the same data distribution , the context set is split from the target set . Now , given the context set ( X , Y ) , we would like to estimate a function that can make predict the labels Y∗ for target samples X∗ . In general , we define a stochastic process by a random function f : X → Y . Given the context set , we define the joint distribution over the function values { f ( x1 ) , · · · , f ( xn ) } , which in Gaussian processes is a multivariate Gaussian distribution parameterized by a kernel function . The rationale of neural processes is to extract knowledge from the context set to specify the prior over the prediction function . Instead of a kernel function , neural processes adopt a deep neural network to define the prior distribution . Specifically , the model introduces a latent variable z to account for uncertainty in the predictions of Y∗ . The observed context set is encoded into the latent variable which is conditioned on the context set , i.e . it follows the prior distribution p ( z|X , Y ) . The latent variable z is a high-dimensional random vector parameterising the stochastic process by f ( X∗ ) = g ( X∗ , z ) . The function g ( · , · ) is an extra fixed and learnable decoder function , which is also implemented by a neural network . Thus , the neural process model can be formulated as follows : p ( Y∗|X∗ , X , Y ) = ∫ p ( Y∗|X∗ , z ) p ( z|X , Y ) dz . ( 1 ) The graphical model for neural processes is shown in Figure 1 ( a ) . The neural process model is optimized using amortized variational inference . Let q ( z|X∗ , Y∗ ) be a variational posterior of the latent variable z . The evidence lower-bound ( ELBO ) for neural processes is given as follows : log p ( Y∗|X∗ , X , Y ) ≥ Eq ( z|X∗ , Y∗ ) [ p ( Y∗|X∗ , z ) ] − DKL [ q ( z|X∗ , Y∗ ) ||p ( z|X , Y ) ] . ( 2 ) In neural processes , the function space defined by deep neural networks allows the model to extract deep features while retaining a probabilistic interpretation ( Jakkala , 2021 ) . In multi-task learn- ing , usually different tasks can be from different domains and have their specific data distributions ( Lawrence & Platt , 2004 ; Long et al. , 2017 ) . Due to the complex data structure of multi-task learning , it is not straightforward to explore task relatedness in such function spaces . In this paper , we aim to extend the methodology of the neural process to the scenarios of multi-task learning to learn the shared knowledge among tasks for improving individual tasks . 3 MULTI-TASK NEURAL PROCESSES . The common setup of multi-task learning is that there are multiple related tasks for which we would like to improve the learning of each individual one by sharing information across the different tasks ( Williams et al. , 2007 ) . Multi-task learning has been studied under different settings ( Requeima et al. , 2019 ; Williams et al. , 2007 ; Lawrence & Platt , 2004 ; Yu et al. , 2005 ; Long et al. , 2017 ) . In this paper , we tackle the multi-input multi-output setting , where each task has a different distribution while different tasks share the same target space ( Lawrence & Platt , 2004 ; Yu et al. , 2005 ; Long et al. , 2017 ; Zhang et al. , 2020 ) . It aims to improve the overall performance of all multiple tasks simultaneously different from the sequential multi-task leaning ( Requeima et al. , 2019 ; Garnelo et al. , 2018a ) . This is a challenging scenario due to the domain shift between tasks , which makes it sub-optimal to directly apply neural processes by incorporating the data from other related tasks into the context of each individual task . 3.1 HIERARCHICAL CONTEXT MODELING . Multi-task learning considers the estimation of random functions fl , l = 1 , ... , L for each of the L related tasks . Each task l has its own training data , which is split into a context set Dl = ( Xl , Yl ) and a target set D∗l = ( X ∗ l , Y ∗ l ) . X ∗ l ∈ Rn ∗ l×d and Xl ∈ Rnl×d are inputs while Y∗l ∈ Rn ∗ l×C and Yl ∈ Rnl×C are outputs . d and C are the sizes of the input space and output space , respectively . n∗l and nl are the sizes of respectively the target and context set for the l-th task . Thus , we obtain the task-specific latent variables fl = fl ( Xl ) ∈ Rnl×C . We use { Dl } Ll=1 to denote all context sets in the dataset , which for brevity we represent as { Dl } and likewise will do for other sets . Formulated this way the goal of multi-task learning is to predict { Y∗l } for given { X∗l } simultaneously with the assistance information { Dl } from all tasks . To this end , we construct a joint prediction distribution with respect to the latent random functions { fl } as : p ( { Y∗l } | { X∗l } , { Dl } ) = L∏ l=1 p ( Y∗l |X∗l , { Dl } ) = L∏ l=1 ∫ p ( Y∗l |X∗l , fl ) p ( fl|M ) dfl . ( 3 ) Here , to enable shared knowledge to be transferred among tasks , we introduce a global variable M which works as a container to collect the useful information from { Dl } of all tasks . In contrast to neural processes for single tasks , the global variable M provides the contextual information from all tasks for each individual task . The concrete formation of M depends on the learning scenarios . For regression tasks , M ∈ RL∗d and each row corresponds to one task , which is the average of all feature vectors from a task . For classification tasks , M ∈ RL∗C∗d where each vector Ml , c is the average of all features of each category from the corresponding task . Similar to Gaussian processes , we assume that the function value is Y∗l = fl ( X ∗ l ) + , where ∼ N ( 0 , σ2 ) is the observation noise . For regression tasks , we can define the predictive likelihood on the target set as p ( Y∗l |X∗l , fl ) = N ( Y∗l |fl ( X∗l ) , σ2 ) . For classification tasks , we use log p ( Y∗l |X∗l , fl ) = ∑n∗l i=1 y ∗ l , ilog ( fl ( x ∗ l , i ) ) as the log-likelihood function , where x ∗ l , i is the i-th target sample from the l-th task . In order to combine Gaussian processes and neural networks in the context of multi-task learning , we define p ( fl|M ) in ( 3 ) as a deep neural network in place of a Gaussian distribution parameterized by a kernel function . To be more specific , we assume that fl is parameterized by a random variable ψl by defining fl ( X ) = Xψ > l . We specify a data dependent prior by conditioning ψl on the global variable M. In this way , we incorporate the transferable knowledge into the learning of the prediction function of the current tasks . Thus , the predictive distribution for the l-th task over its target set can be formulated as follows : p ( Y∗l |X∗l , { Dl } ) = ∫ p ( Y∗l |X∗l , ψl ) pθ ( ψl|M ) dψl . ( 4 ) In effect , ψl denotes the parameters of classifier for classification tasks or regressors for regression tasks . Particularly , the introduced latent variable ψl ∈ RC×d denotes the task specific classifier , where d is the dimension of the input feature and C is the number of classes in the dataset . As done in ( Requeima et al. , 2019 ) , we generate each column of ψl independently from the context samples of the corresponding class . In our case , each column of ψl encodes the context information of its class from all tasks pθ ( ψl|M ) = ∏C c=1 pθ ( ψl , c|Mc ) . Directly aggregating M as done in neural processes for single tasks is not applicable for multi-task learning due to the distribution shift between tasks . The data from related tasks should be processed and adapted to the current task as the contextual information . To this end , we introduce a higherlevel latent variable αl to extract the shared knowledge from M , which is conditioned on the data Dl of each task . This results in a hierarchical Bayesian modeling of functions in the neural process : pθ ( ψl|M ) = ∫ pθ1 ( ψl|αl , M ) pθ2 ( αl|Dl ) dαl , ( 5 ) whereαl is the latent variable to control the access to shared knowledge for each task , which is used to explore the relevant knowledge to the task l. pθ1 ( ψl|αl , M ) and pθ2 ( αl|Dl ) are prior distributions of the latent variable ψl and αl , respectively , which are parameterized with neural networks . To be specific , we define pθ1 ( ψl|αl , M ) = N ( ψl|µ ( ml ) , Σ ( ml ) ) . Here ml contains the relevant knowledge to the task l , which is adapted from the global variable M by a deterministic function ml = h ( αl , M ) , where h ( · ) is a learnable function implemented with a neural network . By substituting ( 5 ) into ( 4 ) , we obtain the model of multi-task neural processes as follows : p ( { Y∗l } | { X∗l } , { Dl } ) = L∏ l=1 ∫ ∫ p ( Y∗l |X∗l , ψl ) pθ1 ( ψl|αl , M ) pθ2 ( αl|Dl ) dψldαl . ( 6 ) The designed hierarchical context modeling provides a principled way to explore task relatedness in the function space , which allows task-specific function variables to leverage the shared knowledge from related tasks . We provide theoretical proof in Appendix to show that the proposed multi-task neural processes are a valid stochastic processes , which completes the theory of multi-task neural processes . The graphical model of the multi-task neural processes is shown in Figure 1 ( b ) .

Summary: The paper proposes a multi-task learning model with neural processes. The model is based on a hierarchical construction whereby each task is conditioned on global and local information. The paper derives a “hierarchical” ELBO for this model that is evaluated with MC sampling. Experiments are presented to validate the proposed model.