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: This work proposes a multi-task learning architecture for neural processes termed the Multi-task process (MTP). The MTP model conditions task-specific latent variables on a global latent variable that is responsible for information sharing between the tasks, and is able to handle both the isotopic and the heterotopic cases. The model was instantiated with an attentive neural process architecture and the generative model of the MTP was shown to correspond to a stochastic process by a Kolmogorov Extension Theorem argument.