B AYESIAN
C ONTEXT A GGREGATION FOR N EURAL
P ROCESSES
Anonymous authors
Paper under double-blind review
A BSTRACT
Formulating scalable probabilistic regression models with reliable uncertainty estimates has been a long-standing challenge in machine learning research.
Recently, casting probabilistic regression as a multi-task learning problem in terms of conditional latent variable (CLV) models such as the Neural Process (NP) has shown promising results.
In this paper, we focus on context aggregation, a central component of such architectures, which fuses information from multiple context data points.
So far, this aggregation operation has been treated separately from the inference of a latent representation of the target function in CLV models.
Our key contribution is to combine these steps into one holistic mechanism by phrasing context aggregation as a Bayesian inference problem.
The resulting Bayesian
Aggregation (BA) mechanism allows principled handling of task ambiguity, which is key for efﬁciently processing context information.
We demonstrate on a range of challenging experiments, including dynamics modeling and image completion, that
BA consistently improves upon the performance of traditional aggregation methods while remaining computationally efﬁcient and fully compatible with existing
NP-based models.
I NTRODUCTION
Estimating statistical relationships between physical quantities from measured data is of central importance in all branches of science and engineering and devising powerful regression models for this purpose forms a major ﬁeld of study in statistics and machine learning.
When judging representative power, neural networks (NNs) are arguably the most prominent member of the regression toolbox.
NNs cope well with large amounts of training data and are computationally efﬁcient at test time.
On the downside, standard NN variants do not provide uncertainty estimates over their predictions and tend to overﬁt on small datasets.
Gaussian processes (GPs) may be viewed as complementary to NNs as they provide reliable uncertainty estimates but their cubic (quadratic) scaling with the number of context data points at training (test) time in their basic formulation affects the application on tasks with large amounts of data or on high-dimensional problems.
Recently, a lot of interest in the scientiﬁc community is drawn to combinations of aspects of NNs and
GPs.
Indeed, a prominent formulation of probabilistic regression is as a multi-task learning problem formalized in terms of amortized inference in conditional latent variable (CLV) models, which results in NN-based architectures which learn a distribution over target functions.
Notable variants are given by the Neural Process (NP) (Garnelo et al., 2018b) and the work of Gordon et al.
(2019), which presents a unifying view on a range of related approaches in the language of CLV models.
Inspired by this research, we study context aggregation, a central component of such models, and propose a new, fully Bayesian, aggregation mechanism for CLV-based probabilistic regression models.
To transform the information contained in the context data into a latent representation of the target function, current approaches typically employ a mean aggregator and feed the output of this aggregator into a NN to predict a distribution of global latent parameters of the function.
Hence, aggregation and latent parameter inference have so far been treated as separate parts of the learning pipeline.
Moreover, when using a mean aggregator, every context sample is assumed to carry the same amount of information.
Yet, in practice, different input locations have different task ambiguity and, therefore, samples should be assigned different importance in the aggregation process.
In contrast, our Bayesian aggregation mechanism treats context aggregation and latent parameter inference as one holistic mechanism, i.e., the aggregation directly yields the distribution over the latent parameters of the target function.
Indeed, we formulate context aggregation as Bayesian inference of latent parameters using
Gaussian conditioning in the latent space.
Compared to existing methods, the resulting aggregator improves the handling of task ambiguity, as it can assign different variance levels to the context samples.
This mechanism improves predictive performance, while it remains conceptually simple and does not introduce noticeable computational overhead.
Furthermore, our Bayesian aggregator can also be applied to deterministic model variants like the Conditional NP (CNP) (Garnelo et al., 2018a).
In summary, our contributions are (i) a novel Bayesian Aggregation (BA) mechanism for context aggregation in NP-based models for probabilistic regression, (ii) its application to existing CLV architectures as well as to deterministic variants like the CNP, and (iii) an exhaustive experimental evaluation, demonstrating BA’s superiority to traditional mean aggregation.
R ELATED W ORK
Prominent approaches to probabilistic regression are Bayesian linear regression and its kernelized counterpart Gaussian processes (GPs) (Rasmussen and Williams, 2005).
The formal correspondence of GPs with inﬁnite-width Bayesian NNs (BNNs) has been established in Neal (1996) and Williams (1996).
A broad range of research aims to overcome the cubic scaling behaviour of GPs in the number of context points, e.g., through sparse GP approximations (Smola and Bartlett, 2001; Lawrence et al., 2002; Snelson and Ghahramani, 2005; Quiñonero-Candela and Rasmussen, 2005), by deep kernel learning (Wilson et al., 2016), by approximating the posterior distribution of BNNs (MacKay, 1992; Hinton and van Camp, 1993; Gal and Ghahramani, 2016; Louizos and Welling, 2017), or, by adaptive Bayesian linear regression, i.e., by performing inference over the last layer of a NN which introduces sparsity through linear combinations of ﬁnitely many learned basis functions (LazaroGredilla and Figueiras-Vidal, 2010; Hinton and Salakhutdinov, 2008; Snoek et al., 2012;
Calandra et al., 2016).
An in a sense complementary approach aims to increase the data-efﬁciency of deep probabilistic regression architectures by a fully Bayesian treatment of hierarchical latent variable models (“DeepGPs”)
(Damianou and Lawrence, 2013).
A parallel line of research studies probabilistic regression in the multi-task setting.
Here, the goal is to formulate models which are data-efﬁcient on an unseen target task by training them on data from a set of related source tasks.
Bardenet et al.
; Yogatama and Mann (2014); Golovin et al.
(2017) study multi-task formulations of GP-based models.
More general approaches of this kind employ the meta-learning framework (Schmidhuber, 1987; Thrun and Pratt, 1998; Vilalta and Drissi, 2005), where a model’s training procedure is formulated in a way which incentivizes it to learn how to solve unseen tasks rapidly with only a few context examples (“learning to learn”, “few-shot learning” (Fei-Fei et al., 2006; Lake et al., 2011)).
A range of such methods trains a meta-learner to learn how to adjust the parameters of the learner’s model (Bengio et al., 1991; Schmidhuber, 1992), an approach which has recently been applied to few-shot image classiﬁcation (Ravi and Larochelle, 2017) or to learn data-efﬁcient optimization algorithms (Hochreiter et al., 2001; Li and Malik, 2016;
Andrychowicz et al., 2016;
Chen et al., 2017; Perrone et al., 2018; Volpp et al., 2019).
Other branches of meta-learning research aim to learn similarity metrics to determine the relevance of context samples for the target task (Koch et al., 2015;
Vinyals et al., 2016;
Snell et al., 2017; Sung et al., 2017), or explore the application of memory-augmented neural networks for meta-learning (Santoro et al., 2016).
Finn et al.
(2017) propose model-agnostic meta-learning (MAML), a general framework for fast parameter adaptation for gradient-based learning methods.
A successful formulation of probabilistic regression as a few-shot learning problem in a multi-task setting is enabled by recent advances in the area of probabilistic meta-learning methods which allow a quantitative treatment of the uncertainty arising due to task ambiguity, a feature particularly relevant for few-shot learning problems.
One line of work speciﬁcally studies probabilistic extensions of MAML (Grant et al., 2018; Ravi and Larochelle, 2017; Rusu et al., 2018;
Finn et al., 2018;
Kim et al., 2018).
Further important approaches are based on amortized inference in multi-task
CLV models (Heskes, 2000; Bakker and Heskes, 2003; Kingma and Welling, 2013; Rezende et al., 2014; Sohn et al., 2015), which forms the basis of the Neural Statistician proposed by Edwards and
Storkey (2017) and of the NP model family (Garnelo et al., 2018b; Kim et al., 2019b; Louizos et al., 2019).
Gordon et al.
(2019) present a unifying view on many of the aforementioned probabilistic architectures.
Building on the conditional NPs (CNPs) proposed in Garnelo et al.
(2018a), a range of
NP-based architectures, such as Garnelo et al.
(2018b)
; Kim et al.
(2019b), consider combinations of deterministic and CLV model architectures.
Recently, Gordon et al.
(2020) extended CNPs to include translation equivariance in the input space, yielding state-of-the-art predictive performance.
In this paper, we also employ a formulation of probabilistic regression in terms of a multi-task CLV model.
However, while in previous work context aggregation mechanisms (Zaheer et al., 2017;
Wagstaff et al., 2019) were merely viewed as a necessity to consume context sets of variable size, we take inspiration from Becker et al.
and emphasize the fundamental connection of latent parameter inference with context aggregation and, hence, base our model on a novel Bayesian aggregation mechanism.
P RELIMINARIES
We present the standard multi-task CLV model which forms the basis for our discussion and present traditional mean context aggregation (MA) and the variational inference (VI)-inspired likelihood approximation as employed by the NP model family (Garnelo et al., 2018a;
Kim et al., 2019b), as well as an alternative Monte Carlo (MC)-based approximation.
Problem Statement.
We frame probabilistic regression as a multi-task learning problem.
Let
F denote a family of functions f (cid:96) :
R d x → R d y with some form of shared statistical structure.
z (cid:96) θ x c(cid:96),n y c(cid:96),n x t(cid:96),m y t(cid:96),m
N (cid:96)
M (cid:96)
L
We assume to have available data sets
[Equation]
,i , y (cid:96),i ) } i of evaluations y (cid:96),i ≡ f (cid:96) ( x (cid:96),i ) + ε from a subset of functions (“tasks”)
{ f (cid:96) } L(cid:96)
=1 ⊂ F with additive Gaussian noise ε ∼ N (cid:0) 0 , σ 2 n (cid:1) .
From this data, we aim to learn the posterior predictive distribution p ( y (cid:96) | x (cid:96) , D c(cid:96) ) over a (set of) y
(cid:96) , given the corresponding (set of) inputs x (cid:96) as well as a context set D c(cid:96) ⊂
D (cid:96) .
The Multi-Task CLV Model.
We formalize the multitask learning problem in terms of a CLV model (Heskes, 2000; Gordon et al., 2019) as shown in Fig.
The model employs task-speciﬁc global latent variables z (cid:96) ∈ R d z , as well as a task-independent latent variable θ , capturing the shared statistical structure between tasks.
To learn θ , we split the data into context sets D c(cid:96)
[Equation]
[Equation]
= and target sets D t(cid:96) ≡ { x t(cid:96),m , y t(cid:96),m } M
[Equation]
m =1 and maximize the posterior predictive likelihood function
L
[Equation]
In what follows, we omit task indices (cid:96) to avoid clutter.
Likelihood Approximation.
Marginalizing over the task-speciﬁc latent variables z is intractable for reasonably complex models, so one has to employ some form of approximation.
The NP-family of models (Garnelo et al., 2018b; Kim et al., 2019b) uses an approximation of the form
[Equation]
Being derived using a variational approach, this approximation utilizes an approximate posterior distribution q φ ( z | D c )
≈ p ( z | D c , θ ) .
Note, however, that it does not constitute a proper evidence lower bound for the posterior predictive likelihood since the intractable conditional prior p ( z | D c , θ ) has been replaced by q φ ( z | D c ) in the nominator of the rightmost term (Le et al., 2018).
A second possibility to approximate the integral, employed for instance in Gordon et al.
(2019), is to replace the intractable latent posterior distribution by an approximate distribution q φ ( z | D c )
≈ p ( z | D c , θ )
[Figure]
and to use a Monte-Carlo (MC) approximation of the resulting integral based on K latent samples, i.e.,
[Equation]
Note that both approaches employ approximations q φ ( z | D c ) of the latent posterior distribution p ( z | D c , θ ) and, as indicated by the notation, amortize inference in the sense that one single set of parameters φ is shared between all context data points.
This enables efﬁcient inference at test time, as no per-data-point optimization loops are required.
As is standard in the NP literature (Garnelo et al., 2018b; Kim et al., 2019b), we represent q φ ( z | D c ) and p (cid:0) y tm | z, x tm , θ (cid:1) by NNs and refer to them as the encoder (enc, parameters φ ) and decoder (dec, parameters θ ) networks, respectively.
These networks set the means and variances of factorized Gaussian distributions, i.e.,
[Equation]
Context Aggregation.
The latent variable z is global in the sense that it depends on the whole context set D c .
Therefore, some form of aggregation mechanism is required to enable the encoder to consume context sets D c of variable size.
To represent a meaningful operation on sets, such an aggregation mechanism has to be invariant to permutations of the context data points.
Zaheer et al.
(2017) characterize possible aggregation mechanisms w.r.t.
this permutation invariance condition, resulting in the structure of traditional aggregation mechanisms depicted in Fig. 2(a).
Each context data tuple ( x cn , y cn ) is ﬁrst mapped onto a latent observation r n = enc r,φ ( x cn ,
y cn ) ∈ R d r .
Then, a permutation-invariant operation is applied to the set { r n } Nn =1 to obtain an aggregated latent observation ¯ r .
One prominent choice, employed for instance in Garnelo et al.
(2018a), Kim et al.
(2019b), and Gordon et al.
(2019), is to take the mean, i.e.,
[Equation]
Subsequently, ¯ r is mapped onto the parameters µ z and σ 2 z of the approximate posterior distribution q φ ( z | D c ) using additional encoder networks, i.e., µ z = enc µ z ,φ (¯ r ) and σ 2 z = enc σ 2 z ,φ (¯ r ) .
Note that three encoder networks are employed here: (i) enc r,φ to map from the context pairs to r n , (ii) enc µ z ,φ to compute µ z from the aggregated mean ¯ r and (iii) enc σ 2 z ,φ to compute the variance σ 2 z from ¯ r .
In what follows, we refer to this aggregation mechanism as mean aggregation (MA) and to the networks enc µ z ,φ and enc σ 2 z ,φ collectively as “ ¯ r -to- z -networks”.
B AYESIAN C ONTEXT A GGREGATION
We propose Bayesian Aggregation (BA), a novel context data aggregation technique for CLV models which avoids the detour via an aggregated latent observation ¯ r and directly treats the object of interest, namely the latent variable z , as the aggregated quantity.
This reﬂects a central observation for CLV models with global latent variables: context data aggregation and hidden parameter inference are fundamentally the same mechanism .
Our key insight is to deﬁne a probabilistic observation model p ( r | z ) for r which depends on z .
Given a new latent observation r n = enc r,φ ( x cn , y cn ) , we can update p ( z ) by computing the posterior p ( z | r n ) = p ( r n | z ) p ( z ) /p ( r n ) .
Hence, by formulating context data aggregation as a Bayesian inference problem, we aggregate the information contained in D c directly into the statistical description of z based on ﬁrst principles.
4. B AYESIAN C ONTEXT A GGREGATION VIA G AUSSIAN C ONDITIONING
BA can easily be implemented using a factorized Gaussian observation model of the form
[Equation]
Note that, in contrast to standard variational auto-encoders (VAEs) (Kingma and Welling, 2013), we do not learn the mean and variance of a Gaussian distribution, but we learn the latent observation r n (which can be considered as a sample of p ( z ) ) together with the uncertainty σ 2 r n of this observation.
This architecture allows the application of Gaussian conditioning while this is difﬁcult for VAEs.
Indeed, we impose a factorized Gaussian prior p
( z ) ≡ N (cid:0) z | µ z, 0 , diag (cid:0) σ 2 z, 0 (cid:1)(cid:1) and arrive at a
Gaussian aggregation model which allows to derive the parameters of the posterior distribution q φ ( z | D c ) in closed form (cf.
App. 7.1):
[Equation]
Here (cid:9) , (cid:12) and (cid:11) denote element-wise inversion, product, and division, respectively.
These equations naturally lend themselves to efﬁcient incremental updates as new context data ( x cn , y cn ) arrives by using the current posterior parameters µ z, old and σ 2 z, old in place of the prior parameters, i.e.,
[Equation]
BA employs two encoder networks, enc r,φ and enc σ 2 r ,φ , mapping context tuples to latent observations and their uncertainties, respectively.
In contrast to MA, it does not require ¯ r -to- z -networks, because the set { r n } Nn =1 is aggregated directly into the statistical description of z by means of Eq.
(8), cf.
Fig.
2(b).
Note that our factorization assumptions avoid the expensive matrix inversions that typically occur in Gaussian conditioning and which are difﬁcult to backpropagate.
Using factorized distributions renders BA cheap to evaluate with only marginal computational overhead in comparison to MA.
Furthermore, we can easily backpropagate through BA to compute gradients to optimize the parameters of the encoder and decoder networks.
As the latent space z is shaped by the encoder network, the factorization assumptions are valid because the network will ﬁnd a space where these assumptions work well.
Note further that BA represents a permutation-invariant operation on D c .
Discussion.
BA includes MA as a special case.
Indeed, Eq.
(8) reduces to the mean-aggregated latent observation Eq.
(6) if we impose a non-informative prior and uniform observation variances σ 2 r n ≡ 1 .
This observation sheds light on the beneﬁts of a Bayesian treatment of aggregation.
MA assigns the same weight 1 /N to each latent observation r n , independent of the amount of information contained in the corresponding context data tuple ( x cn , y cn ) , as well as independent of the uncertainty about the current estimation of z .
Bayesian aggregation remedies both of these limitations: the inﬂuence of r n on the parameters µ z, old and σ 2 z, old describing the current aggregated state is determined by the relative magnitude of the observation variances σ 2 r n and the latent variances σ 2 z, old , cf.
Eq.
This emphasizes the central role of the learned observation variances σ 2 r n : they allow to quantify the amount of information contained in each latent observation r n .
BA can therefore handle task ambiguity more efﬁciently than MA, as the architecture can learn to assign little weight (by predicting high observation variances σ 2 r n ) to context points ( x cn , y cn ) located in areas with high task ambiguity, i.e., to points which could have been generated by many of the functions in F .
Conversely, in areas with little task ambiguity, i.e., if ( x cn , y cn ) contains a lot of information about the underlying function, BA can induce a strong inﬂuence on the posterior latent distribution.
In contrast,
MA has to ﬁnd ways to propagate such information through the aggregation mechanism by encoding it in the mean-aggregated latent observation ¯ r .
4. L IKELIHOOD A PPROXIMATION WITH B AYESIAN C ONTEXT A GGREGATION
We show that BA is versatile in the sense that it can replace traditional MA in CLV-based NP architectures as proposed, e.g., in Garnelo et al.
(2018b) and Gordon et al.
(2019), which employ sampling-based likelihood approximations as discussed in Sec. 3, as well as in deterministic variants like the Conditional NP (CNP) (Garnelo et al., 2018a).
Sampling-Based Likelihood Approximations.
BA is naturally compatible with both the VIinspired and MC likelihood approximations for CLV models.
Indeed, BA deﬁnes a Gaussian latent distribution from which we can easily obtain samples z in order to evaluate Eq.
(2) or Eq.
(3) using the decoder parametrization Eq.
Bayesian Context Aggregation for Conditional Neural Processes.
BA motivates a novel, alternative, method to approximate the posterior predictive likelihood Eq.
(1), resulting in a deterministic loss function which can be efﬁciently optimized for θ and φ in an end-to-end fashion.
To this end, we employ a Gaussian approximation of the posterior predictive likelihood of the form
[Equation]
This is inspired by GPs which also deﬁne a Gaussian likelihood.
Maximizing this expression yields the optimal solution µ y = ˜
µ y , Σ y =
˜Σ y , with ˜ µ y and ˜Σ y being the ﬁrst and second moments of the true posterior predictive distribution.
This is a well-known result known as moment matching , a popular variant of deterministic approximate inference used, e.g., in Deisenroth and Rasmussen (2011) and Becker et al.
˜ µ y and ˜Σ y are functions of the moments µ z and σ 2 z of the latent posterior p ( z | D c , θ ) which motivates the following decoder parametrization:
[Equation]
Here, µ z and σ 2 z are given by the BA Eqs.
Note that we deﬁne the Gaussian approximation to be factorized w.r.t.
individual y tm , an assumption which simpliﬁes the architecture but could be dropped if a more expressive model was required.
This decoder can be interpreted as a “moment matching network”, computing the moments of y given the moments of z .
Indeed, in contrast to decoder networks of CLV-based NP architectures as deﬁned in Eq.
(5), it operates on the moments µ z and σ 2 z of the latent distribution instead of on samples z which allows to evaluate this approximation in a deterministic manner.
In this sense, the resulting model is akin to the CNP which deﬁnes a deterministic, conditional model with a decoder operating on the mean-aggregated latent observation ¯ r .
However, BA-based models trained in this deterministic manner still beneﬁt from BA’s ability to accurately quantify global parameter uncertainty which yields signiﬁcantly improved predictive likelihoods.
In what follows, we refer to this approximation scheme as direct parameter-based (PB) likelihood optimization.
Discussion.
The concrete choice of likelihood approximation or, equivalently, model architecture depends mainly on the intended use-case.
Sampling-based models are generally more expressive as they can represent complex, i.e., structured, non-Gaussian, posterior predictive distributions.
Moreover, they yield true function samples while deterministic models only allow approximate function samples through auto-regressive (AR) sampling schemes.
Nevertheless, deterministic models exhibit several computational advantages.
They yield direct probabilistic predictions in a single forward pass, while the predictions of sampling-based methods are only deﬁned through averages over multiple function samples and hence require multiple forward passes.
Likewise, evaluating the MC-based likelihood approximation Eq.
(3) during training requires to draw multiple ( K ) latent samples z .
While the VI-inspired likelihood approximation Eq.
(2) can be optimized on a single function sample per training step through stochastic gradient descent (Bishop, 2006), it is ﬂawed as it is not a proper evidence lower bound and furthermore requires to feed target sets through the encoder which can impede the training for small context sets D c as discussed in detail in App. 7.2.
E XPERIMENTS
We present experiments to compare the performances of BA and of MA in NP-based models.
To provide a complete picture, we evaluate all combinations of likelihood approximations (PB/deterministic
[Table]
Eq.
(10), VI-inspired Eq.
(2), MC Eq.
and aggregation methods (BA Eq.
(8), MA Eq.
(6)), resulting in six different model architectures, cf.
Fig.
4 in App.
Two of these architectures correspond to existing members of the NP family: MA + deterministic is equivalent to the CNP (Garnelo et al., 2018a), and MA + VI-inspired corresponds to the Latent-Path NP (LP-NP) (Garnelo et al., 2018b), i.e., the NP without a deterministic path.
We also provide results for the Attentive
Neural Process (ANP) (Kim et al., 2019b), which employs a hybrid approach, combining LP-NP with an attention mechanism in a parallel deterministic path.
2 We emphasize, however, that the main purpose of our experiments is to provide a clear-cut comparison of BA and MA and note that
BA could also be used in hybrid models like ANP, but do not pursue this idea further and leave it for future research.
We further point the reader to App. 7.3, where we present evaluations of an
NP-architecture using MA together with a self-attentive encoder network as well as to App. 7.4, where we discuss NP-based regression on the context of general probabilistic regression methods and compare our results with standard (multi-task) GP approaches.
The performance of NP-based models depends heavily on the encoder and decoder network architectures as well as on the latent space dimensionality d z .
To assess the inﬂuence of the aggregation mechanism independently from all other confounding factors, we consistently optimize the encoder and decoder network architectures, the latent-space dimensionality d z , as well as the learning rate of the Adam optimizer (Kingma and Ba, 2015), independently for all model architectures and for all experiments using the Optuna (Akiba et al., 2019) framework, cf.
App.
If not stated differently, we report performance in terms of the mean posterior predictive log-likelihood over 256 test tasks with 256 data points each, conditioned on context sets containing N ∈ { 0 , 1 , . .
. , N max } data points.
For sampling-based methods (VI-inspired, MC, ANP), we report the joint log-likelihood over the test sets using a Monte-Carlo approximation with 25 latent samples, cf.
App.
We average the resulting log-likelihood values over 10 training runs with different random seeds and report 95% conﬁdence intervals.
We publish source code to reproduce the experimental results online.
GP Samples.
We evaluate the architectures on synthetic functions drawn from GP priors with different kernels (RBF, weakly periodic, Matern-5/2), as proposed by Gordon et al.
(2020), cf.
App.
We generate a new batch of functions for each training epoch.
The results (Tab. 1) show that our BA consistently outperforms MA, independent of the model architecture.
Interestingly, despite employing a factorized Gaussian approximation, our deterministic PB approximation performs at least on-par with the traditional VI-inspired approximation which tends to perform [Table]
[Table]
particularly poorly for small context sets, reﬂecting the intricacies discussed in Sec.
As expected, the MC approximation yields the best results in terms of predictive performance, as it is more expressive than the deterministic approaches and does not share the problems of the VI-inspired approach.
As shown in Tab.and Tab. 10, App.
7.6, our proposed PB likelihood ap-
[Table]
[Figure]
MC
[Figure]
proximation is much cheaper to evaluate compared to both sampling-based approaches which require multiple forward passes per prediction.
We further observe that BA tends to require smaller encoder and decoder networks as it is more efﬁcient at propagating context information to the latent state as discussed in Sec. 4.1.
The hybrid ANP approach is competitive only on the Matern-5/2 function class.
Yet, for a fair assessment, we refer the reader to Tab. 11, App. 7.6, demonstrating that the attention mechanism greatly improves performance in terms of MSE but does not perform well in terms of the likelihood.
Further, it leads to unnaturally “shaky” predictions, cf.
Figs. 8 to 13, App. 7.6.
Quadratic Functions.
We further seek to study the performance of BA with very limited amounts of training data.
To this end, we consider two quadratic function classes, each parametrized by three real parameters from which we generate limited numbers L of training tasks.
The ﬁrst function class is deﬁned on a one-dimensional domain, i.e., x ∈ R , and we choose L = 64 , while the second function class, as proposed by Perrone et al.
(2017), is deﬁned on x ∈ R 3 with L = 128 , cf.
App.
As shown in Tab. 3, BA again consistently outperforms MA, often by considerably large margins, underlining the efﬁciency of our Bayesian approach to aggregation in the regime of little training data.
On the 1 D task, all likelihood approximations perform approximately on-par in combination with
BA, while MC outperforms both on the more complex 3 D task.
Fig.
3 compares prediction qualities.
Dynamics of a Furuta Pendulum.
We study BA on a realistic dataset given by the simulated dynamics of a rotary inverted pendulum, better known as the Furuta pendulum (Furuta et al., 1992), which is a highly non-linear dynamical system, consisting of an actuated arm rotating in the horizontal plane with an attached pendulum rotating freely in the vertical plane, parametrized by two masses, three lengths, and two damping constants.
The regression task is deﬁned as the one-step-ahead prediction of the four-dimensional system state with a step-size of ∆ t = 0 .
1 s , as detailed in App.
The results (Tab. 4) show that
BA improves predictive performance also on complex, non-synthetic regression tasks with higher-dimensional input- and output spaces.
Further, they are consistent with our previous ﬁndings regarding the likelihood approximations, with MC being strongest in terms of predictive likelihood, followed by our efﬁcient deterministic alternative PB.
[Figure]
[Table]
2D Image Completion.
Our ﬁnal experiment is a 2 D image completion task on the MNIST dataset (LeCun and Cortes, 2010), as proposed by Garnelo et al.
(2018a;b) and Kim et al. (2019a).
Here, the inputs x are pixel locations in images showing handwritten digits, and we regress onto the corresponding pixel intensities y , cf.
App.
Interestingly, we found that architectures without deterministic paths were not able to solve this task reliably which is why we only report results for deterministic models.
As shown in Tab. 5, BA improves performance in comparison to by a large margin.
This reﬂects the
[Table]
fact that BA’s ability to quantify the information content of a context tuple is particularly beneﬁcial on this task, as pixels in certain areas in the middle of the images typically convey much more information about the identity of the digit than, e.g., pixels located near the borders.
C ONCLUSION AND O UTLOOK
We proposed a novel Bayesian Aggregation (BA) method for NP-based models, combining context aggregation and hidden parameter inference in one holistic mechanism which enables efﬁcient handling of task ambiguity.
BA is conceptually simple, compatible with existing NP-based model architectures, and consistently improves performance compared to traditional mean aggregation.
It introduces only marginal computational overhead, simpliﬁes the architectures in comparison to existing CLV models (no ¯ r -to- z encoder), and tends to require less complex encoder and decoder network architectures.
Our experiments further demonstrate that the VI-inspired likelihood approximation traditionally used to train NP-based models should be abandoned in favor of a MC-based approach, and that our proposed PB likelihood approximation represents an efﬁcient deterministic alternative with strong predictive performance.
We believe that a range of existing models, e.g., the ANP, can beneﬁt from
BA, especially when reliable quantiﬁcation of uncertainty is crucial.
Also, more complex Bayesian aggregation models are conceivable, opening interesting avenues for future research.
A PPENDIX
We present the derivation of the Bayesian aggregation update equations Eqs.
(8), (9) in more detail.
To foster reproducibility, we describe all experimental settings as well as the hyperparameter optimization procedure used to obtain the results reported in Sec. 5, and publish the source code online.
4 We further provide additional experimental results and visualizations of the predictions of the compared architectures.
7. D ERIVATION OF THE B AYESIAN A GGREGATION U PDATE E QUATIONS
We derive the full Bayesian aggregation update equations without making any factorization assumptions.
We start from a Gaussian observation model of the form
[Equation]
where r n and Σ r n are learned by the encoder network.
If we impose a Gaussian prior in the latent space, i.e.,
[Equation]
we arrive at a Gaussian aggregation model which allows to derive the parameters of the posterior distribution, i.e., of
[Equation] [Equation]
in closed form using standard Gaussian conditioning (Bishop, 2006):
[Equation]
As the latent space z is shaped by the encoder network, it will ﬁnd a space where the following factorization assumptions work well (given d z is large enough):
[Equation]
[Equation]
This yields a factorized posterior, i.e.,
[Equation]
with
[Equation]
Here (cid:9) , (cid:12) and (cid:11) denote element-wise inversion, product, and division, respectively.
This is the result
Eq.
(8) from the main part of this paper.
7. D ISCUSSION OF VI- INSPIRED L IKELIHOOD A PPROXIMATION
To highlight the limitations of the VI approximation, we note that decoder networks of models employing the PB or the MC likelihood approximation are provided with the same context information at training and test time: the latent variable (which is passed on to the decoder in the form of latent samples z (for MC) or of parameters µ z , σ 2 z describing the latent distribution (for PB)) is in both cases conditioned only on the context set D c .
In contrast, in the variational approximation Eq.
(2), the expectation is w.r.t.
q φ conditioned on the union of the context set D c and the target set D t .
As
D t is not available at test time, this introduces a mismatch between how the model is trained and
[Table]
[Table]
how it is used at test time as the decoder is trained on samples from q φ ( z | D c ∪ D t ) but evaluated on samples from q φ ( z | D c ) .
This is not a serious problem when the model is evaluated on context sets with sizes large enough to allow accurate approximations of the true latent posterior distribution.
Small context sets, however, usually contain too little information to infer z reliably.
Consequently, the distributions q φ ( z | D c ) and q φ ( z | D c ∪ D t ) typically differ signiﬁcantly in this regime.
Hence, incentivizing the decoder to yield meaningful predictions on small context sets requires intricate and potentially expensive additional sampling procedures to choose suitable target sets D t during training.
As a corner case, we point out that it is not possible to train the decoder on samples from the latent prior, because the right hand side of Eq.
(2) vanishes for D c
= D t = ∅ .
7. S ELF -A TTENTIVE E NCODER A RCHITECTURES
Kim et al.
(2019b) propose to use attention-mechanisms to improve the quality of NP-based regression.
In general, given a set of key-value pairs { ( x n , y n ) } Nn =1 , x n ∈ R d x , y n ∈ R d y , and a query x ∗ ∈ R d x , an attention mechanism A produces a weighted sum of the values, with the weights computed from the keys and the query:
[Equation]
There are several types of attention mechanisms proposed in the literature (Vaswani et al., 2017), each deﬁning a speciﬁc form of the weights w n .
Laplace attention adjusts the weights according to their spatial distance
[Equation]
Similarly, dot-product attention computes
[Equation]
A more sophisticated mechanism is multihead attention , which employs a set of 3 H learned linear mappings (cid:8) L Kh (cid:9)
H h
=1 , (cid:8) L Vh (cid:9)
H h
=1 , (cid:110)
L Qh (cid:111)
H h
=1 .
For each h , these mappings are applied to keys, values, and queries, respectively.
Subsequently, dot-product attention is applied to the set of transformed key-value pairs and the transformed query.
The resulting H values are then again combined by a further learned linear mapping L O to obtain the ﬁnal result.
Self-attention (SA) is obtained by setting the set of queries equal to the set of keys.
Therefore, SA produces again a set of N weighted values.
Combining SA with an NP-encoder, i.e., applying SA to the set { f x ( x n ) , r n } Nn =1 of inputs x n and corresponding latent observations r n (where we also consider a possible nonlinear transformation of the inputs f x ) and subsequently applying MA yields an interesting baseline for our proposed BA.
Indeed, similar to BA, SA computes a weighted sum of the latent observations r n .
Note, however, that SA weights each latent observation according to some form of spatial relationship of the corresponding input with all other latent observations in the context set.
In contrast, BA’s weight for a given latent observation is based only on features computed from the context tuple corresponding to this very latent observation and allows to incorporate an estimation of the amount of information contained in the context tuple into the aggregation (cf.
Sec. 4.1).
This leads to several computational advantages of BA over SA: (i)
In general, SA scales quadratically in the number N of context tuples, as it has to be evaluated on all N 2 pairs of context tuples.
In contrast,
BA scales linearly with N .
(ii) BA allows for efﬁcient incremental updates when context data arrives sequentially (cf.
Eq.
(9)), while using SA does not provide this possibility: it requires to store and encode the whole context set D c at once and to subsequently aggregate the whole set of resulting (SA-weighted) latent observations.
Tabs.
6, 7 show a comparison of BA with SA in the encoder combined with MA for various different attention mechansims.
The results show that Laplace SA and dot-product SA cannot improve the performance compared to MA without an SA mechanism in the encoder.
In contrast, multihead SA leads to signiﬁcant improvements in predictive performance and can often even compete with BA.
We emphasize, however, that BA does not use SA in the encoder in this comparison and therefore still beneﬁts from the aforementioned computational advantages.
In situations where computational disadvantages can be accepted in favour of increased predictive performance, a combination of BA with SA in the encoder can be interesting.
Note that BA relies on a second encoder output σ 2 r n (in addition to the latent observation r n ) which assesses the information content in each context tuple ( x n , y n ) .
As each SA-weighted r n is informed by the other latent observations in the context set, obviously, one would have to also process the set of σ 2 r n with some form of SA mechanism.
We are conﬁdent that this is indeed possible in a theoretically well-founded manner, but note that it is not immediately obvious how to do this properly.
Therefore, we leave a combination of SA and BA for future research.
7. N EURAL P ROCESS - BASED M ODELS IN THE C ONTEXT OF S CALABLE P ROBABILISTIC R EGRESSION
We discuss in more detail how NP-based models relate to other existing methods for (scalable) probabilistic regression such as (multi-task) GPs (Rasmussen and Williams, 2005; Bardenet et al., 2013; Yogatama and Mann, 2014; Golovin et al., 2017), Bayesian neural networks (BNNs) (MacKay, 1992; Gal and Ghahramani, 2016), and DeepGPs (Damianou and Lawrence, 2013).
NPs are motivated in the seminal papers (Garnelo et al., 2018a;b; Kim et al., 2019b) as well as in our
Sec. 1 as models which combine the computational efﬁciency of neural networks with well-calibrated uncertainty estimates (like those of GPs).
Indeed, NPs scale linearly in the number N of context and M of target data points, i.e., like O ( N + M ) while GPs scale like O ( N 3 + M 2 ) .
Furthermore,
NPs are shown to exhibit well-calibrated uncertainty estimates.
In this sense, NPs can be counted as members of the family of scalable probabilistic regression methods.
A central aspect of NP training which distinguishes NPs from a range of standard methods is that they are trained in a multi-task fashion.
This means that NPs rely on data from a set of related source tasks from which they automatically learn powerful priors and the ability to adapt quickly to unseen target tasks.
This multi-task training procedure of NPs scales linearly in the number L of source tasks, which makes it possible to train these architectures on large amounts of source data.
Applying
GPs in such a multi-task setting can be challenging, especially for large numbers of source tasks
[Table]
(Bardenet et al., 2013; Yogatama and Mann, 2014; Golovin et al., 2017).
Similarly, Bayesian Neural
Networks (BNNs) (MacKay, 1992; Gal and Ghahramani, 2016) as well as DeepGPs (Damianou and
Lawrence, 2013) are in their vanilla forms speciﬁcally designed for the single-task setting.
Therefore,
GPs, BNNs, and DeepGPs are not directly applicable in the NP multi-task setting, which is why they are typically not considered as baselines for NP-based models as discussed in (Kim et al., 2019b).
Indeed, the experiments presented in (Garnelo et al., 2018a;b; Kim et al., 2019b) focus mainly on evaluating NPs in the context of few-shot probabilistic regression, i.e., on demonstrating the data-efﬁciency of NPs on the target task.
In contrast, the application of NPs in situations with large ( > 1000 ) numbers of context/target points per task has to the best of our knowledge not yet been investigated in detail in the literature.
Furthermore, it has not been studied how to apply NPs in situations where only a single or very few source tasks are available.
The focus of our paper is a clear-cut comparison of the performance of our BA with traditional MA in the context of NP-based models.
Therefore, we also consider experiments similar to those presented in (Garnelo et al., 2018a;b; Kim et al., 2019b) and leave further comparisons with existing methods for (multi-task) probabilistic regressions for future work.
Nevertheless, to illustrate this discussion, we provide two simple GP-based baseline methods: (i) a vanilla GP, i.e., the hyperparameters are optimized on each target task individually and the source data is not used, and (ii) a naive but easily interpretable example of a multi-task GP which optimizes one set of hyperparameters on all source tasks and uses it for predictions on the target tasks without further adaptation.
The results in Tab. 8 show that those GP-based models can only compete with
NPs on function classes where either the inductive bias as given by the kernel functions ﬁts the data well, or on function classes which exhibit a relatively low degree of variablity.
On more complex function classes NPs produce predictions of much better quality, as they incorporate the source data more efﬁciently.
7. E XPERIMENTAL D ETAILS
We provide details about the data sets as well as about the experimental setup used in our experiments in Sec. 5.
7.5. D ATA G ENERATION
In our experiments, we use several classes of functions to evaluate the architectures under consideration.
To generate training data from these function classes, we sample L random tasks (as described in Sec. 5), and N tot random input locations x for each task.
For each minibatch of training tasks, we uniformly sample a context set size N ∈ { n min , . . . , n max } and use a random subset of N data points from each task as context sets D c .
The remaining M = N tot − N data points are used as the target sets D t (compare App.
7.5.3 for the special case of the VI-inspired likelihood approximation.).
Tab. 9 provides details about the data generation process.
GP Samples.
We sample one-dimensional functions f : R → R from GP priors with three different stationary kernel functions as proposed by Gordon et al.
A radial basis functions (RBF) kernel with lenghtscale l = 1 . 0 :
[Equation]
[Equation]
A weakly periodic kernel:
[Equation]
A Matern-5/2 kernel with lengthscale l = 0 .
[Equation]
Quadratic Functions.
We consider two classes of quadratic functions.
The ﬁrst class f Q, 1 D : R → R is deﬁned on a one-dimensional domain and parametrized by three parameters a, b, c ∈ R :
[Equation]
The second class f Q, 3 D : R 3 → R is deﬁned on a three-dimensional domain and also parametrized by three parameters a, b, c ∈ R : [Equation]
[Equation]
This function class was proposed in Perrone et al.
For both function classes we add Gaussian noise with standard deviation σ n to the evaluations, cf.
Tab. 9.
Furuta Pendulum Dynamics.
We consider a function class obtained by integrating the non-linear equations of motion governing the dynamics of a Furuta pendulum (Furuta et al., 1992; Cazzolato and Prime, 2011) for a time span of ∆ t = 0 .
1 s. More concretely, we consider the mapping
[Equation]
where Θ = (cid:104) θ arm ( t ) , θ pend ( t ) , ˙ θ arm ( t ) , ˙ θ pend ( t ) (cid:105) T denotes the four-dimensional vector describing the dynamical state of the Furuta pendulum.
The Furuta pendulum is parametrized by seven parameters (two masses, three lengths, two damping constants) as detailed in Tab. 9.
During training, we provide L = 64 tasks, corresponding to 64 different parameter conﬁgurations.
We consider the free system and generate noise by applying random torques at each integration time step ( ∆ t Euler = 0 .
001 s) to the joints of the arm and pendulum drawn from Gaussian distributions with standard deviations σ τ, pend , σ τ, arm , respectively.
2D Image Completion.
For this task, we use the MNIST database of 28 · 28 images of handwritten digits (LeCun and Cortes, 2010), and deﬁne 2 D functions mapping pixel locations x 1 , x 2 ∈ { 0 , . . .
27 } (scaled to the unit square) to the corresponding pixel intensities y ∈ { 0 , . . .
, 255 } (scaled to the unit interval), cf.
Tab. 9.
One training task corresponds to one image drawn randomly from the training set (consisting of 60000 images) and for evaluation we use a subset of the test set (consisting of 10000 images).
7.5. M ODEL A RCHITECTURES
We provide the detailed architectures used for the experiments in Sec. 5 in Fig.
For ANP we use multihead attention and refer the reader to Kim et al.
(2019b) for details about the architecture.
[Table]
( x cn ,y cn )
MLP
MLP r n σ 2 r n
Bayesianaggregation µ z σ 2 z
MLP
MLP x t x t µ y σ 2 y (a) BA + PB.
( x cn ,y cn )
MLP r n
Mean aggregation ¯ r
MLP
MLP x t x t µ y σ 2 y (b) MA + det.
(CNP) ( x cn ,y cn )
MLP
MLP r n σ 2 r n
Bayesianaggregation µ z σ 2 z z
MLP
MLP x t x t µ y σ 2 y (c) BA + VI-inspired, BA + MC.
( x cn ,y cn )
MLP r n
Mean aggregation ¯ r
MLP
MLP µ z σ 2 z z
MLP
MLP x t x t µ y σ 2 y (d) MA + VI-inspired (LP-NP), MA + MC.
7.5. H YPERPARAMETERS AND H YPERPARAMETER O PTIMIZATION
To arrive at a fair comparison of our BA with MA, it is imperative to use optimal model architectures for each aggregation method and likelihood approximation under consideration.
Therefore, we optimize the number of hidden layers and the number of hidden units per layer of each encoder and decoder MLP (as shown in Fig. 4), individually for each model architecture and each experiment.
For the ANP, we also optimize the multihead attention MLPs.
We further optimize the latent space dimensionality d z and the learning rate of the Adam optimizer.
For this hyperparameter optimization, we use the Optuna framework (Akiba et al., 2019) with TPE Sampler and Hyperband pruner (Li et al., 2017).
We consistently use a minibatch size of 16 tasks and S
= 10 latent samples to evaluate the
MC likelihood approximation during training.
To evaluate the VI-inspired likelihood approximation, we sample target set sizes between N tot and N in each training epoch, cf.
Tab. 9.
7.5. E VALUATION P ROCEDURE
To evaluate the performance of the various model architectures we generate L = 256 unseen test tasks with target sets D t(cid:96) consisting of M = 256 data points each and compute the average posterior predictive log-likelihood
[Equation]
1 M (cid:80)
[Equation]
D c(cid:96) , θ (cid:17) , given context sets D c(cid:96) of size N .
Depending on the architecture, we approximate the posterior predictive log-likelihood according to:
• For BA + PB likelihood approximation:
[Equation]
• For MA + deterministic loss (= CNP):
[Equation]
• For architectures employing sampling-based likelihood approximations (VI-inspired, MCLL)
we report the joint log-likelihood over all data points in a test set, i
[Equation]
e.
[Equation]
where z (cid:96),s ∼ q φ ( z | D (cid:96) ) .
We employ S = 25 latent samples.
To compute the log-likelihood values given in tables, we additionally average over various context set sizes N as detailed in the main part of this paper.
We report the mean posterior predictive log-likelihood computed in this way w.r.t.
10 training runs with different random seeds together with 95% conﬁdence intervals
[Table]
We average over the same context set sizes as used to compute the posterior predictive log-likelihood, cf.
Sec. 5, and again use S = 25 latent samples to compute the mean prediction of sampling-based methods.
Our BA consistently improves predictive performance compared to MA not only in terms of the likelihood (as shown in Sec. 5), but also in terms of MSE.
Furthermore, while ANP tends to perform poorly in terms of likelihood (cf.
Sec. 5), it’s MSE is improved greatly by the attention mechanism
[Table]
7. A DDITIONAL E XPERIMENTAL R ESULTS
We provide additional experimental results accompanying the experiments presented in Sec. 5:
• Results for relative evaluation runtimes and numbers of parameters of the optimized network architectures on the full GP suite of experiments, cf.
Tab.
The posterior predictive mean squared error on all experiments, cf.
Tab. 11.
• The context-size dependent results for the predictive posterior log-likelihood for the 1 D and 3 D Quadratic experiments, the Furuta dynamics experiment, as well as the 2 D image completion experiment, cf.
Fig.
• More detailed plots of the predictions on one-dimensional experiments ( 1 D Quadratics (Figs. 6, 7), RBF-GP, (Figs. 8, 9), Weakly Periodic GP (Figs. 10, 11), and Matern-5/2 GP (Figs.
[Figure]
[Figure]
l og p ( y t 1 : M | x t 1 : M , D c )
Quadratic1D  20 N −−−−−
Quadratic3D  20 N
FurutaDynamics 200 N   2DImageCompletion
Aggregator BA
MA Loss
PB/det.
VI-inspired
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]