Practical Conditional Neural Process Via Tractable Dependent Predictions

1 INTRODUCTION . Conditional Neural Processes ( CNP ; Garnelo et al. , 2018a ) are a scalable and flexible family of metalearning models capable of producing well-calibrated uncertainty estimates . CNPs naturally handle off-the-grid and missing data and are trained using a simple-to-implement maximum-likelihood procedure . At test time , CNPs require significantly less computation and memory than other metalearning approaches , such as gradient-based fine tuning ( Finn et al. , 2017 ; Triantafillou et al. , 2019 ) , making them ideal for resource and power-limited applications , such as mobile devices . Further , CNPs can be combined with attention ( Kim et al. , 2019 ) , or equivariant networks which account for symmetries in the task at hand ( Gordon et al. , 2020 ; Kawano et al. , 2021 ; Holderrieth et al. , 2021 ) , achieving impressive performance on a variety of problems . Despite these favourable qualities , CNPs are severely limited by the fact that they do not model dependencies in their output ( fig . 1 ) . Limitations of CNPs : More specifically , given two target input locations xm and xm′ , CNPs model their respective outputs ym and ym′ independently . In this paper , we refer to such predictions as mean-field . The inability to model dependencies hurts the predictive performance of CNPs and ren- ders it impossible to produce coherent function samples . Since many downstream tasks require dependent function samples , this excludes mean-field CNPs form a range of applications . In heatwave or flood prediction for example , we need to evaluate the probability of the event that the temperature or precipitation remains above some threshold , throughout a region of space and time . As illustrated by fig . 1 , mean-field predictions model every location independently , and may assign unreasonably low probabilities to such events . If we were able to draw coherent samples from the predictive , the probabilities of such events and similar useful quantities could be more reasonably estimated . Limitations of existing models with dependencies : To address the above , follow-up work has introduced Neural Processes ( NPs ; Garnelo et al. , 2018b ; Kim et al. , 2019 ; Foong et al. , 2020 ) , which use latent variables to model output dependencies . However , the likelihood for these models is not analytically tractable , so approximate inference is required for training ( Le et al. , 2018 ; Foong et al. , 2020 ) . Alternatively , Bruinsma et al . ( 2021 ) recently introduced a variant of the CNP called the Gaussian Neural Process , which we will refer to as the FullConvGNP , which directly parametrises the covariance of a Gaussian predictive over the output variables . In this way the FullConvGNP models statistical dependencies in the output , and can be trained by an exact maximum-likelihood objective , without requiring approximations . However , for D-dimensional data , the architecture of the FullConvGNP involves 2D-dimensional convolutions , which can be very costly , and , forD > 1 , poorly supported by most Deep Learning libraries . Contributions : In this work ( i ) we introduce Gaussian Neural Processes ( GNPs ) , a class of model which directly parametrises the covariance of a Gaussian predictive process , thereby circumventing the costly convolutions of the FullConvGNP , and is applicable to higher-dimensional input data . GNPs have analytic likelihoods making them substantially easier to train than their latent variable counterparts ; ( ii ) we show that GNPs can be easily applied to multi-output regression , as well as composed with invertible marginal transformations to model non-Gaussian data ; ( iii ) we demonstrate that modelling correlations improves performance on experiments with both Gaussian and non-Gaussian synthetic data , including a downstream estimation task that mean-field models can not solve ; ( iv ) we demonstrate that GNPs outperform their mean-field and latent variable counterparts on real-world electroencephalogram ( EEG ) data and climate data ; ( v ) in climate modelling , GNPs outperform a standard ensemble of widely used methods in statistical downscaling , while providing spatially coherent temperature samples which are necessary for climate impact studies . 2 CONDITIONAL & GAUSSIAN NEURAL PROCESSES . Background : We present CNPs from the viewpoint of prediction maps ( Foong et al. , 2020 ) . A prediction map π is a function which maps ( 1 ) a context set ( xc , yc ) where xc = ( xc,1 , . . . , xc , N ) are the inputs and yc = ( yc,1 , . . . , yc , N ) the outputs and ( 2 ) a set of target inputs xt = ( xt,1 , ... , xt , M ) to a distribution over the corresponding target outputs yt = ( yt,1 , ... , yt , M ) : π ( yt ; xc , yc , xt ) = p ( yt|r ) , ( 1 ) where r = r ( xc , yc , xt ) is a vector which parameterises the distribution over yt . For a fixed context ( xc , yc ) , using Kolmogorov ’ s extension theorem ( Oksendal , 2013 ) , the collection of finitedimensional distributions π ( yt ; xc , yc , xt ) for all xt,1 , . . . , xt , M ∈ RM , M ∈ N , defines a stochas- tic process if these are consistent under ( i ) permutations of any entries of ( xt , yt ) and ( ii ) marginalisations of any entries of yt . Prediction maps include , but are not limited to , Bayesian posteriors . One familiar example of such a map is the Bayesian Gaussian process ( GP ; Rasmussen , 2003 ) posterior π ( yt ; xc , yc , xt ) = N ( yt ; m , K ) , ( 2 ) where m = m ( xc , yc , xt ) and K = k ( xc , xt ) are given by the usual GP posterior expressions . Another prediction map is the CNP Garnelo et al . ( 2018a ) : π ( yt ; xc , yc , xt ) = ∏M m=1 p ( yt , m|rm ) , ( 3 ) where each p ( yt , m|rm ) is an independent Gaussian and rm = r ( xc , yc , xt , m ) is parameterised by a DeepSet ( Zaheer et al. , 2017 ) . CNPs are permutation and marginalisation consistent and thus correspond to valid stochastic processes . However , CNPs do not respect the product rule in general ( Foong et al. , 2020 ) . Nevertheless , CNPs and their variants ( Gordon et al. , 2020 ) have been demonstrated to give competitive performance and robust predictions in a variety of tasks and are a promising class of meta-learning models . Gaussian Neural Processes : A central problem with CNP predictive distributions is that they are mean-field : eq . ( 3 ) does not model correlations between yt , m and yt , m′ for m ̸= m′ . However , many tasks require modelling dependencies in the output variable . To remedy this , we consider parameterising a correlated multivariate Gaussian π ( yt ; xc , yc , xt ) = N ( yt ; m , K ) ( 4 ) where , instead of the expressions for the Bayesian GP posterior , we use neural networks to parameterise the mean m = m ( xc , yc , xt ) and covariance K = K ( xc , yc , xt ) . We refer to this class of models as Gaussian Neural Processes ( GNPs ) . The first such model , the FullConvGNP , was introduced by Bruinsma et al . ( 2021 ) with promising results . Unfortunately , the FullConvGNP relies on 2D-dimensional convolutions for parameterising K , applying the sequence of computations ( xc , yc ) 1−→ ( x̃ , h ) 2−→ r = PSD ( CNN2D ( h ) ) 3−→ Kij = ∑L l=1 ψ ( xt , i , x̃l ) rl ψ ( x̃l , xt , j ) ( 5 ) where 1 maps ( xc , yc ) to a 2D-dimensional grid h at locations x̃ = ( x̃1 , . . . , x̃L ) , x̃l ∈ R2D , using a SetConv layer ( Gordon et al. , 2020 ) , 2 maps h to r through a CNN with 2D-dimensional convolutions , followed by a PSD map which ensures r is positive-definite , and 3 aggregates r using an RBF ψ . The CNN at 2 requires expensive 2D-dimensional convolutions , which are challenging to scale to higher dimensions ( see appendix B ) . To overcome this difficulty , we propose parameterising m and K by mi = f ( xt , i , r ) , Kij = k ( g ( xt , i , r ) , g ( xt , j , r ) ) ( 6 ) where r = r ( xc , yc ) , f and g are neural networks with outputs in R and RDg , and k is an appropriately chosen positive-definite function . Note that , since k models a posterior covariance , it can not be stationary . The special case where Kij = σ2i Iij is diagonal corresponds to a mean-field CNP as presented in ( Garnelo et al. , 2018a ) . Equation ( 6 ) defines a class of GNPs which , unlike the FullConvGNP , do not require costly convolutions . GNPs can be readily trained via the log-likelihood θ∗ = argmaxθ log π ( yt ; xc , yc , xt ) , ( 7 ) where θ collects all the parameters of the neural networks f , g , and r. In this work , we consider two methods to parameterise K , which we discuss next . Linear covariance : The first method we consider is the linear covariance Kij = g ( xt , i , r ) ⊤g ( xt , j , r ) ( 8 ) which can be seen as a linear-in-the-parameters model with Dg basis functions and a unit Gaussian distribution on their weights . This model meta-learns Dg context-dependent basis functions , which approximate the true distribution of the target , given the context . By Mercer ’ s theorem ( Rasmussen , 2003 ) , up to regularity conditions , every positive-definite function k can be decomposed as k ( z , z′ ) = ∑∞ d=0 ϕd ( z ) ϕd ( z ′ ) ( 9 ) where ( ϕd ) ∞d=1 is a set of orthogonal basis functions . We therefore expect eq . ( 8 ) to be able to recover arbitrary ( sufficiently regular ) GP predictives as Dg grows large . Further , the linear covariance has the attractive feature that sampling from it scales linearly with the number of target locations . A drawback is that the finite number of basis functions may limit its expressivity . Kvv covariance : An alternative covariance which sidesteps this issue , is the kvv covariance Kij = k ( g ( xt , i , r ) , g ( xt , j , r ) ) v ( xt , i , r ) v ( xt , j , r ) , ( 10 ) where k is the Exponentiated Quadratic ( EQ ) covariance with unit lengthscale and v is a scalaroutput neural network . The v modulate the magnitude of the covariance , which would otherwise not be able to shrink near the context . Unlike linear , kvv is not limited by a finite number of basis functions , but the cost of drawing samples from it scales cubically in the number of target points . Multi-output regression : Extending this approach to the multi-output setting where yt , m ∈ RDy with Dy > 1 , can be achieved by learning functions m1 , . . . , mDy and g1 , . . . , gDy for each dimension of the output variable . We can represent covariances across different target points and different target vector entries , by passing those features through either the linear or the kvv covariance Kijab = ga ( xt , i , r ) ⊤gb ( xt , j , r ) , ( 11 ) Kijab = k ( ga ( xt , i , r ) , gb ( xt , j , r ) ) va ( xt , i , r ) vb ( xt , j , r ) , ( 12 ) where Kijab denotes the covariance between entry a of yt , i and entry b of yt , j . Neural architectures : This discussion leaves room for choosing f , g , and r , producing different models belonging to the GNP family , of which the FullConvGNP is also a member . For example , we may choose these to be DeepSets , attentive Deepsets or CNNs , giving rise to Gaussian Neural Processes ( GNPs ) , Attentive GNPs ( AGNPs ) or Convolutional GNPs ( ConvGNPs ) respectively . Particularly , in the ConvGNP , the feature function g takes the form ( xc , yc ) 1−→ ( x̃ , h ) 2−→ r = CNND ( h ) 3−→ g ( xt , i , r ) = ∑L l=1 ψ ( xt , i , xr , l ) rl , ( 13 ) where , crucially , h are values on a D-dimensional grid at x̃ = ( x̃1 , . . . , x̃L ) , x̃l ∈ RD , and 2 uses D-dimensional rather than a 2D-dimensional CNN . This renders the ConvGNP much cheaper than the FullConvGNP in both compute and memory , while retaining translation equivariance ( see appendix A.1 for proof ) , making the former a scalable alternative to the latter .

Summary: The authors present a class of neural process models that are able to produce correlated predictions while amenable to exact, simple and scalable maximum likelihood optimization supporting multiple outputs. By using invertible transformations (gaussian copula), the model is able to capture non-Gaussian output distributions. Experiments with artificial and real data (EEG and climate), highlight the predictive ability of the proposed model.