A UTOREGRESSIVE C ONDITIONAL N EURAL P ROCESSES
Anonymous authors
Paper under double-blind review
A BSTRACT
Conditional neural processes (CNPs; Garnelo et al., 2018a) are attractive metalearning models which produce well-calibrated predictions and are trainable via a simple maximum likelihood procedure.
Although CNPs have many advantages, they are unable to model dependencies in their predictions.
Various works propose solutions to this, but these come at the cost of either requiring approximations or being limited to Gaussian predictions.
In this work, we instead propose to change how CNPs are deployed at test time, without any modifications to the model or training procedure .
Instead of making predictions independently for every target point, we autoregressively define a joint predictive distribution using the chain rule of probability, taking inspiration from the neural autoregressive density estimator (NADE) literature.
We show that this simple procedure allows factorised Gaussian
CNPs to model highly dependent, non-Gaussian predictive distributions.
Perhaps surprisingly, in an extensive range of tasks with synthetic and real data, we show that CNPs in autoregressive (AR) mode not only significantly outperform non-AR
CNPs, but are also competitive with more sophisticated models that are significantly more expensive and challenging to train.
This performance is remarkable since AR
CNPs are not trained to model joint dependencies.
Our work provides an example of how ideas from neural distribution estimation can benefit neural processes, motivating research into the AR deployment of other neural process models.
I NTRODUCTION
Conditional neural processes (CNPs; Garnelo et al., 2018a) are a family of meta-learning models which combine the flexibility of deep learning with the uncertainty awareness of probabilistic models.
They are trained to produce well-calibrated predictions via a simple maximum-likelihood procedure, and naturally handle off-the-grid and missing data, making them ideally suited for tasks in climate science and healthcare.
Since their introduction, attentive (ACNP; Kim et al., 2019) and convolutional (ConvCNP; Gordon et al., 2020) variants have also been proposed.
Unfortunately, existing CNPs do not model statistical dependencies (Figure 1; left).
This harms their predictive performance and makes it impossible to draw coherent function samples, which are necessary in downstream estimation tasks (Markou et al., 2022).
Various approaches have been proposed to address this.
Garnelo et al.
(2018b) introduced the latent neural process (LNP), which uses a latent variable to induce dependencies and model non-Gaussianity.
However, this renders the likelihood intractable, necessitating approximate inference.
Another approach, which maintains tractability at the cost of only allowing Gaussian
[Table]
predictions, is the fully convolutional Gaussian neural process (FullConvGNP; Bruinsma et al., 2021).
It uses a neural network to define the mean and covariance function of a predictive Gaussian process (GP) which models dependencies.
However, it uses a much more complex architecture and is only practically applicable to problems with one-dimensional inputs, limiting its adoption compared to the more lightweight CNP.
Recently, Markou et al.
(2022) proposed the Gaussian neural process (GNP), which is considerably simpler but sacrifices performance relative to the FullConvGNP.
In this paper we propose a much simpler method for modelling dependencies with neural processes that has been largely overlooked: autoregressive (AR) sampling.
AR sampling requires no changes to the architecture or training procedure.
Instead, we change how the CNP is deployed at test time , extracting predictive information that would ordinarily be ignored.
Instead of making predictions at all target points simultaneously, we autoregressively feed samples back into the model.
AR CNPs trade the fundamental property of consistency under marginalisation and permutation , which is foundational to many neural process models, for non-Gaussian and correlated predictions.
In Tablewe place AR CNPs within the framework of other neural process models.
Our key contributions are:
• We show that CNPs used in AR mode capture rich, non-Gaussian predictive distributions and produce coherent samples (Figure 1).
This is remarkable as the CNPs used have Gaussian likelihoods, are not trained to model joint dependencies or non-Gaussianity, and are significantly cheaper to train than LNPs and FullConvGNPs (Figure 2).
• We prove that given sufficient data and model capacity, the performance of AR CNPs is at least as good as that of GNPs, which explicitly model correlations in their predictions.
Viewing AR CNPs as a type of neural density estimator (Uria et al., 2016), we highlight their connections to a range of existing methods in the deep generative modelling literature.
• In an extensive range of Gaussian and non-Gaussian regression tasks, we show that AR CNPs are consistently competitive with, and often significantly outperform, all other neural process models in terms of predictive log-likelihood.
• We deploy AR CNPs on a range of tasks involving real-world climate data.
To handle the high-resolution data in a computationally tractable manner, we introduce a novel multiscale architecture for ConvCNPs.
We also combine AR ConvCNPs with a beta-categorical mixture likelihood, producing strong results compared to other neural processes.
Our work represents a promising first application of this procedure to the simplest class of neural processes, and motivates future work on applications of AR sampling to other neural process models.
A UTOREGRESSIVE C ONDITIONAL N EURAL P ROCESSES
Meta-learning.
We first define the problem set-up.
Let X be a compact input space and let Y be the output space.
Let D N
=
( X × Y )
N be the collection of all sets of N input–output pairs, and let
D = (cid:83) ∞ N =0 D N .
We call elements D ∈ D data sets and denote D =
( x , y ) where x ∈ X N , y ∈ Y
N are the inputs and outputs respectively.
In meta-learning we are given a collection of data sets ( D m )
Mm =1 , called a meta–data set , with the individual data sets D m called tasks (Vinyals et al., 2016).
Every task D m is split up D m = D (c)
m ∪ D (t) m into a context set D (c) m =
( x (c) m , y (c) m ) and a target set D (t) m =
( x (t) m , y (t) m ) .
Here x (c) m are called the context inputs , y (c) m the context outputs , x (t) m the target inputs , and y (t) m the target outputs .
Our goal is to devise an algorithm which takes in a context set
D (c) m and produces the best possible prediction for the target outputs y (t) m given target inputs x (t) m .
Neural processes.
Let P be the set of all Y -valued stochastic processes on X .
Neural processes (NPs) directly and flexibly parametrise a map π θ :
D → Q where Q ⊆ P and where θ are learnable parameters.
CNPs set
Q to be the collection of GPs f such that cov( f ( x ) , f ( y ))
=
for x ̸ = y .
GNPs let Q be the collection of continuous GPs.
Latent NPs (LNPs; Garnelo et al., 2018b) let Q be a collection of non-Gaussian processes by making use of a latent variable.
Let P x (t) m π ( D (c) m ) denote the finite-dimensional distribution of the process π ( D (c) m ) evaluated at inputs x (t) m , and denote its density
[Figure]
N LL / d a t a p o i n t
FullConvGNP ConvLNP
ConvCNP
AR ConvCNP
[Figure]
by q θ ( • | x (t) m , D (c) m ) .
To learn the parameters θ , NPs seek to maximise:
[Equation]
For CNPs and GNPs, L M can be computed exactly since q θ is Gaussian.
However, for LNPs, L M must be approximated (Garnelo et al., 2018b; Foong et al., 2020), typically impacting performance.
Autoregressive CNPs.
Our proposal is to take an existing CNP and run it in an autoregressive fashion, feeding predictions for earlier outputs back into the model.
Inspired by the product rule, we define the joint predictive as a product of conditionals, modelling each conditional with a CNP.
For example, in the case of three target points, q θ ( y (t) 3 | y (t) 1:2 , D (c) m )
q θ ( y (t)
| y (t) 1 , D (c) m )
q θ ( y (t) 1 | D (c) m ) .
To enable a theoretical analysis of this procedure, we now proceed to set up more formal notation.
Suppose that π θ : D → Q is an NP, and we wish to predict at some target inputs x (t) given a context set D (c) .
Standard NPs would output the predictive P x (t) π θ ( D (c) ) which, for CNPs, would be a factorised Gaussian.
We propose to instead roll out the NP autoregressively, as described in Proc.
Procedure 2.1 (Autoregressive application of neural process) .
For a neural process π θ , context set
D (c) =
( x (c) , y (c) ) , and target inputs x (t) , let AR x (t) ( π θ , D (c) ) be the distribution defined as follows:
[Equation]
where a ⊕ b concatenates two vectors a and b .
See Figure 7 in Appendix C for an illustration. Since earlier samples y (t) i feed back into later applications of π θ , the whole sample y (t) is correlated, even if π θ does not model dependencies between target outputs , as with CNPs.
At test time, when evaluating the corresponding the density q (AR) θ of AR x (t) ( π θ , D (c) ) at y (t) , we use the formula
[Equation]
Whilst any NP can be used in AR, we focus on CNPs as they are the computationally cheapest class.
Understanding the infinite data limit.
To better understand why AR CNPs successfully model dependencies, we analyse the idealised case of infinite data and model capacity.
Let p ( f ) be the law of the data-generating stochastic process, and let p ( ε ) be the law of a stochastic process representing observation noise, defined by letting ε ( x ) be a vector of i.i.d.
noise variables for all x .
We assume y (c) m =
y m ( x (c) m ) and y (t) m
=
y m ( x (t) m ) where y m ( • ) =
f m ( • ) + ε m ( • ) , (4) ( f m )
Mm =1 are i.i.d.
draws from p ( f ) , and ( ε m )
Mm =1 are i.i.d.
draws from p ( ε ) .
Define the prediction map π y : D → P as the mapping from a data set to the posterior over y , π y ( D ) =
p ( y | D ) .
Then
L M is a Monte Carlo approximation of the following infinite-sample objective (Foong et al., 2020):
[Equation]
Under appropriate regularity assumptions, L ∞ ( π ) is maximised over all π when the expected KL divergence term is zero, which occurs if and only if π = π y .
In practice, NPs do not maximise L ∞ ( π ) over all π , but (i) use a finite-sized meta–data set and (ii) restrict Q ⊆ P .
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Equation] Here π M is an NP trained on the practical objective (1), which, in the limit of infinite data, approximates the so-called ideal NP π ∞ .
The ideal NP depends on the choice of Q , i.e. the class of NPs under consideration, and, in turn, approximates π y .
For CNPs and GNPs, using the fact that minimising
KL( p, q ) over q matches moments (Minka, 2001), we can readily determine and even practically compute the ideal NP for these two classes of NPs.
The ideal CNP predicts a diagonal-covariance
GP whose mean function and marginal variance function match π y : π ∞ ( D )
=
GP ( m, k ) where m ( x ) =
E [ y ( x ) | D ] , and k ( x, x ′ ) =
V [ y ( x ) | D ] if x = x ′ and k ( x, x ′ )
=
otherwise.
On the other hand, the ideal GNP predicts a GP whose mean function and full covariance function match π y : π ∞ ( D ) =
GP ( m, k ) where m ( x ) =
E [ y ( x ) | D ] , k
( x, x ′ )
=
cov( y ( x ) , y ( x ′ ) | D ) .
The main result of this subsection is that the ideal CNP, despite not modelling correlations, becomes superior to the ideal GNP when deployed in AR mode :
Proposition 2.1 (Advantage of AR CNPs over GNPs) .
Assume appropriate regularity conditions on y .
Let π C be the ideal CNP and let π G be the ideal GNP.
Then, for all inputs x and data sets D ∈ D ,
[Equation]
We provide a proof in Appendix A. Intuitively, the advantage of AR CNPs comes from their ability to model non-Gaussian dependencies.
Proposition 2.1 shows that to outperform the GNP, it suffices to train a CNP to model the marginals of π y , and rely on the AR procedure to induce dependencies.
A visualisation of the ideal CNP and the ideal CNP applied autoregressively can be seen in Figure 3.
Consistency and the AR design space.
As shown in Table 1, AR CNPs give up the fundamental property of consistency , as the distributions { AR x ( π θ , D (c) m ) : x ∈ X N , N ∈ N } are not consistent under permutation or marginalisation: permuting x can change the distribution, as can introducing and marginalising new target points.
This violates the conditions of the Kolmogorov extension theorem (Oksendal, 2013), preventing the distributions from defining a consistent stochastic process.
There is thus a large design space involved when deploying AR CNPs, where choices that have no effect on the predictions of other NPs can now significantly affect performance.
One such choice is how many points to sample at a time.
Sampling one at a time induces dependencies between all points, but requires N forward passes.
Alternatively, we could divide the N inputs in x (t) into blocks of K points each, and sample each block with a single CNP forward pass.
This requires
N/K forward passes, with points in the same block conditionally independent.
If K = N , this is the standard CNP prediction, and if K = 1 , we recover Procedure 2.1.
This provides a knob for practitioners to trade off between faster, consistent, but less expressive standard CNP predictions, and slower, less consistent, but more expressive AR predictions.
In this paper, we use full AR mode with
K = 1 , and leave an investigation of block AR sampling to future work.
Obtaining smooth samples.
Due to the lack of consistency in AR mode, the spacing chosen between target points can significantly affect performance.
For example, care must be taken so the number of target points is not much greater than the size of the context sets seen during train time, to avoid confronting the model with an out-of-distribution context set size at test time.
This raises the question of how to sample functions on a very fine grid.
Furthermore, since CNPs do not differentiate between epistemic and aleatoric uncertainty, it is not clear how to obtain smooth, noiseless samples, that is, samples for f uncorrupted by the i.i.d.
noise ε in (4).
The following proposition shows that, for a smooth sample corrupted by additive noise, the smooth component can be approximated with the predictive mean conditioned on noisy observations:
Proposition 2.2 (Recovery of smooth samples) .
Let X ⊆
R be compact, and let f be a stochastic process with surely continuous sample paths and sup x ∈X ∥ f ( x ) ∥
L 2 < ∞ .
Let ( ε n ) n ≥ 0 be i.i.d.
(potentially non-Gaussian) random variables such that E [ ε 0 ]
= 0 and V [ ε 0 ] < ∞ .
Consider any sequence ( x n ) n ≥ 1 ⊆ X , and let x ∗ ∈
X be a limit point of ( x n ) n ≥ 1 .
If y ( x ∗ ) =
f ( x ∗ ) + ε 0 and y n = f ( x n ) + ε n are noisy observations of f , then
[Equation]
We provide a proof in Appendix B. Equation (8) suggests the following two-step procedure for obtaining smooth samples from AR CNPs.
Step 1: Let x 1: n be a number target inputs that does not exceed the number of points seen during training.
Sample y 1: n ∼ AR x 1: n ( π θ , D (c) m ) .
This sample includes observation noise.
Step 2: Remove the noise from the sample by passing it through the model once more: N ( µ 1: n , D )
=
P x 1: n π θ ( D (c) m ⊕ ( x 1: n , y 1: n )) .
Here the predictive mean µ 1: n forms the noiseless sample.
To produce a sample at arbitrarily many inputs, one may also evaluate
N ( µ ′ 1: n , D ) =
P x ′ 1: n π θ ( D (c) m ⊕ ( x 1: n , y 1: n )) where x ′ 1: n is arbitrary.
This result of this procedure is illustrated in Figure 4, and was used to generate the noiseless samples shown in Figure 1 (right).
Figure 7 in Appendix C also illustrates this two-step procedure in a pictorial step-by-step fashion.
[Figure]
C ONNECTIONS TO O THER N EURAL D ISTRIBUTION E STIMATORS
Various paradigms have been developed for neural distribution estimators, including normalising flows (Dinh et al., 2014), generative adversarial networks (GANs; Goodfellow et al., 2014), variational autoencoders (VAEs; Kingma & Welling, 2013), autoregressive models (Uria et al., 2016), and diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020).
We argue that NPs should be viewed as neural distribution estimators alongside these models, revealing interesting connections between them.
Figure 5 visualises the landscape of neural distribution estimators, highlighting the dimensionality of the target distribution, and the use of attention or convolution.
Our proposed AR CNPs can be viewed both as a kind of NP, and as an instance of autoregressive models.
As such, they inherit the strengths of AR models, such as their ability to model complex dependencies with a tractable likelihood, but also some of their weaknesses, most notably slow test-time sampling.
Slow sampling is the main drawback of applying AR CNPs to large target sets, though it may be possible to adapt techniques for speeding up other AR models to AR CNPs (Ramachandran et al., 2017a).
One major difference between AR CNPs and existing AR models is that AR CNPs model an uncountably infinite set of variables, which allows querying arbitrary input locations but leads to the consistency issues discussed in Section 2.
Another distinguishing feature of the AR CNP is that it is trained to be order agnostic with respect to the product rule decomposition of the joint distribution: there is no preferred ordering of the input values during train time.
DEformer
(Alcorn & Nguyen, 2021), EoNADE
(Uria et al., 2014)
[Table]
and
N e u r a l p r o ce ss e s
N o r m a li
s i
ng
fl o w s
G e n e r a ti v e a dv e r s a r i a l n e t w o r k
s
V a r i a ti on a l a u t o e n c od e r s
A u t o r e g r e ss i v e m od e l s 1,
[Table]
[Table] [Table]
[Figure]
XLnet (Yang et al., 2019) are other examples of AR models designed for order agnosticism.
The AR
CNP shares design choices with other AR models to achieve this goal: (i) a shared architecture is used to produce each conditional distribution, similar to WaveNet (Oord et al., 2016a) and PixelCNN (Oord et al., 2016b); (ii) the datapoint index is given as input to the network as in the DEformer model (Alcorn & Nguyen, 2021); and (iii) training maximises a log-likelihood including all decompositions of the joint distribution, similar to EoNADE (Uria et al., 2014) and
XLnet (Yang et al., 2019).
In recent work, Nguyen & Grover (2022) proposed the Transformer NP (TNPs), which uses a causally-masked transformer architecture with an autoregressive likelihood.
In contrast, rather than proposing a new AR architecture, our work focuses on running existing CNPs in AR mode to obtain coherent samples and improved likelihoods, without modifying the architecture or training procedure.
In prior work, Volpp et al.
(2021) used AR sampling in order to visualise samples from CNPs.
However, their work focuses on proposing a novel context aggregation mechanism for NPs, and they do not evaluate the likelihood of CNPs autoregressively or investigate any performance gains.
Figure 5 also shows the connections betweeen NPs and VAEs and normalising flows.
Just as VAEs typically use decoders that parametrise a factorised distribution, and rely on the latent variable to induce dependencies, the same is true for LNPs.
Again, the key difference with standard VAEs is that LNPs model a distribution over an uncountable number of variables.
In addition, models like conditional BRUNO (Korshunova et al., 2020) and copula GNPs (Markou et al., 2022) combine ideas from NPs and normalising flows, by transforming a stochastic process with an invertible transformation.
Finally, GAN models such as Spatial GAN (Jetchev et al., 2016) and ∞− GAN (Lu et al., 2020) model countable numbers of variables, such as images of arbitrary size.
Inspecting
Figure 5, we see that GANs are the only class of models depicted that do not currently have an NP version, or a version that models an uncountable number of variables.
This suggests adversarial training of NPs as an interesting avenue for future investigation.
T HE P ERFORMANCE OF A UTOREGRESSIVE N EURAL P ROCESSES
In this section we investigate the performance of AR CNPs on synthetic and real data.
Across a wide range of tasks, the AR CNP is competitive with much more sophisticated approaches.
Throughout, we train LNPs with both the ELBO (Garnelo et al., 2018b) and ML objective (Foong et al., 2020).
Code for reproducing our experiments can be found at https://github.com/user/repo .
For all experiments, we use a random ordering of the target points in Proc. 2.1; see App.
D for a justification.
4. S YNTHETICALLY G ENERATED G AUSSIAN AND N ON -G AUSSIAN D ATA
Synthetic experiment setup.
We evaluate an extensive collection of NP models on a wide range of Gaussian and non-Gaussian synthetic regression tasks.
We consider tasks with functions drawn from (i) different GPs; (ii) a non-Gaussian sawtooth process (as in Figure 1); (iii) a non-Gaussian mixture task, where, with equal probability, we sample the function from one of three possible GPs or the sawtooth process.
We also consider various versions of the tasks for different input and output dimension d x , d y , with dependencies across the output channels.
To ensure a fair comparison, we choose architectural details to make the parameter counts comparable between all models.
Results.
Table 2 highlights the best performing models on some representative tasks; for further results across all twenty synthetic tasks and further experimental details, see Appendix H. The AR procedure dramatically improves the performance of the ConvCNP, with the AR ConvCNP being the best performing model for most tasks, except on the Gaussian EQ task where it performs marginally worse than the FullConvGNP.
In particular, the AR ConvCNP outperforms the FullConvGNP and
ConvGNP on non-Gaussian tasks, in agreement with Proposition 2.1, while having a faster training time than the other convolutional models (Figure 2).
For the sawtooth task, Figure 10 in Appendix H. illustrates that predictions by the AR ConvCNP can be multi-modal and non-Gaussian, even when using a Gaussian likelihood.
Finally, we note that in tasks with d x = 2 , where the FullConvGNP cannot be used (as discussed in Section 1), the AR ConvCNP far outperforms all competing approaches.
4. S IM - TO -R EAL T RANSFER WITH THE L OTKA –V OLTERRA E QUATIONS
Predator-prey data.
We next investigate sim-to-real transfer, where the models are trained on simulated data and tested on real data.
NPs are well-suited to this setting, since a large meta-data set can be easily generated to train them.
We consider the hare–lynx data set, which is a population time series of Snowshoe hares and Canadian lynx (MacLulich, 1937).
To generate simulated data, we use a stochastic version of the Lotka–Volterra equations (Lotka, 1910; Volterra, 1926):
[Equation]
[Table] [Table]
[Table]
Under these equations, the prey population X t grows exponentially with rate α , the predator population Y t decays exponentially with rate γ , and the predators hunt the prey.
W (1) and W (2) are independent Brownian motions introducing noisy behaviour.
These equations generate non-Gaussian data with both within-channel as well as cross-channel dependencies.
We simulate the Lotka-Volterra equations on a dense grid, and use them to generate meta–data sets in three different ways.
Interpolation : we randomly subsample the data into context and target sets.
Forecasting : we choose a random time, before which all data are contexts, and all future data are targets.
Reconstruction: we randomly choose between the X t or Y t , split the chosen series as in forecasting, and append the other series to the context.
In training, for every batch, we choose one of these tasks uniformly at random.
Results.
Table 3 shows the results of the best performing models.
The AR ConvCNP performs best both on the simulated as well as the real data, again demonstrating that running CNPs in AR mode improves performance and can even outperform strong GNP and LNP baselines.
For full experimental details and additional results see Appendix I.
4. E LECTROENCEPHALOGRAM EXPERIMENTS
Electroencephalogram data.
We next trained various NPs on real time series data consisting of electroencephalogram (EEG) measurements (Zhang et al., 1995b), following Markou et al.
Each time series consists of 256 regularly spaced measurements across 7 correlated channels.
For each channel, we randomly select a number of the 256 points uniformly at random to be target points, and use the remaining ones as context points, independently across the channels.
Results.
After training, we test the models on this interpolation task and also on a reconstruction task, where we set a part of a channel as target and the remainder as context.
In Table 4, we observe that the AR ConvCNP is competitive with the FullConvGNP, despite having significantly shorter training times and fewer parameters.
Both the AR ConvCNP and the FullConvGNP outperform the
ConvCNP and the ConvLNP.
Full experimental detail are in Appendix J.
4. E NVIRONMENTAL M ODELLING
Data fusion.
A ubiquitous challenge in environmental sciences is fusing spatio-temporal data from disparate sources (Chang & Bai, 2018; Lahat et al., 2015).
This challenge, which we refer to as
[Table]
data fusion , arises in diverse applications including climate monitoring (Gettelman et al., 2022), air quality estimation (Ferrer-Cid et al., 2020), land cover mapping (Robinson et al., 2021), population estimation (Lu et al., 2010), and hydrology (Hosseini & Kerachian, 2017).
NPs naturally handle multiple sources of both on-grid and off-grid data, and can make predictions at arbitrary locations (Vaughan et al., 2022; Markou et al., 2022), making them particularly well suited to data fusion.
Joint and non-Gaussian statistics.
Another challenge in environmental sciences is estimating the joint probability of events that span space and time.
For example, a renewable energy provider may want to know the compound risk of both low wind speeds at their offshore wind farm and high cloud coverage over their solar panel farm, possibly at different locations.
For such tasks, it is essential to model statistical dependencies to obtain well-calibrated uncertainty estimates.
Furthermore, many environmental variables, such as cloud cover, are bounded in output space and are poorly modelled by Gaussian predictives.
CNPs are most commonly applied with Gaussian likelihoods, but can be easily used with arbitrary likelihood functions.
While GAN-based approaches (Ravuri et al., 2021) can capture both joint as well as non-Gaussian statistics, they are unsuitable for integrating off-grid with on-grid data, as may be required for data fusion.
The AR ConvCNP thus fills a gap in the environmental modelling toolbox by enabling data fusion while capturing statistical dependencies and working with non-Gaussian likelihoods.
Here, we assess the AR ConvCNP on two common environmental modelling tasks, namely data assimilation and statistical downscaling .
Data assimilation.
Data assimilation is the task of combining observations of the Earth system to produce predictions on a regular grid, called a reanalysis .
Reanalyses are typically generated by fitting the trajectories of physics-based climate models to observations (Hersbach et al., 2020;
Gettelman et al., 2022).
AR ConvCNPs can not only provide a substantial speedup over conventional data assimilation systems, but also permit rigorous uncertainty estimation.
To explore the ability of the AR ConvCNP to emulate a data assimilation system for a non-Gaussian variable, we train convolutional NP models to predict simulated daily-average cloud cover fraction over Antarctica.
We use reanalysis data from ECMWF ERA5 (Hersbach et al., 2020) as ground truth.
Cloud cover ranges in [0 , 1] , with observations frequently taking values of 0 or 1 (Figure 13).
We evaluate the performance of NPs using either a Gaussian likelihood or a more appropriate beta-categorical mixture model with three components.
Two of these are discrete components that capture values of exactlyor 1, while the last is a beta distribution which captures continuous values in (0 , 1) .
This provides a robust way of handling 0 and 1 values, unlike the existing Copula GNP model (Markou et al., 2022) which can have its output constrained in (0 , 1) but places zero density at the endpoints.
Results.
In Table 5 we see that the AR ConvCNP significantly outperforms competing NPs for both the Gaussian and beta-categorical likelihoods.
Figure 6 shows samples drawn from the models, after observing context points on half of the space.
The AR ConvCNP displays remarkable ability to extrapolate rich, non-stationary, multi-scale structure, such as sudden changes in cloud cover over the Ross Ice Shelf coastline at the bottom of the figure.
By comparison, the ConvLNP and
ConvGNP produce blurry, lower frequency samples.
Unlike GPs, convolutional NP models have a fixed receptive field induced by the CNN architecture used for the encoder, which is computationally expensive to increase.
Away from the context points on the left, samples from the non-AR models will be independent of the observations, reverting to some mean representation of the data (Fig.
6c-e).
This highlights a further benefit of AR CNPs: successive AR applications increase the receptive field, enabling rich, conditional sample structure to extrapolate far away from observed data.
Further commentary and model samples are provided in Appendix K.3.
See Appendix K for full details.
Environmental downscaling.
The spatial resolutions of physics-based reanalyses are limited by their vast computational demands, making them unsuitable for capturing local and extreme events (Stocker et al., 2013; Maraun et al., 2017).
Statistical downscaling addresses this issue by leveraging additional information to produce fine-grained predictions (Maraun & Widmann, 2018).
Recently,
NPs have been shown to outperform a large ensemble of existing climate downscaling approaches
(Vaughan et al., 2022).
We compare the AR ConvCNP to the MLP ConvCNP of Vaughan et al.
and the MLP ConvGNP of (Markou et al., 2022) in a temperature downscaling task over Germany.
In this task, the context data consist of low-resolution ERA-Interim reanalysis data and high-resolution topography, and the target data consist of weather station observations from the ECA&D dataset.
We also consider a second setup where we reveal some station observations to aid the downscaling process.
As Appendix L.2 explains, the MLP ConvCNP and MLP ConvGNP cannot be naively extended to include these station observations.
We therefore introduce a novel multiscale architecture, which we use to run the ConvCNP in AR mode.
See Appendix L for full experimental details.
Results.
In Table 6 we observe that the AR ConvCNP matches the performance of the ConvGNP, which is remarkable as the latter has been previously demonstrated to outperform a range of stateof-the-art downscaling approaches (Markou et al., 2022; Vaughan et al., 2022).
When additional observations from weather stations are revealed, the AR ConvCNP significantly outperforms the
MLP ConvGNP in both metrics.
D OWNSCALING
[Table]
L IMITATIONS AND C ONCLUSION
We have shown that the AR procedure can be readily applied to improve the performance of CNPs, producing coherent samples and dramatically improved likelihoods.
Surprisingly, in an extensive range of experiments, this simple approach often outperforms more complicated methods which rely on latent variables or that explicitly model correlations.
We demonstrate the effectiveness of our approach on data sets of real-world interest by applying AR CNPs on climate data fusion tasks, modelling [0 , 1] -constrained data with a beta-categorical likelihood and introducing a novel multiscale architecture.
Notably, AR CNPs fill a gap in the climate modelling toolbox by enabling joint, non-Gaussian predictives, which could be used to better estimate the magnitude of compound risks.
We also position AR CNPs within the larger neural density estimator literature, showing the fruitfulness of combining NPs with other modelling paradigms.
More generally, AR CNPs equip the NP framework with a new knob where modelling complexity and computational expense at training time can be traded for computational expense at test time.
In particular, the higher quality samples and better likelihoods obtained by applying NPs autoregressively come with the additional cost of performing a forward pass for every element in the target set.
This can be prohibitively expensive for large target sets, and constitutes the primary practical drawback of using AR CNPs.
In addition, since AR CNPs do not define a consistent stochastic process, design choices for the AR procedure may affect the quality of the results.
Thus practitioners need to avoid choosing target sets that lead to pathological behaviour, such as when the spatial density of the target inputs is too high.
However, the flexibility of this design space also presents an opportunity: as an example, in Appendix M we show that auxiliary target points can be used to further improve predictions.
Finally, promising avenues for future work include applying the AR procedure to other
NPs besides CNPs, and investigating the efficacy of the block sampling scheme presented in Section 2.
R EPRODUCIBILITY S TATEMENT
All our experiments are carried out using either synthetic or publicly available datasets.
The EEG data set is available through the UCI database, 2 and the environmental data are also publicly available through the European Climate Data Service.
We make the code necessary to reproduce our experiments 4 as well as instructions for obtaining, preprocessing and modelling the Antarctica data 5 publicly available.
Proofs for Propositions 2. and 2.2 are given in Appendix A and Appendix B respectively.
Details on the model architectures and the experimental setup can be found in Appendices F to H for the synthetic datasets, Appendix I the sim-to-real transfer experiments, Appendix J for the EEG experiments, Appendix K for the data assimilation experiment, and Appendix L for the downscaling experiment.
E THICS S TATEMENT
Training CNPs autoregressively improves their performance dramatically, but we do not foresee adverse societal impacts as a result of this work.
That being said, the problem of capturing the statistical trends present in a dataset must be performed with care, especially in safety critical applications, where the stakes of making incorrect and confident predictions can have severe consequences.
We view the AR procedure as a useful tool, rather than a panacea, for capturing such behaviours, and hope this work encourages further research into building effective but reliable models to this end.
We also note that while training CNPs is computationally cheaper than alternative NP models, AR sampling itself incurs a substantial computational cost, and thus energy cost, at test time.
Running
AR sampling on a large scale could lead to greater power demands for these models, resulting in larger carbon footprints which are undesirable.
However, we believe the potential benefits for environmental modelling could outweigh this cost, while leveraging methods to make AR CNPs more computationally efficient should help alleviate this issue.
A P ROOF OF P ROPOSITION 2.
Additional notation.
If y 1 ⊕ y 2 ∼ P x 1 ⊕ x 2 π ( D ) , then denote the distribution of y 1 | y 2 by P x 1 | x 2 π ( D ) .
Note that P x 1 | x 2 π ( D ) depends on y 2 , because it is the distribution of y 1 | y 2 , even though the notation does not make this dependence explicit.
The “appropriate regularity conditions”.
Let P Nλ be the collection of distributions on R N that (a) have a density with respect to the Lebesgue measure and (b) have a covariance matrix which is strictly positive definite.
Let P Nλ ,G ⊆
P Nλ be the subcollection of distributions which are Gaussian.
Then, by
Corollary B.1 by Bruinsma et al.
(2021), for all µ ∈ P
Nλ such that inf
[Equation]
where N ( µ ) denotes the Gaussian distribution with mean vector and covariance matrix equal to those of µ .
In the proposition, by appropriate regularity conditions on y , we mean the assumption that, for all inputs x and D ∈ D , P x π y ( D ) is in P | x | λ and such that inf ν ∈P | x | λ ,G KL
( P x π y ( D ) , ν ) < ∞ .
Assume the appropriate regularity conditions on y .
We now list three technical observations.
(1) Note that P x 1 | x 2 π y ( D ) is the distribution of y ( x 1 ) | D, ( x 2 , y 2 ) , so we have the identity P x 1 | x 2 π y ( D ) =
P x 1 π y ( D ⊕ ( x 2 , y 2 )) .
Therefore, for all inputs x 1 , inputs x 2 , and D ∈ D , P x 1 | x 2 π y ( D ) is in P | x 1 | λ and such that inf ν KL( P x 1 | x 2 π y ( D ) , ν )
[Equation]
∞ .
The ideal CNP π C matches the means and marginal variances of the true posterior predictives (Section 2).
Hence, for all x ∈ X and D ∈ D , P x π C ( D ) is in P 1 λ, G .
The ideal GNP π G matches the mean vectors and covariance matrices of the true posterior predictives (Section 2).
Hence, for all inputs x and D ∈ D , P x π G ( D ) is in P | x | λ, G ; which means that, for all x 1 ∈ X , inputs x 2 , and D ∈ D , P x 1 | x 2 π G ( D ) is in P 1 λ, G .
In the proof, to apply and use (10), we implicitly use these observations.
Proposition 2.1 (Advantage of AR CNPs over GNPs) .
Assume appropriate regularity conditions on y .
Let π C be the ideal CNP and let π G be the ideal GNP.
Then, for all inputs x and data sets D ∈ D ,
[Equation]
Proof of Proposition 2.1.
Let x be some inputs and let D ∈ D be some data set.
We will argue that, for all n = 1 , . . .
, | x | ,
[Equation]
Assuming this inequality, the result follows directly from the chain rule for the KL divergence in combination with the definition of AR x
(Procedure 2.1):
[Equation]
where the expectations are over y 1:( n − 1) ∼ P x 1:( n − 1) π y ( D ) .
To prove the inequality, note that, conditional on y 1:( n − 1) , using (10),
[Equation]
By the property of π C that it matches the mean and marginal variance of the true posterior (Section 2),
N
[Equation]
Therefore, arg
[Equation]
Noting that P x n | x 1:( n − 1)
π G ( D )
∈ P 1 λ ,G , we obtain the desired inequality.
B P ROOF OF P ROPOSITION 2.
Proposition 2.2 (Recovery of smooth samples) .
Let X ⊆
R be compact, and let f be a stochastic process with surely continuous sample paths and sup x ∈X ∥ f ( x ) ∥
L 2 < ∞ .
Let ( ε n ) n ≥ 0 be i.i.d.
(potentially non-Gaussian) random variables such that E [ ε 0 ]
= 0 and V [ ε 0 ] < ∞ .
Consider any sequence ( x n ) n ≥ 1 ⊆ X , and let x ∗ ∈
X be a limit point of ( x n ) n ≥ 1 .
If y ( x ∗ ) =
f ( x ∗ ) + ε 0 and y n = f ( x n ) + ε n are noisy observations of f , then
[Equation]
almost surely
[Equation]
Proof of Proposition 2.2.
Consider the increasing filtration F n = σ ( y 1 , . .
y n ) with limit F ∞ =
σ ( (cid:83) ∞
n =1 F n ) .
Also let T n = σ ( ε n +1 , ε n +2 , . . . ) and consider the tail σ -algebra T = (cid:84) ∞ n =1 T n .
Let ( x n i ) ∞
i =1 be a subsequence of ( x n ) ∞ n =1 such that x n i → x ∗ .
Let g n = 1 n (cid:80) ni
=
y i .
Since g n is a function of y 1 , . .
, y n , it is F n –measurable and therefore F ∞ –measurable.
Note that
[Equation]
By sure continuity of f , the first term converges to f ( x ∗ ) surely.
By the strong law of large numbers (Example 5.6.1; Durrett, 2010), the second term converges to zero on a tail event A ∈ T of probability one.
We conclude that 1 A f ( x ∗ ) is σ ( F ∞ , T ) –measurable.
Therefore, by almost sure convergence of
L 2 –bounded martingales (Theorem 5.4.5; Durrett, 2010),
[Equation]
where all equalities hold almost surely.
C I LLUSTRATION OF THE AR PROCEDURE
Figure 7 depicts the AR sampling procedure (Procedure 2.1) and procedure to produce smooth samples (Proposition 2.2) using the ConvCNP trained on the EQ data process from Section 4.
Model fit
Step 1: Draw noisy samples using AR sampling (Procedure 2.1) ...
Step 2: Denoise sample by passing it through the model (Proposition 2.2)
[Figure]
D N UMBER AND O RDER OF T ARGET P OINTS
When deploying a conditional neural process (CNP) autoregressively (AR; Procedure 2.1), the number and ordering of the target points matters.
In this appendix, we describe our observations of the effects of the number and ordering of the target points on the quality of the predictions.
In short, our recommendation is to choose a different random ordering for every sample, and to not let the number/density of target points exceed that at training time.
D. E FFECTS OF THE N UMBER OF T ARGET P OINTS
During the AR sampling procedure, the AR CNP is evaluated at context sets of increasing size.
Our experience is that, as long as the sizes of these context sets do not exceed the sizes seen at training time, the predictions should not be significantly affected by changes in the number of target points.
However, if the AR sampling procedure evaluates the model at context sets of larger sizes than seen during training time, then that presents the model with an out-of-distribution situation.
What happens then comes down to how well the neural networks generalise.
Our experience is that the predictions quickly start to break down.
A notable exception of this rule of thumb are convolutional-deep-set–based models, such as the
Convolutional Conditional Neural Process (ConvCNP; Gordon et al., 2020).
For these models, the magnitude of the density channel is what determines whether the models generalises or not.
This means that it is not the total number of points that matters, but rather the density of the points.
Therefore, the AR ConvCNP can be evaluated at arbitrarily many target points, as long as the density of these points does not significantly exceed the density of context points seen at training time.
Once the density exceeds the density of the training data, the model is presented with an out-of-distribution situation, and what happens then again comes down to how well neural networks generalise.
Figure Appendix D.1 illustrates this observation.
When the density target points does not exceed the training data (50 and 100 points), the predictions look calibrated.
However, once the density of target points comes close or exceed the training data (200, 500, and 1000 points), bias starts to creep into the predictions.
Although the number/density of points in the AR sampling procedure should not exceed that at training time, AR CNPs can still produce high-quality samples at arbitrarily many target points by following the trick outlined at the end of the two-step procedure below Proposition 2.2.
D. E FFECTS OF THE O RDERING OF T ARGET P OINTS
Our experience is that, as long as the number of target points (or density) does not exceed that at training time, the ordering of the target point does not really matter.
Appendix D.1 also demonstrates this.
When the density of the target points does not exceed the training data (50 and 100 points), sampling randomly or left to right does not really matter.
However, once the density of the target points comes close to or exceeds the training data (200, 500, and 1000 points), we observe a difference in performance between sampling randomly and sampling left to right.
Across all numbers of target points, a random ordering seems to perform most robustly.
Our recommendation is therefore to choose a different random ordering of the target points for every sample.
D. A NALYSIS OF AR CNP S FOR CNP S WITH G AUSSIAN M ARGINALS
In this subsection, we argue that, for CNPs with Gaussian marginals, predictions in the first few AR steps might be poor, but predictions in later AR steps tend to be more accurate.
Choosing a different random ordering for every sample therefore “averages out” the effects from these first few AR steps.
When evaluating a CNP with Gaussian marginals in AR mode, every conditional prediction in the
AR process is Gaussian.
Let us consider the process of producing an AR sample.
For the first target input x 1 , we run the CNP forward to obtain a distribution for the corresponding target output y 1 .
In reality, the true posterior most likely is non-Gaussian, which means that the prediction for the first target point may be poor.
Nevertheless, we sample this Gaussian, append the sample ( x 1 , y 1 ) to the context set, and run the CNP forward again.
Because we now feed the earlier sample y 1 through the non-linear network, the marginal predictive for the next target output y 2 (having integrated out y 1 )
[Figure]
is non-Gaussian.
As we perform more AR steps, the marginal predictions of later points become increasingly non-Gaussian, increasing the model’s flexibility.
We see that, for a given ordering of the target inputs, the prediction for the first target input is likely poor (because it is Gaussian), and (in the best case) the predictions become more and more accurate as we take more AR steps (because they become more and more non-Gaussian).
This is exactly what is happening in Figure 3: the left prediction is Gaussian and therefore a poor approximation, and, as we go to the right and take more and more AR steps, the prediction becomes more and more non-Gaussian and therefore more accurate.
If we were to feed the target inputs in right to left, then the same phenomenon would happen.
The right prediction would be a Gaussian and a very poor approximation, and, as we go to the left and take more AR steps, the prediction would become more non-Gaussian and therefore more accurate.
More generally, for a given ordering of the target points, the ordering will produce high quality predictions if the conditional distributions of the AR factorisation match the corresponding conditional distributions of the true posterior.
Since the conditionals of the AR CNP are typically Gaussian by design, this means that the ordering is “good” if the corresponding conditionals of the true posterior are close to Gaussian.
So when is a conditional of the posterior close to Gaussian?
Let us assume that the true underlying process is a sum of a non-Gaussian process (constituting epistemic uncertainty) and independent
Gaussian noise (constituting aleatoric uncertainty).
Generally, a conditional will have both epistemic and aleatoric uncertainty, so a Gaussian will be a bad fit.
However , as we condition the conditionals of the true generative process on more and more data, the underlying function will be pinned down more and more accurately, meaning that the conditional will consist mostly of aleatoric uncertainty, which is Gaussian.
Therefore, as we condition on more and more data, we expect the conditionals to become more and more Gaussian.
This again suggests that the samples in the first few AR steps might be a poor fit (because the corresponding conditionals of the true posterior are not yet Gaussian), but that samples in later AR steps should be a better fit (because the corresponding conditionals are then close to Gaussian).
To summarise, an ordering of the target points is “good” if the corresponding conditionals of the true posterior are also close to Gaussian.
Under the assumption that the ground-truth process is a non-Gaussian process with additive Gaussian noise, conditionals tend to be close to Gaussian if they are conditioned on many data points.
As a consequence, the earlier conditionals in the AR factorisation tend to be poor fits to the ground-truth posterior, whereas later conditionals tend to produce better fits.
Choosing a different random ordering for every sample therefore “averages out” the effects from the first few AR steps.
E D ETAILS FOR F IGURE
The generative process visualised in the top panel of figure 3 is defined by the following mixture distribution:
[Equation]
Given this mixture distribution, the (Gaussian) ideal CNP can be computed in closed form by computing the first two moments of p true :
[Equation]
where
[Equation]
The updated mixture weights for the posterior distribution p true ( y | x, D (c) ) given a context set D (c) can be computed via Bayes rule and p CNP ( y | x, D (c) ) can be computed given the updated mixture weights.
Note that in figure 3 the prior mixture weights are a 1 = a 3 = 0 .
25 and a 2 = 0 .
5 , means are given by
[Equation]
and the target locations are x = 1 , 2 , 4 , and 6 .
The bottom four panels of figure 3 show squared exponential kernel density estimates of 30 000 samples drawn from the generative distribution p true ( y 1 , y 2 , y 4 , y 6 ) , the ideal CNP p CNP ( y 1 , y 2 , y 4 , y 6 ) , and the ideal CNP applied in AR mode from left to right
[Equation]
F D ESCRIPTION OF M ODELS
The architectures follow the descriptions from the respective papers they are introduced.
Although these descriptions are for one-dimensional inputs and outputs, the architectures are readily generalised to multidimensional inputs and outputs; we will explicitly mention wherever that generalisation requires extra care.
All architectures use ReLU activation functions.
All GNPs, in addition to a covariance matrix over the target points, also output heterogeneous observation noise along the marginal means; the total covariance over the target points is thus the sum of the covariance by the model and a diagonal matrix formed from these observation noises.
Conditional neural process (CNP; Garnelo et al., 2018a).
Set the dimensionality of the encoding to
K = 256 .
Parametrise the encoder with a three-hidden-layer multi-layer perceptron (MLP) of width 256 ; and parametrise the decoder with a six-hidden-layer MLP of width 256 .
For multidimensional outputs, let the decoder have width 512 .
For multidimensional outputs where outputs can have context points at different inputs, produce a separate encoding for every output and concatenate these into one big encoding.
These encoders may or may not share parameters.
In our experiments, for two-dimensional outputs, parametrise separate encoders; for higher-dimensional outputs, apply the same encoder.
Gaussian neural process (GNP; Markou et al., 2022).
Use the same choices for K , the encoder, and the decoder as the CNP.
Set the rank of the kernel map to R = 64 .
As mentioned in the introduction, let the decoder produce one extra dimension which forms heterogeneous observation noise.
For multidimensional outputs, the same caveats as for the CNP apply.
Latent neural process (LNP; Garnelo et al., 2018b).
The LNP builds off the CNP.
Call the existing encoder the deterministic encoder .
The NP adds one more encoder called the stochastic encoder .
The stochastic encoder mimics the deterministic encoder, but outputs a K -dimensional vector of means and a K -dimensional vector of marginal variances.
These are used to sample a
K -dimensional Gaussian latent variable (the stochastic encoding ).
The decoder now additionally takes in the stochastic encoding.
For multidimensional outputs, the same caveats as for the CNP apply.
Attentive conditional neural process (ACNP; Kim et al., 2019).
The ACNP builds off the CNP.
It replaces the deterministic encoder enc θ : D → R K with an eight-head attentive encoder enc (att) θ : D ×
X → R K (Vaswani et al., 2017).
Unlike the original deterministic encoder enc θ , the new attentive encoder enc (att) θ also takes in the target input.
Let D (c) =
( x (c) , y (c) ) ∈
D be a context set of size N and let x (t) ∈
X be a target input.
We now descibe the computation of enc (att) θ ( D (c) , x (t) ) .
Parametrise ϕ x : X → ( R 32 ) 8 and ϕ xy : X
× Y → ( R 32 ) 8 both with three-hidden-layer MLPs of width 256 .
Compute the keys : ( k h,n ) 8 h =1 = ϕ x ( x (c) n ) for n = 1 , . . .
, N, (35) the values : ( v h,n ) 8 h =1 = ϕ xy ( x (c) n , y (c) n ) for n
= 1 , . .
, N, (36) the query : ( q h ) 8 h =1 =
ϕ
x
( x (t) ) .
Then compute v (q) h = N (cid:88)
[Equation]
Concatenate v (q) =
( v (q) 1 , . . .
, v (q) 8 ) ∈
R 256 and q = ( q 1 , . .
. , q 8 ) ∈ R 256 .
Let L : R 256 → R 256 be a linear layer; let ϕ (res) : R
→ R 256 be a one-hidden-layer MLP of width 256 ; and let norm 1 and norm 2 be two layer normalisation layers with learned pointwise transformations (Ba et al., 2016).
Then enc
[Equation]
) where z = norm
[Equation]
For multidimensional outputs, the same caveats as for the CNP apply.
Attentive Gaussian neural process (AGNP).
The AGNP build off the GNP.
It replaces the deterministic encoder with the same eight-head attentive deterministic encoder of the ACNP.
Attentive neural process (ALNP; Kim et al., 2019).
The ALNP build off the LNP.
It replaces the deterministic encoder with the same eight-head attentive deterministic encoder of the ACNP.
Convolutional Conditional Neural Process (ConvCNP; Gordon et al., 2020).
Set the discretisation to an evenly spaced grid at a certain density (the points per unit ) spanning a bit more (the margin ) than the most extremal context and target inputs.
The points per unit and margin are specified separately for every experiment.
Initialise the length scales of all Gaussian kernels to twice the interpoint spacing of the discretisation.
Divide the data channel by the density channel.
Parametrise dec θ with a
U-Net (Ronneberger et al., 2015).
Before the U-turn, let the U-Net have six convolutional layers with kernel size five, stride two, and 64 output channels; and six more such layers, but using transposed convolutions, after the U-turn.
The layers after the U-turn additionally take in the outputs of the layers before the U-turn in reversed order; this is the U-net structure (Figure 1; Ronneberger et al., 2015).
For multidimensional outputs where outputs can have context points at different inputs, produce a separate data and density channel for every output and concatenate these into one big encoding; use separate length scales for every application of enc θ .
Convolutional Gaussian neural process (ConvGNP; Markou et al., 2022).
Use the same choices for the discretisation, length scales, and CNN architecture as for the ConvCNP.
Set the rank of the kernel map to R = 64 .
As mentioned in the introduction, let the decoder produce one extra channel which forms heterogeneous observation noise.
For multidimensional outputs, the same caveat as for the ConvCNP applies.
Fully convolutional Gaussian neural process (FullConvGNP; Bruinsma et al., 2021).
For the mean architecture and the kernel architecture, use the same choices for the discretisation, length scales, and CNN architecture as for the ConvCNP.
Implement the source channel with the identity matrix and apply the matrix transform Z (cid:55)→ ZZ T to ensure positive definiteness.
Let the decoder produce one extra channel which forms heterogeneous observation noise.
For multidimensional outputs, in addition to the caveat for the ConvCNP, two additional caveats apply.
First, for D odimensional outputs, let the decoder produce D 2 o channels rather than just one.
These channels should be interpreted as all covariance and cross-covariance matrices between all outputs.
Second, when applying the matrix transform Z (cid:55)→ ZZ T , these channels should first be assembled into one total covariance matrix.
Convolutional latent neural process (ConvLNP; Foong et al., 2020).
The ConvLNP builds off the
ConvCNP.
The ConvLNP replaces the CNN architecture by two copies of this architecture placed in sequence.
In between the two architectures, there is a sampling step: the first architecture outputschannels, comprising 16 means and 16 marginal variances, which are used to sample a 16-dimensional
Gaussian latent variable; and the second architecture then takes in this sample.
Autoregressive Conditional Neural Processes (AR CNPs).
The AR CNP, AR ACNP, and AR
ConvCNP use the architectures described above.
Rolling out an AR CNP according to Procedure 2. requires an ordering of the target points.
In all experiments, we choose a random ordering of the target points.
G T RAINING , C ROSS -V ALIDATION , AND E VALUATION P ROTOCOLS
The following description applies to the synthetic experiments (Section 4.1), the predator–prey experiments (Section 4.2), the EEG experiments (Section 4.3), and the environmental downscaling experiments (Section 4.4).
For the environmental data assimilation experiments, a different protocol was used; we refer the reader to Appendix K for full details of the environmental data assimilation experiments.
A task consists of a context set and target set.
How precisely the context and target sets are generated is specific to an experiment.
To train a model, we consider batches of 16 tasks at a time, compute an objective function value, and update the model parameters using ADAM (Kingma & Ba, 2015).
The learning rate is specified separately for every experiment.
We define an epoch to consist of 2 14 ≈ 16 k tasks.
We typically train a model for between 100 and 1000 epochs.
For an experiment, we split up the meta–data set into a training set , a cross-validation set , and an evaluation set .
The model is trained on the training set.
During training, after every epoch, the model is cross-validated on the cross-validation set.
Cross-validation uses 2 12 fixed tasks.
These 2 12 are fixed, which means that cross-validation always happens with exactly the same data.
The cross-validation objective is a confidence bound computed from the model objective.
Suppose that model objective over all 2 12 cross-validation tasks has empirical mean ˆ µ and empirical variance ˆ σ 2 .
If a higher model objective is better, then the cross-validation objective is given by ˆ µ − 1 .
96 · ˆ σ/ √ 2 12 .
The model with the best cross-validation objective is selected and used for evaluation.
Evaluation is performed with the evaluation set and also uses 2 12 tasks.
Conditional neural processes and Gaussian neural processes are trained, cross-validated, and evaluated with the neural process ELBO objective proposed by Garnelo et al.
(2018a).
We normalise the terms in the neural process objective by the target set sizes.
Latent-variable neural processes (LNPs) are trained, cross-validated, and evaluated with the ELBO objective proposed by Garnelo et al.
(2018b) using five samples, also normalised by the target set size.
When training LNPs with the ELBO objective, but not when cross-validating and evaluating, the context set is subsumed in the target set.
Additionally, LNPs are trained, cross-validated, and evaluated with the ML objective proposed by Foong et al.
(2020), again normalised by the target set size.
When training and cross-validating
LNPs with the ML objective, we use twenty samples; and when evaluating, we use 512 samples.
For completeness, LNPs trained with the ELBO objective are also evaluated with the ML objective using 512 samples.
To stabilise the numerics for GNPs, we increase the regularisation of covariance matrices for one epoch.
To encourage LNPs to fit, we fix the variance of the observation noise of the decoder to 10 −for the first three epochs.
H D ETAILS OF S YNTHETIC E XPERIMENTS H. D ESCRIPTION OF E XPERIMENTS
We synthetically generate data sets by randomly sampling from five different choices for the groundtruth stochastic process f .
Let the inputs be d x -dimensional.
Then define the following stochastic processes:
EQ: a Gaussian process with an exponentiated quadratic (EQ) kernel:
[Equation]
where ℓ > 0 is a length scale; Matérn– 52 : a Gaussian process with a Matérn– 52 kernel:
[Equation]
where k ( r ) =
(1 + √ 5 r + 53 r 2 )
e − r and ℓ > 0 is a length scale; weakly periodic: a Gaussian process with a weakly periodic kernel:
[Equation]
where ℓ d > 0 is a length scale specifying how quickly the periodic pattern changes, ℓ p > 0 a length scale of the periodic pattern, and p > 0 the period; and where the application of sin is elementwise; sawtooth: a sawtooth process with a random frequency, direction, and phase:
[Equation]
where ω ∼ Unif(Ω) is the frequency of the sawtooth wave, u ∼ Unif( { x ∈ R d x : ∥ x ∥ 2 = 1 } ) the direction, and ϕ ∼ Unif([0 , 1]) the phase; mixture: with equal probability, sample f from the EQ process, Matérn– 52 process, weakly periodic process, or sawtooth process.
We will call these stochastic processes the data processes .
The data processes are stochastic processes with d x -dimensional inputs and one-dimensional outputs.
We will turn them into processes with d ydimensional outputs according to the following procedure: sample from the one-dimensional-output prior d y times; and, for these d y samples, take d y different linear combinations.
We choose the parameters of the data processes based on the input dimensionality d x :
[Equation]
with c = √ d x .
Scaling with the input dimensionality aims to roughly ensure that data with onedimensional inputs and data with two-dimensional inputs are equally difficult.
Figure 9 illustrates the sawtooth data process in all four configurations.
We will construct data sets by sampling inputs uniformly at random from X =
[ − 2 , 2] d x and then sampling outputs from one of the data processes.
We will colloquially call X the training range .
For the EQ, Matérn– 52 , and weakly periodic process, but not for the sawtooth process 6 , we also add independent Gaussian noise with variance 0 .
The numbers of context and target points are as follows.
For the EQ, Matérn– 52 , and weakly periodic process, the number of context points is chosen[-
 The sawtooth process is already challenging enough.-] uniformly at random from { 0 , . .
, 30 · d x } and the number of targets points is fixed to 50 · d x .
For the sawtooth and mixture process, the number of context points is chosen uniformly at random from { 0 , . .
, 30 } if d x = 1 and { 0 , . . .
, 75 · d x } otherwise; and the number of targets points is fixed to· d x .
In the case of a multidimensional-output data process, we separately sample the number and positions of the context and target inputs for every output dimension.
For every data process and each of the four configurations, we evaluate every model in three different ways.
First, we evaluate the model on data generated exactly like the training data.
This task is called interpolation and abbreviated “int.”
in the tables of results.
The interpolation task measures how well a model fits the data and is the primary measure of performance.
Second, we evaluate the model on data with inputs sampled from [2 , 6] d x .
This task is called out-of-input-distribution (OOID) [Figure] [Figure] [Figure]
2 x 1.
[Figure] 2 x 1. [Figure] [Figure]
[Figure]
[Figure]
interpolation and abbreviated “OOID” in the tables of results.
OOID interpolation measures how well a model generalises to data sampled from other regions of the input space.
Third, we evaluate the model on data with context inputs sampled from [ − 2 , 2] d x and target inputs sampled from [2 , 6] d x .
This task is called extrapolation and abbreviated “ext.”
in the tables of results.
The extrapolation task measures how well predictions based on data in the training range generalise to other regions of the input space.
For this experiment, the learning rate is 3 · 10 − 4 , the margin is 0 . 1 , and the points per unit is 64 .
We trained the models for 100 epochs.
Due to an error in the cross-validation procedure, we did not use cross-validation, but used the model at epoch 100.
For the kernel architecture of the FullConvGNP, we reduce the points per unit and the number of channels in the U-Net by a factor two.
For the ConvLNP with two-dimensional inputs, we reduce the number of outputs channels in the U-Net by a factor √ 2 ; and, for training and cross-validation, we reduce the number of samples of the ELBO objective to one and the number of samples for the ML objective to five.
H. M ULTI -M ODALITY OF P REDICTIONS BY AR C ONV CNP
Figure 10 demonstrates multi-modality of predictions by the AR ConvCNP trained on the sawtooth process.
Note that the prediction is bimodal for one and two observations, and collapses to a single mode upon observing the third observation.
H. F ULL R ESULTS
We the show the full results for all data sets and tasks in Tables 7 to 18.
The AR ConvCNP consistently shows very strong performance compared to other NP models.
Note that the FullConvGNP takes much longer to train than the ConvCNP (Figure 2), and cannot be applied to tasks with 2-dimensional input spaces. [Figure] n s i [Figure]
[Figure]
[Figure]
D e n s i t y
[Figure]
e n s i t y
[Figure]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
I D ETAILS OF S IM - TO -R EAL T RANSFER E XPERIMENTS I. D ESCRIPTION OF E XPERIMENT
Our goal will be to make predictions for the famous hare–lynx data set.
The hare–lynx data set is a time series from 1845 to 1935 recording yearly population counts of a population of Snowshoe hares and a population of Canadian lynx (MacLulich, 1937).
A digital version extracted from the original graph by MacLulich (1937) is available by Hundley.
7 Hundley, the author of this digital source, says that other authors caution that the hare–lynx data is actually a composition of multiple time series, and presents the data with that caution.
We, therefore, also present the data with this caution.
Figure 11a visualises the hare–lynx data set.
To make predictions for the hare–lynx data set, we use the Lotka–Volterra equations (Lotka, 1910;
Volterra, 1926), also called the predator–prey equations.
The Lotka–Volterra equations are an idealised mathematical model for the population counts of a prey population and a predator population: prey population:
[Equation]
[Equation]
predator
[Equation]
These differential equations say that the prey population naturally grows exponentially with rate α , and that the predator population naturally decays exponentially with rate δ .
In addition, the predators hunt the prey.
The resulting additional growth in the predator population and the resulting additional decrease in the prey population are both proportional to the product of the densities.
In this idealised mathematical form, the population counts converge to a smooth, noiseless limit cycle and then perfectly track this limit cycle ever after.
This is unlike real-world predator–prey population counts, which exhibit noisy behaviour and imperfect cycles.
We therefore consider a stochastic version of the
Lotka–Volterra equations, given by the following two coupled stochastic differential equations:
[Equation]
where W (1) and W (2) are two independent Brownian motions.
Compared to the Lotka–Volterra equations, (47) and (48) have two additional terms, σX νt d W (1) t and σY νt d W (2) t , which introduce noisy behaviour.
In these terms, multiplying by X νt and Y νt makes the noise go to zero when X t and
Y t become small, ensuring that X t and Y t remain positive.
In addition, we multiply by a parameter σ > 0 to control the magnitude of the noise, and we raise X t and Y t to a power ν > 0 to control how quickly the noise grows as X t and Y t grow.
Namely, X t naturally grows exponentially, so, by adding noise of magnitude proportional to X t , we risk large spikes in the prey population.
To moderate this behaviour, we choose ν to be strictly less than one.
Finally, to obtain a variety of magnitudes of population counts, we multiply the realisation with a scale η .
After simulating from (47) and (48) a few times, we settle on ν = 16 .
For the remainder of the parameters, we simply manually play around with (47) and (48), settle on parameter ranges that look reasonable, and randomly sample parameters from those intervals.
Table 19 summarises the sampling distributions for all parameters of (47) and (48).
Figure 11b shows four samples from the proposed stochastic model.
To generate a meta–data set, we simulate (47) and (48) on a dense grid spanning 110 years, throw away the first 10 years, and retain between 150 and 250 data points for X t and Y t .
The numbers of data points and the locations of the data points are sampled separately for X t and Y t .
Hence, whereas the hare–lynx data is regularly spaced and the populations are always simultaneously observed, our simulator generates data at arbitrary and nonsimultaneous points in time.
We split these data sets into context and target sets in three different ways.
To train the models, for every batch, we randomly choose one of the interpolation, forecasting or reconstruction tasks by rolling a three-sided die.
We will also perform these tasks on the real hare–lynx data; in that case, for interpolation, we let the number of target points per output be between one and fifteen.
The tasks on simulated and real data are similar, but slightly differ in the number of context and target points.
To deal with the positivity of population counts, we transform the marginals of all models to distributions on (0 , ∞ ) by pushing the marginals through x (cid:55)→ log(1 + x ) .
[Figure] [Figure]
[Table]
For this experiment, the learning rate is 1 · 10 − 4 , the margin is 1 , and the points per unit is 4 .
We trained the models for 200 epochs.
The convolutional models use a U-Net architecture with seven layers instead of six where, in the first layer, the stride is one instead of two.
For the kernel architecture of the FullConvGNP, we reduce the points per unit and the number of channels in the U-Net by a factor two.
[Table]
I. F ULL R ESULTS
In Table we present the full results for the sim-to-real experiments. [Figure]
J D ETAILS OF E LECTROENCEPHALOGRAPHY E XPERIMENTS
We explore an electroencephalography data set collected from 122 subjects (Begleiter, 2022).
There are two groups of subjects: alcoholic and control.
Every subject was subjected to a single stimulus or two stimuli, and their response was measured with 64 electrodes placed on a subject’s scalp.
These measurements are in terms of trials , where a trial consists of 256 samples of the electrodes spanning one second.
The data sets contains up to 120 trials for each subject.
The data is available at https://archive.ics.uci.edu/ml/datasets/eeg+database and the collection is described in detail by Zhang et al.
(1995a).
In this experiment, we focus only on seven frontal electrodes: FZ, F1, F2, F3, F4, F5, and F6.
Figure 12 illustrates a trial of a subject, showing the samples for these seven electrodes.
We randomly split all subjects into three sets: an evaluation set consisting of ten subjects, a crossvalidation set consisting of ten other subjects, and a training set consisting of all remaining subjects.
For each of these sets, we create a meta–data set by aggregating the trials for all subjects.
We split every trial into a context and target set in the same three ways as for the predator–prey experiment.
First, for all seven electrodes separately, randomly designate between 50 and 256 points to be the target points and let the remainder (between 0 and 206) be the context points.
This task is called interpolation and is the primary measure of performance.
Additionally, randomly choose one of the seven electrodes and, for that choice, split the data in two exactly like for forecasting.
For all other electrodes, append all data to the context set.
This task is called reconstruction and measures the model’s ability to infer a signal for one electrode from the others.
We train all models on the interpolation task, and evaluate the models on the interpolation and reconstruction task.
For this experiment, the learning rate is 5 · 10 − 5 , the margin is 0 . 1 , and the points per unit is 256 .
We trained the models for 1000 epochs.
For the FullConvGNP, the learning rate is 2 · 10 − 4 .
The training run for the FullConvGNP was terminated after 84 hours, reaching epoch 127.
The convolutional models use a U-Net architecture where, in the first layer, the stride is one instead of two.
In addition, the number of channels are adjusted as follows: the ConvCNP and ConvGNP use 128 channels, the ConvLNP uses 96 channels, and the FullConvGNP uses 64 channels.
The length scales of the Gaussian kernels of the convolutional model is initialised to 0 .
To scale to seven outputs, the deep set–based and attentive models reuse the same encoder for every output dimension.
K D ETAILS OF E NVIRONMENTAL D ATA A SSIMILATION E XPERIMENT
In this section we provide further details on the Antarctic cloud cover data assimilation experiment described in Section 4.4.
K. D ATA CONSIDERATIONS
Data sources.
Daily-averaged cloud cover reanalysis data was obtained from ERA5 (Hersbach et al., 2020).
An Antarctic land mask and elevation field was obtained from the BedMachine dataset (Morlighem, 2020).
Figure 13 shows an empircal density of the ERA5 cloud cover values calculated over the models’ training period of 2000-2013.
The spikes at 0 and 1 correspond largely to values of exactly 0 and exactly 1.
This motivates the beta-categorical likelihood described in Section 4.4.
Data preprocessing.
The cloud cover data and land/elevation auxiliary data were regridded from lat/lon to a Southern Hemisphere Equal Area Scalable Earth 2 (EASE2) grid at 25 km resolution and cropping to a size of 280 × 280.
This centres the data on the South Pole.
Data normalisation.
We normalised the data before passing it to the convolutional NP models.
The cloud cover and land mask data already took appropriate normalised values in [0 , 1] .
The elevation field was normalised from metres to values in [0 , 1] .
The input coordinates x were normalised from metres to take values in [ − 1 , 1] .
K. M ODEL CONSIDERATIONS
Here we provide details on the training procedure and architectures for each of the convolutional NP models in the Antarctic data assimilation experiment.
Generating the training, validation, and test tasks.
Following meta-learning principles, we collect data from day τ into a task D τ .
Each task D τ was generated by first drawing the integer number of simulated cloud cover context points N (c) ∼ Unif { 1 , 2 , . . .
Letting N (c) vary encourages the model to learn to deal with both data-sparse and data-rich scenarios.
The number of target points N (t) was fixed to a value of 2,000.
Next, the input locations x (c) τ and x (t) τ were sampled uniformly at random across the entire 280 ×input space and the corresponding y τ values were sampled without observation noise.
For the training dates, the random seed used for generating D τ is updated every epoch, allowing for an infinitely growing simulated training data set.
In contrast, for the validation and test dates, the random seeds were held fixed so that metrics computed over the validation and test sets are not stochastic.
Training procedure.
Each model was trained for 150 epochs on 14 years of data from 2000–2013.
An Adam optimiser was used with a learning rate of 5 × 10 − 5 and a batch size of 2.
For the loss functions we use a negative log-likelihood loss function for the ConvCNP and ConvGNP.
For the
ConvLNP we use the ELBO objective and fix the variance of the observation noise to 0 .
01 for the first four epochs.
Validation data from 2014–2017 was used for checkpointing the model weights using the per-datapoint predictive log-likelihood.
The two year period of 2018–2019 data was reserved for the test set.
The time taken to train each model on a Tesla V100 GPU is as follows:
• ConvCNP: 25 . 0 hours, • ConvGNP: 27 .
5 hours, • ConvLNP: 43 . 6 hours.
Architectures.
For each model, the U-Net component of the encoder uses 5 × 5 convolutional kernels with the following sequence of channel numbers (d.s.
= 2 × 2 downsample layer, u.s. =
× 2 upsample layer): 16 d.s. −−→ 32 d.s
. −−→ 64 d.s.
−−→ 128 u.s. −−→ 64 u.s. −−→ 32 u.s. −−→ 16 .
We use bilinear resize operators for the upsampling layers to fix checkerboard artifacts that we encountered when using standard zero-padding upsampling (Odena et al., 2016).
We use a margin of 0 . 1 and 150 points per unit for the encoder’s internal discretisation.
The length scales of the Gaussian kernels for both the encoder and decoder
SetConv layers are set to 1 / 150 and held fixed during training.
These architecture choices result in a receptive field of 0 .
433 in normalised input space, or roughly 1 .
500 km in raw input space, spanning around 20% of the region in Figure 6 in either dimension.
For the ConvGNP we use 128 basis functions for the low-rank covariance parameterisation described in Markou et al. 2021.
For the ConvLNP we use an 8-dimensional latent variable and evaluate the ML objective (Foong et al., 2020) using 8 latent samples.
The number of learnable parameters for each model is as follows:
• ConvCNP: 618 k , • ConvGNP: 621 k , • ConvLNP: 1 . 234 k (increase due to second UNet architecture after the latent variable).
The difference in parameters from switching to a beta-categorical likelihood from a Gaussian likelihood is negligible.
Input data.
Each model receives two context sets as input.
The first contains observations of the simulated ERA5 daily-average cloud cover.
The second contains a set of 6 gridded auxiliary variables.
These are elevation, land mask, cos(2 π × day of year / 365) , sin(2 π × day of year / 365) , x 1 , and x 2 .
The elevation and land mask auxiliary fields allow the models to predict spatial non-stationarity.
For example, the convolutional filters of the model’s encoder could learn how cloud cover around the
Antarctic coastline behaves differently to the centre of the continent.
The cos and sin variables inform the model at what time of year
D τ corresponds to, helping with learning seasonal variations in the data.
The x 1 and x 2 inputs again help with breaking translation equivariance in the convolutional filters by informing the model where in input space the data corresponds to. [Equation]
A NTARCTIC CLOUD COVER MODEL SAMPLES
Figure 14 gives a detailed breakdown of sample extrapolation ability, showing seven samples from the four Antarctic cloud cover models.
The AR ConvCNP samples display remarkable structure and variation while still closely interpolating the context points.
The samples also provide interesting scenarios in the gaps between the context points on the left hand side.
In contrast, the ConvCNP samples are incoherent, underestimating the probability of joint events.
The ConvLNP samples were generated by sampling from the latent variable and then computing the mean of the marginal distributions.
As is visible in Figure 14, the ConvLNP displays low sample variance with respect to the latent variable.
However, to the best of our knowledge we used a faithful reproduction of the original ConvLNP model, so we leave a more rigorous treatment of this undesirable behaviour to future work.
For all the non-AR models, the limited receptive field size leads to samples on the right hand side becoming independent of the context observations on the left hand side after roughly 750 km of distance from them.
This results in the models defaulting to some mean representation of the data.
It is interesting to see that all the non-AR models display similar marginal mean structure, with greater cloud cover towards the centre of the continent and lower cloud cover towards the coastline, followed by increased cloud levels over the Southern Ocean.
The AR samples were drawn on a sparse 70 x 70 grid spanning the entire input space to save compute time.
The ConvCNP model was then conditioned on these AR samples and the predictive mean was computed over the dense 280 x 280 target space.
It took 14 minutes to generate these AR ConvCNP samples on a Tesla V100 GPU.
L D ETAILS OF C LIMATE D OWNSCALING E XPERIMENTS L. D ESCRIPTION OF E XPERIMENT
The MLP ConvGNP (Markou et al., 2022) can be used to successfully model dependencies between outputs in a statistical downscaling task, improving log-likelihoods over the MLP ConvCNP
(Vaughan et al., 2022) and
enabling coherent samples.
In this experiment, we demonstrate that the AR ConvCNP can also be used for this purpose.
The goal of this experiment is to estimate the maximum daily temperature at 589 weather stations around Germany.
To estimate these temperatures, we follow Vaughan et al.
and use 25 coarsegrained ERA-Interim reanalysis variables (Dee et al., 2011) in combination with 1 km –resolution elevation data (Earth Resources Observation and Science Center, U.S. Geological Survey, U.S.
Department of the Interior, 1997).
We also consider a second setup where we reveal some of the weather station observations.
These revealed weather station observations can be used by the models to aid downscaling performance.
The ERA-Interim reanalysis variables considered are tabulated in Table 21.
In contrast to previous downscaling work, which degrade reanalysis data to between 2 ◦ and 2.5 ◦ , we opt to use the ERAReanalysis data at the native 0.75 ◦ resolution, consistent with the latest high-resolution climate models with horizontal resolution ranging from 0.5 ◦ to 1.0 ◦ .
All variables are spatially subset to between 6 ◦ to 16 ◦ longitude and 47 ◦ to 55 ◦ latitude, covering Germany.
The weather station data are a subselection from of the European Climate Assessment & Dataset (Tank et al., 2002) and are available at https://www.ecad.eu ; we use the blended data.
Like for the ERA-Reanaysis variables, we take the weather stations located within the aforementioned square.
The locations of the weather stations around Germany are visualised in Figure 15.
The 1 km –resolution elevation data is taken from the United States Geological Survey GTOPO30 elevation data set available at https://doi.org/10.5066/F7DF6PQS .
This provides global elevation data at 30-arc second resolution, which is approximately 1 km .
Following the VALUE framework (Maraun et al., 2015), we consider all days of the years 1979– and split these years into five folds.
We use the first four folds (spanning 1979–2003) for training, holding out the last 1000 days for cross-validation; and use the fifth fold (spanning 2003–2008) for evaluation.
L. M ULTISCALE C ONVOLUTIONAL A RCHITECTURE
Deploying the AR ConvCNP in this downscaling experiment comes with a significant challenge.
Because the elevation data has a fine resolution of 1 km , we expect that predictions by the AR
ConvCNP will vary roughly also on this length scale.
In the autoregressive sampling procedure (Procedure 2.1), samples from the model will be fed back into the model.
Therefore, the AR ConvCNP must handle context data which varies on a 1 km spatial scale, which means that the discretisation of the AR ConvCNP must roughly be a 1 km –resolution grid.
Unfortunately, making the discretisation this fine is prohibitively expensive and imposes prohibitive memory requirements.
It is this limitation that prevents us from extending the Vaughan et al.
(2022)’s
MLP ConvCNP and Markou et al.
(2022)’s
MLP ConvGNP to include additional weather station observations.
We must therefore innovate on the AR ConvCNP design to come up with a convolutional architecture that can handle such a fine discretisation at reasonable computational expense.
The architecture that we propose is a multiscale architecture operating on multiple spatial length scales.
Let us divide the context set D
= D lr ∪ D mr ∪ D hr into a low-resolution component D lr , a medium-resolution component D mr , and a high-resolution component D hr .
Let the low-resolution component D lr consist of features of the context set that vary on a long spatial length scale, the medium-resolution component D mr of features that vary on a medium-range spatial length scale, and the high-resolution component D hr of features that vary on a short spatial length scale.
The central assumption of the architecture is that predictions for target points depend on precise short-length-scale details D hr nearby, but that this dependence weakens as we move away from the target point, starting to depend more on broad-stroke long-length-scale components D lr .
For example, predictions might depend on detailed orographic information nearby, but more on general orographic shapes farther away.
[Table]
position in the year t transformed encoded with t (cid:55)→ (cos(2 πt
[Table]
, sin
[Table]
πt
[Table]
Figure 16 depicts the multiscale architecture.
The architecture is a cascade of three convolutional deep sets, parametrised by three CNNs; please see the caption.
The low-resolution CNN handles the context data D lr with a long spatial length scale.
Because these features have a long spatial length scale, the CNN can get away with a low-resolution discretisation.
The output of the low-resolution
CNN then feeds into a medium-resolution CNN.
The medium-resolution CNN handles the context data D mr with a medium spatial length scale and has a medium-resolution discretisation.
Finally, the output of the medium-resolution CNN feeds into a high-resolution CNN.
This CNN handles the context data D hr with a short spatial length scale and has a high-resolution discretisation
[Table]
The key to the computational efficiency of this architecture is that we construct the high-resolution discretisation only locally to the target points: a small square covering 0 .
25 ◦ more than the most extremal target points.
If the target points are confined to a small region, then the high-resolution grid will also be small, covering only 0 .
25 ◦ more than that region.
Crucially, the high-resolution
[Equation]
(cid:32)(cid:34)
[Equation]
density
[Equation]
(cid:35)(cid:33)
[Equation]
(cid:32)(cid:34)
[Equation]
data
[Equation]
density
[Equation]
(cid:35)(cid:33)
z
[Equation]
[Equation]
= CNN
[Equation]
(cid:18)(cid:20) data ( D lr ) density ( D lr ) (cid:21)(cid:19) 0 .
01 ◦ ( ✓ ), local to target points ( ✗ ) 0 .
75 ◦ ( ✗ ), covering all of Germany ( ✓ ) 0 .
[Figure]
( ✓ ), covering a medium-sized square ( ✓ )
[Figure]
[Figure]
grid will not be constructed over all of Germany, like it would if we were to more naively apply the ConvCNP with a high-resolution discretisation, incurring prohitive computational cost.
Even though the high-resolution grid is only constructed locally to the target points, the model can still capture long-range dependencies via the medium-resolution and low-resolution grids.
Namely, the medium-resolution grid is a square covering 5 ◦ more than the most extremal target points, and the low-resolution grid covers all of Germany; see Figure 16.
To utilise this computational gain, the target points must be confined to a small region.
This perfectly synergises with the autoregressive sampling procedure (Procedure 2.1), because this procedure evaluates the model one target point at a time.
The training procedure, however, must be adapted.
During training, we subsample the target points to ensure that the target set is always confined to a small square, which is described in
Appendix L.4.
During the autoregressive sampling procedure, the AR ConvCNP takes in earlier AR samples from the model.
In the architectures of the MLP ConvCNP and MLP ConvGNP, these is no natural context data to which these samples can be appended.
Therefore, in addition to the ERA-Interim reanalysis variables and the elevation data, we also let the AR ConvCNP take in observed weather stations as context data.
We will append the earlier AR samples to these weather station context data.
To have the model make appropriate use of the weather station context set, we must randomly divide the weather stations observations over the context and target set, which we describe in Appendix L.4.
We let the low-resolution context data D lr consist of the 25 coarse-grained ERA-Interim reanalysis variables, and let the medium-resolution D mr and high-resolution context data D hr both consist of the weather station observations (and earlier AR samples) and the 1 km –resolution elevation data.
When the 1 km –resolution data is fed to the medium-resolution CNN, the data loses some detail, because the internal discretisation of the medium-resolution CNN is coarser than the data; however, when it is fed to the high-resolution CNN, the data retains its detail.
The same holds for the weather station observations (and earlier AR samples).
L. A RCHITECTURES
MLP ConvCNP and MLP ConvGNP (Vaughan et al., 2022; Markou et al., 2022).
The MLP
ConvCNP and MLP ConvGNP are a respectively a ConvCNP and ConvGNP where the decoder dec θ = fuse θ ◦ dec ′ θ is decomposed into a convolutional architecture dec ′ θ followed by a pointwise
MLP fuse θ :
[Equation]
In this architecture, the ERA-Interim variables are incorporated via the convolutional architecture, producing the encoding z ( • ) .
On the other hand, as (49) shows, the 1 km –resolution elevation data is included via the pointwise MLP fuse θ .
Parametrise dec ′ θ with a seven-layer residual convolutional neural network (He et al., 2016).
Every residual layer involves one depthwise-separable convolutional filter (Chollet, 2017) with kernel size three followed by a pointwise MLP.
Every layer has 128 channels, and the network also outputschannels.
The discretisation for dec ′ θ is the grid of the ERA-Interim reanalysis variables.
Parametrise fuse θ with a three-hidden-layer MLP of width 128.
AR ConvCNP.
The AR ConvCNP does not use the pointwise MLP fuse θ to incorporate the 1 km – resolution elevation data.
Instead, it is a normal ConvCNP where the convolutional architecture is implemented by the multi-scale architecture described in Figure 16.
Parametrise CNN lr with a depthwise-separable residual convolutional neural network like in the MLP
ConvCNP and MLP ConvGNP, but use six layers instead of seven.
Let CNN lr output 64 channels.
The discretisation for CNN lr is the grid of the ERA-Interim reanalysis variables.
Parametrise CNN mr with a U-Net (Ronneberger et al., 2015) using an architecture similar to what we have been using.
Before the U-turn, let the U-Net have five convolutional layers with kernel size five, stride one for the first layer and stride two afterwards, 64 output channels for the first three layers andoutput channels afterwards.
After the U-turn, instead of using transposed convolutions, use regular convolutions combined with an upsampling layer using bilinear interpolation.
Let CNN mr output 64 channels.
The receptive field of CNN mr is approximately 10 ◦ .
The discretisation for CNN mr is centred around the target points with margin 5 ◦ .
Parametrise CNN hr with a U-Net like for CNN hr , but with four convolutional layers before the U-turn.
The receptive field of CNN hr is approximately 0 .
5 ◦ .
The discretisation for CNN hr is centred around the target points with margin 0 .
25 ◦ .
L. T RAINING D ETAILS
MLP ConvCNP and MLP ConvGNP.
The MLP ConvCNP and MLP ConvGNP are trained with learning rate 2 . 5 · 10 − 5 for 500 epochs.
For the MLP ConvGNP, to encourage the covariance to fit, we fix the variance of the decoder to 10 − 4 I for the first ten epochs.
AR ConvCNP.
The AR ConvCNP is trained with learning rate 1 · 10 − 5 for 500 epochs.
During training and cross-validation, the target points are subsampled to lie in a 3 ◦ × 3 ◦ square.
For training, the number of target points is ensured to be at least ten; and for cross-validation, at least one.
The size of the cross-validation set is increased ten fold.
Sampling of data.
For the MLP ConvCNP and MLP ConvGNP, since these architectures cannot take in weather station observations, all weather stations are used as context data.
For the AR ConvCNP, a data set is split into a context and target set by randomly selecting n points as context points and letting the remainder be target points.
Specifically, the number of context points n is sampled from p ( n ) ∝ e − 0 .
01 n .
This splitting is done after subsampling the 3 ◦ × 3 ◦ square.
M A LTERNATE AR P ROCEDURE WITH A UXILIARY D ATA
We propose an additional procedure which uses autoregressive sampling with auxiliary data to generate more expressive marginal predictives.
The input points of the auxiliary data are chosen randomly, and then sampled autoregressively before sampling the target points.
Finally, we discard the sampled values for the auxiliary data, but retain the samples for the target points.
Adding auxiliary points in this way allows the model to roll out autoregressively with more steps, even if the target set is small (or just a single point).
We describe the procedure below:
Procedure M.1 (Autoregressive application of neural process with auxiliary data) .
For a neural process π θ , context set D (c) =
( x (c) , y (c) ) , a target input x (t) , a distribution r over X , a number of auxiliary data points R ∈ N , and a number of trajectories M ∈ N , let AuxAR x (t) ( π θ , D (c) , r, R, M ) be the distribution defined as follows.
We first autoregressively sample the auxiliary data trajectories at random locations sampled from r :
[Equation]
Next, conditioned on the auxiliary data, we sample the target point of interest to make predictions.
We then marginalise out the auxiliary data by averaging over the M trajectories:
[Equation]
This procedure introduces three hyperparameters: the distribution from which to draw inputs r , the length of trajectories R , and the number of trajectories to sample M .
In the following experiments, we set the distribution r to be uniform over the training domain with no dependence on the context set or target point of interest: r = Uniform([ b, h ]) , where b and h are the lower and upper bounds of the training domain, respectively.
One could experiment with other choices for the distribution r .
The trajectory length R is chosen between 0 8 and 8 , and the number of trajectories M is chosen between 1 and 128 .
M. G ENERATED D ATA
We create three data generating processes for our experiments: a mixture of functions, random sawtooth functions, and random audio-like functions.
The first two experiments have multi-modal true marginals, whereas the last has heavy-tailed marginals.
Function mixture.
The function mixture data are generated by choosing one of the following three functions, the first with a probability of 14 , the second with a probability of 12 , and the third with a probability of 14 : [Equation]
[Equation]
Sawtooth.
The sawtooth data are generated from the following function:
[Equation]
We sample the frequency ω ∼ Unif([3 , 5]) , the direction d as either − 1 or 1 with equal probability, and the shift as ϕ ∼ Unif([ 1 ω , 1]) .
Synthetic Audio.
Synthetic audio data are generated by convolving a Dirac comb 9 with a truncated decaying sum of sinusoids:
[Equation] [Equation] [Equation]
The Dirac comb is defined as Comb T (
[Table] [Table]
where
[Equation]
We truncate the waves up to the period length, because otherwise the convolution with the Dirac comb would lead to increasing amplitude, resulting in a non-stationary process.
M. T RAINING
For the function mixture experiment, no training is required because we use the analytically derived ideal CNP π ∞ as our model.
See Section 2.
For the sawtooth and synthetic audio experiments, we train ConvCNP models.
We train each model for 100 epochs using 1024 batches per epoch with a batch size of 16.
We discretise the encoder by evaluating 64 points per unit.
We use a margin of 0.1, and a stride length of 2 for each of thelayers of the U-Net.
Each layer has 64 channels.
The receptive field size from this combination of parameters is 6.953.
During training, we sample a number of context points between uniformly at random from {0,...,75}, and we sample exactly 100 target points.
The context points and target points are sampled uniformly from [ − 2 , 2] .
We use the Adam optimizer with a learning rate of 3 × 10 − 4 . M.
[Table]
[Table]
In Table
[Table]
we see that using this procedure improves the held-out log-likelihoods for all of the experiments.
We can better understand the utility of this method by observing the performance for [Table] different context sizes in Tables 24 and 23.
For example, the sawtooth data results in Table 24 show that, for context set sizes 16, the AR method provides no benefits.
The increased flexibility of this
AR method is not needed in this case — a Gaussian predictive models the marginal sufficiently well.
For context set sizes of 4 and 8, on the other hand, we see significant improvements using the AR method.
Similarly, for the synthetic audio data results in Table 23, we see improvements using the
ConvCNP (AuxAR) for all context set sizes except 16.
The autoregressive sampling with auxiliary data method shown here shows promise for improving modeling of processes with multi-modal and heavy-tailed marginal distributions — all with no changes to the training procedure.
The scenarios where this method are most useful are highly contingent upon the context set size, because of its impact how well the marginals are modeled by
Gaussians.