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: There is a long line of recent interesting work on neural processes, a scalable and more flexible alternative to GPs for performing prediction at a set of test points (x1, ..., xm) given a conditioning set ((x, y)_1, ..., (x, y)_n). This mapping is learned via meta-learning.   This paper addresses a core issue of the popular conditional neural process: the predictions at each test point are conditionally independent given the conditioning set. This is an inappropriate modeling assumption for many real-world datasets. In response, the authors propose to go beyond a non-diagonal Gaussian to describe the joint distribution. For example, they use some structured Gaussian covariances (linear, kvv) and also a Gaussian copula model.   They demonstrate both qualitatively and quantitatively that modeling these dependencies improves the performance of the model on a variety of datasets spanning application domains.