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: This paper presents a multi-task neural processes approach in which function priors are derived in a hierarchical Bayesian inference framework to incorporate the shared knowledge into its context of the prediction function. Authors introduced a higher-level latent variable derived from the context data of related tasks to control the sharing of common knowledge between the tasks. Previous works used context data from its own task to generate the task-specific latent variables, whereas this work uses data from other tasks as well. The shared knowledge from all the tasks added to context of each individual task acts as inductive bias for predictions. Their experiment results show better performance of the proposed method than compared ones.