Multi-Task Processes

1 INTRODUCTION . Neural Processes ( NPs ) ( Garnelo et al. , 2018b ) are a class of meta-learning methods that model a distribution of functions ( i.e . a stochastic process ) . By considering a task as a function realized from the underlying stochastic process , they can flexibly adapt to various unseen tasks through inference on functions . The adaptation requires only one forward step of a trained neural network without any costly retraining or fine-tuning , and has linear complexity to the data size . NPs can also quantify their prediction uncertainty , which is essential in risk-sensitive applications ( Gal & Ghahramani , 2016 ) . Thanks to such appealing properties , there have been increasing attempts to improve NPs in various domains , such as image regression ( Kim et al. , 2019 ; Gordon et al. , 2020 ) , image classification ( Requeima et al. , 2019 ; Wang & Van Hoof , 2020 ) , time series regression ( Qin et al. , 2019 ; Norcliffe et al. , 2021 ) , and spatio-temporal regression ( Singh et al. , 2019 ) . In this paper , we explore extending NPs to a multi-task setting where correlated tasks are realized simultaneously from multiple stochastic processes . Many real-world data represent multiple correlated functions , such as different attributes or modalities . For instance , medical data ( Johnson et al. , 2016 ; Harutyunyan et al. , 2019 ) or climate data ( Wang et al. , 2016 ) contain various correlated attributes on a patient or a region that need to be inferred simultaneously . Similarly , in multi-task vision data ( Lin et al. , 2014 ; Zhou et al. , 2017 ; Zamir et al. , 2018 ) , multiple labels of different visual modalities are associated with an image . In such scenarios , it is beneficial to exploit functional correlation by modeling the functions jointly rather than independently , in terms of performance and efficiency ( Caruana , 1997 ) . Unfortunately , naive NPs lack mechanisms to jointly handle a set of multiple functions and can not capture their correlations either . This motivates us to extend NPs to model multiple tasks jointly by exploiting the inter-task correlation . In addition to extending NPs to multi-task settings , we note that handling multi-task data often faces a practical challenge where observations can be incomplete ( i.e . not all the functions share the common sample locations ) . For example , when we collect multi-modal signals from different sensors , the sensors may have asynchronous sampling rates , in which case we can observe signals from only an arbitrary subset of sensors at a time . To fully utilize such incomplete observations , the model should be able to associate functions observed in different inputs such that it can improve the predictive performance of all functions using their correlation . A multivariate extension of Gaussian Processes ( GPs ) ( Álvarez et al. , 2012 ) can handle incomplete observations to infer multiple functions jointly . However , naive GPs suffer from cubic complexity to the data size and needs approximations to reduce the complexity . Also , their behaviour depend heavily on the kernel choice ( Kim et al. , 2019 ) . To address these challenges , we introduce Multi-Task Neural Processes ( MTNPs ) , a new family of stochastic processes that jointly models multiple tasks given possibly incomplete data . We first design a combined space of multiple functions , which allows not only joint inference on the functions but also handling incomplete data . Then we define a Latent Variable Model ( LVM ) of MTNP that theoretically induces a stochastic process over the combined function space . To exploit the inter-task correlation , we introduce a hierarchical LVM consists of ( 1 ) a global latent variable that captures knowledge about all tasks and ( 2 ) task-specific latent variables that additionally capture knowledge specific to each task conditioned on the global latent variable . Inducing each task conditioned on the global latent , the hierarchical LVM allows MTNP to effectively learn and exploit functional correlation in multi-task inference . MTNP also inherits advantages of NP , such as flexible adaptation , scalable inference , and uncertainty-aware prediction . Experiments in synthetic and real-world datasets show that MTNPs effectively utilize incomplete observations from multiple tasks and outperform several NP variants in terms of accuracy , uncertainty estimation , and prediction coherency . 2 PRELIMINARY . 2.1 BACKGROUND : NEURAL PROCESSES . We consider a task f t : X → Yt as a realization of a stochastic process over a function space ( Yt ) X that generates a data Dt = ( XD , Y tD ) = { ( xi , yti ) } i∈I ( Dt ) , where I ( Dt ) denotes a set of data index . Neural Processes ( NPs ) use a conditional latent variable model to learn the stochastic process . Given a set of observations Ct = ( XC , Y tC ) = { ( xi , yti ) } i∈I ( Ct ) , NP infers the target task f t through a latent variable z and models the data Dt by a factorized conditional distribution p ( Y tD|XD , z ) : p ( Y tD|XD , Ct ) = ∫ p ( Y tD|XD , z ) p ( z|Ct ) dz = ∫ ∏ i∈I ( Dt ) p ( yti |xi , z ) p ( z|Ct ) dz . ( 1 ) We refer to the set of observations Ct as a context data and the modeling data Dt as a target data . NP models the generative model p ( Y tD|XD , z ) and the conditional prior p ( z|Ct ) by two neural networks , a decoder pθ and an encoder qφ , respectively . Since the direct optimization of Eq.1 is intractable , the networks are trained by maximizing the following variational lower-bound . log pθ ( Y t D|XD , Ct ) ≥ Eqφ ( z|Dt ) [ log pθ ( Y t D|XD , z ) ] −DKL ( qφ ( z|Dt ) ||qφ ( z|Ct ) ) . ( 2 ) Note that the decoder network qφ is also used as a variational posterior qφ ( z|D ) . The parameter sharing between model prior and variational posterior gives us an intuitive interpretation of the loss function : the KL term acts as a regularizer for the encoder qφ such that the summary of the context is close to the summary of the target . This reflects the assumption that the context and target are generated by the same underlying data-generating process and aids effective test-time adaptation . After training , NP infers the target function according to the latent variable model ( Eq.1 ) . 2.2 EXTENDING TO MULTIPLE TARGET FUNCTIONS . Now we extend the setting to multi-task learning problems where multiple tasks f1 , · · · , fT are realized from T stochastic processes simultaneously , each of which has its own function space ( Yt ) X , ∀t ∈ T = { 1 , 2 , · · · , T } . Let D = ( XD , Y 1 : TD ) = ⋃ t∈T D t be a multi-task target data , where each Dt corresponds to the data of task f t. Then the learning objective for the set of T realized tasks is to model the conditional probability p ( Y 1 : TD |XD , C ) given the multi-task context C = ( XC , Y 1 : T C ) = ⋃ t∈T C t , where each Ct is a set of observations of task f t. The sets C and D can be arbitrarily chosen , but we assume C ⊂ D for simplicity . However , assuming the complete context C for all tasks is often challenged by many practical issues , such as asynchronous sampling across multiple sensors or missing labels in multi-attribute data . To address such challenges , we relax the assumptions on context C and let I ( Ct ) be different across t ∈ T . In this case , an input point xi can be associated with a partial set of output values { yti } t∈Ti , Ti ( T , which is referred incomplete observation . Next , we present two ways to use NPs to model the multi-task data and discuss their limitations . Single-Task Neural Processes ( STNPs ) A straightforward application of NPs to the multi-task setting is assuming independence across tasks and define independent NPs over the function spaces ( Y1 ) X , · · · , ( YT ) X . We refer to this approach as Single-Task Neural Processes ( STNPs ) . Specifically , a STNP has T independent latent variables v1 , · · · , vT , where each vt implicitly represents a task f t. p ( Y 1 : TD |XD , C ) = T∏ t=1 ∫ p ( Y tD|XD , vt ) p ( vt|Ct ) dvt . ( 3 ) Thanks to the independence assumption , STNPs can handle incomplete context by conditioning on each task-specific data Ct independently . However , this approach can only model the marginal distributions for each task , ignoring complex inter-task correlation within the joint distribution of the tasks . Note that this is especially impractical for multi-task settings under the incomplete data since each task f t can be learned only from Ct , ignoring rich contexts available in other data Ct ′ , ∀t′ 6= t. Joint-Task Neural Process ( JTNP ) An alternative approach is to combine output spaces to a product space Y1 : T = ∏ t∈T Yt and define a single NP over the function space ( Y1 : T ) X . We refer to this approach as Joint-Task Neural Processes ( JTNPs ) . In this case , a single latent variable z governs all T tasks jointly . p ( Y 1 : TD |XD , C ) = ∫ p ( Y 1 : TD |XD , z ) p ( z|C ) dz . ( 4 ) JTNPs are amenable to incorporate correlation across tasks through the shared variable z . However , by definition , they require complete context and target for both training and inference . This is because any incomplete set of output values { yti } t∈Ti for an input point xi such that Ti 6= T is not a valid element of the product space Y1 : T . In addition , it relies solely on a single latent variable to explain all tasks , ignoring per-task stochastic factors in each function f t. In what follows , we propose an alternative formulation for jointly handling multiple tasks on incomplete data , which ( 1 ) enables a probabilistic inference on the incomplete data and ( 2 ) is more amenable for learning both task-specific and task-agnostic functional representations . 3 MULTI-TASK NEURAL PROCESSES . In this section , we describe Multi-Task Neural Processes ( MTNPs ) , a family of stochastic processes to model multiple functions jointly and handle incomplete data . We first formulate MTNPs using a hierarchical LVM . Then we propose the training objective and a neural network model . 3.1 FORMULATION . Our objective is to extend NPs to jointly infer multiple tasks from incomplete context . Discussions in Section 2.2 suggest that direct modeling of a distribution over functions of form f : X →∏ t∈T Yt is achievable via JTNP ( Eq . 4 ) , yet it requires complete data in both training and inference . To circumvent this problem , we reformulate the functional form by h : X × T → ⋃ t∈T Yt . Note that this functional form allows us to model the same set of functions as JTNP by f ( xi ) = ( h ( xi , 1 ) , · · · , h ( xi , T ) ) . However , by using the union form we can exploit incomplete data since any partial set of output values { yti } t∈Ti now becomes a set of valid output values at different input points ( xi , t ) , t ∈ Ti . For notational convenience , we denote xti = ( xi , t ) and assume input points in the context C and the target D are embedded by the task indices , i.e. , C = ( X1 : TC , Y 1 : T C ) = ⋃ t∈T C t where Ct = ( XtC , Y t C ) = { ( xti , yti ) } i∈I ( Ct ) and the same for D. Next , we present a latent variable model that induces a stochastic process over functions of form h. To make use of both task-agnostic and task-specific knowledge , we define a hierarchical latent variable model ( Figure 1 ( c ) ) . In this model , the global latent variable z captures shared stochastic factors across tasks using the whole context C , while per-task stochastic factors are captured by the task-specific latent variable vt using Ct and z . It induces the predictive distribution on the target by : p ( Y 1 : TD |X1 : TD , C ) = ∫ ∫ [ T∏ t=1 p ( Y tD|XtD , vt ) p ( vt|z , Ct ) ] p ( z|C ) dv1 : T dz , ( 5 ) where v1 : T : = ( v1 , · · · , vT ) . Similar to Eq . 1 , we assume the conditional independence on p ( Y tD|XtD , vt ) . Note that this hierarchical model can capture and leverage the inter-task correlation by sharing the same z across v1 : T . Also , it is amenable to fully utilize the incomplete data : since the global variable z is inferred from the entire context data C = ⋃ t∈T C t and is conditioned to infer task-specific latent variable vt , each function f t induced by vt exploits the observations available for not only itself Ct , but also for other tasks Ct ′ , ∀t′ 6= t. Next , we show that Eq . 5 induces a stochastic process over the functions of form h : X × T → ⋃ t∈T Yt . Proposition 1 . Consider the following generative process on data D and context C , which is a generalized form of Eq . 5. z ∼ p ( z|C ) , vt ∼ p ( vt|z , t , C ) , yti ∼ p ( yti |xti , vt ) , ∀t ∈ T , ∀i ∈ I ( D ) . ( 6 ) Then under some mild assumptions , there exists a stochastic process over functions of form h : X × T → ⋃ t∈T Yt , where the data D is generated . Proof . We leave the proof in Appendix A.2 . We refer to the resulting stochastic processes as Multi-Task Neural Processes ( MTNPs ) . In the perspective of stochastic process , Eq . 5 allows us to learn functional posterior not only on each task via vt , but also across the tasks via z . Then optimizing Eq . 5 can be interpreted as learning to learn each task captured by vt together with the functional correlation captured by z .

Summary: The paper presents a novel model for multi-task learning with missing data based on Neural Processes (NP). Inter-task correlations are modelled via a shared latent variable. The model has been tested on 1 synthetic and 2 real-world datasets and is experimentally shown to perform better against 4 baselines including 2 naive extensions of NP to the multi-task setting and 2 multi-output Gaussian process models (with spectral kernels).