M ARTINGALE P OSTERIOR N EURAL P ROCESSES
A BSTRACT
A Neural Process ( NP ) estimates a stochastic process implicitly defined with neural networks given a stream of data, rather than pre-specifying priors already known, such as Gaussian processes.
An ideal NP would learn everything from data without any inductive biases, but in practice, we often restrict the class of stochastic processes for the ease of estimation.
One such restriction is the use of a finite-dimensional latent variable accounting for the uncertainty in the functions drawn from NP s. Some recent works show that this can be improved with more “data-driven” source of uncertainty such as bootstrapping.
In this work, we take a different approach based on the martingale posterior, a recently developed alternative to Bayesian inference.
For the martingale posterior, instead of specifying prior-likelihood pairs, a predictive distribution for future data is specified.
Under specific conditions on the predictive distribution, it can be shown that the uncertainty in the generated future data actually corresponds to the uncertainty of the implicitly defined Bayesian posteriors.
Based on this result, instead of assuming any form of the latent variables, we equip a NP with a predictive distribution implicitly defined with neural networks and use the corresponding martingale posteriors as the source of uncertainty.
The resulting model, which we name as Martingale Posterior Neural Process ( MPNP ), is demonstrated to outperform baselines on various tasks.
I NTRODUCTION
A Neural Process ( NP ) (Garnelo et al., 2018a;b) meta-learns a stochastic process describing the relationship between inputs and outputs in a given data stream, where each task in the data stream consists of a meta-training set of input-output pairs and also a meta-validation set.
The NP then defines an implicit stochastic process whose functional form is determined by a neural network taking the meta-training set as an input, and the parameters of the neural network are optimized to maximize the predictive likelihood for the meta-validation set.
This approach is philosophically different from the traditional learning pipeline where one would first elicit a stochastic process from the known class of models (e.g., Gaussian Processes ( GP s)) and hope that it describes the data well.
An ideal NP would assume minimal inductive biases and learn as much as possible from the data.
In this regard, NP s can be framed as a “data-driven” way of choosing proper stochastic processes.
An important design choice for a NP model is how to capture the uncertainty in the random functions drawn from stochastic processes.
When mapping the meta-training set into a function, one might employ a deterministic mapping as in Garnelo et al.
(2018a).
However, it is more natural to assume that there may be multiple plausible functions that might have generated the given data, and thus encode the functional (epistemic) uncertainty as a part of the NP model.
Garnelo et al.
(2018b) later proposed to map the meta-training set into a fixed dimensional global latent variable with a Gaussian posterior approximation.
While this improves upon the vanilla model without such a latent variable (Le et al., 2018), expressing the functional uncertainty only through the Gaussian approximated latent variable has been reported to be a bottleneck (Louizos et al., 2019).
To this end,
Lee et al.
(2020) and Lee et al.
(2022) propose to apply bootstrap to the meta-training set to use the uncertainty arising from the population distribution as a source for the functional uncertainty.
In this paper, we take a rather different approach to define the functional uncertainty for NP s. Specifically, we utilize the martingale posterior distribution (Fong et al., 2021), a recently developed alternative to conventional Bayesian inference.
In the martingale posterior, instead of eliciting a likelihood-prior pair and inferring the Bayesian posterior, we elicit a joint predictive distribution on future data given observed data.
Under suitable conditions on such a predictive distribution, it can be shown that the uncertainty due to the generated future data indeed corresponds to the uncertainty of the Bayesian posterior.
Following this, we endow a NP with a joint predictive distribution defined through neural networks and derive the functional uncertainty as the uncertainty arising when mapping the randomly generated future data to the functions.
Compared to the previous approaches of either explicitly positing a finite-dimensional variable encoding the functional uncertainty or deriving it from a population distribution, our method makes minimal assumptions about the predictive distribution and gives more freedom to the model to choose the proper form of uncertainty solely from the data.
Due to the theory of martingale posteriors, our model guarantees the existence of the martingale posterior corresponding to the valid Bayesian posterior of an implicitly defined parameter.
Furthermore, working in the space of future observations allows us to incorporate the latent functional uncertainty path with deterministic path in a more natural manner.
We name our extension of NP s with the joint predictive generative models as the Martingale Posterior
Neural Process ( MPNP ).
Throughout the paper, we propose an efficient neural network architecture for the generative model that is easy to implement, flexible, and yet guarantees the existence of the martingale posterior.
We also propose a training scheme to stably learn the parameters of MPNP s.
Using various synthetic and real-world regression tasks, we demonstrate that MPNP significantly outperforms the previous NP variants in terms of predictive performance.
B ACKGROUND 2. S ETTINGS AND NOTATIONS
Let X =
R d in be an input space and Y =
R d out be an output space.
We are given a set of tasks drawn from an (unknown) task distribution, τ 1 , τ 2 , . . .
i .
i .
d . ∼ p task ( τ ) .
A task τ consists of a dataset Z and an index set c , where Z = { z i } ni =1 with each z i = ( x i , y i ) ∈
X × Y is a pair of an input and an output.
We assume Z are i.i.d.
conditioned on some function f .
The index set c ⊊ [ n ] where [ n ] := { 1 , . . . , n } defines the context set Z c = { z i } i ∈ c .
The target set Z t is defined similarly with the index t :=
[ n ] \ c .
2. N EURAL PROCESS FAMILIES
Our goal is to train a class of random functions f : X → Y that can effectively describe the relationship between inputs and outputs included in a set of tasks.
Viewing this as a meta-learning problem, for each task τ , we can treat the context Z c as a meta-train set and target Z t as a metavalidation set.
We wish to meta-learn a mapping from the context Z c to a random function f that recovers the given context Z c (minimizing meta-training error) and predicts Z t well (minimizing meta-validation error).
Instead of directly estimating the infinite-dimensional f , we learn a mapping from Z c to a predictive distribution for finite-dimensional observations,
[Equation]
where we are assuming the outputs Y are independent given f and X .
We further restrict ourselves to simple heteroscedastic Gaussian measurement noises,
[Equation]
where µ θ : X → Y and σ 2 θ : X → R + map an input to a mean function value and corresponding variance, respectively.
θ ∈ R h is a parameter indexing the function f , and thus the above predictive distribution can be written as
[Equation]
A NP is a parametric model which constructs a mapping from Z c to θ as a neural network.
The simplest version, Conditional Neural Process ( CNP ) (Garnelo et al., 2018a), assumes a deterministic mapping from Z c to θ as
[Equation]
where δ a ( x ) is the Dirac delta function (which gives zero if x ̸ = a and (cid:82) δ a ( x )d x = 1 ) and f enc is a permutation-invariant neural network taking sets as inputs (Zaheer et al., 2017), parameterized by ϕ enc .
Given a summary θ = r c of a context Z c , the CNP models the mean and variance functions ( µ, σ 2 ) as
[Equation]
where f dec is a feed-forward neural network parameterized by ϕ dec .
Here the parameters ( ϕ enc , ϕ dec ) are optimized to maximize the expected predictive likelihood over tasks, E τ [log p
( Y | X, Z c )] .
Note that in the CNP , the mapping from Z c to θ is deterministic, so it does not consider functional uncertainty or epistemic (model) uncertainty.
To resolve this, Garnelo et al.
(2018b) proposed NP which learns a mapping from an arbitrary subset Z ′ ⊆
Z to a variational posterior q ( θ | Z ′ )
approximating p ( θ | Z ′ ) under an implicitly defined prior p ( θ ) :
[Equation]
With f enc , the Evidence Lower BOund ( ELBO ) for the predictive likelihood is written as
[Equation]
An apparent limitation of the NP is that it assumes a uni-modal Gaussian distribution as an approximate posterior for q ( θ | Z c ) .
Aside from the limited flexibility, it does not fit the motivation of NP s trying to learn as much as possible in a data-driven manner, as pre-specified parametric families are used.
There have been several improvements over the vanilla CNP s and NP s, either by introducing attention mechanism (Vaswani et al., 2017) for f enc and f dec (Kim et al., 2018), or using advanced functional uncertainty modeling (Lee et al., 2020; Lee et al., 2022).
We provide a detailed review of the architectures for such variants in Appendix A. Throughout the paper, we will refer to this class of models as Neural Process Family ( NPF ).
2. M ARTINGALE P OSTERIOR D ISTRIBUTIONS
The martingale posterior distribution (Fong et al., 2021) is a recent generalization of Bayesian inference which reframes posterior uncertainty on parameters as predictive uncertainty on the unseen population conditional on the observed data.
Given observed samples Z = { z i } ni =1 i.i.d.
from the sampling density p 0 , one can define the parameter of interest as a functional of p 0 , that is
[Equation]
where ℓ is a loss function.
For example, ℓ ( z, θ ) =
( z − θ ) 2 would return θ 0 as the mean, and ℓ ( z, θ ) =
− log p ( z | θ ) would return the KL minimizing parameter between p ( · | θ ) and p 0 .
The next step of the martingale posterior is to construct a joint predictive density on Z ′ =
{ z i } Ni = n + for some large N , which we write as p ( Z ′ | Z ) .
In a similar fashion to a bootstrap, one can imagine drawing Z ′ ∼ p
( Z ′ | Z ) , then computing θ ( g N ) where g N ( z )
= 1 N (cid:80) Ni =
δ z i
( z ) .
The predictive uncertainty in Z ′ induces uncertainty in θ ( g N ) conditional on Z .
The key connection is that if p ( Z ′ | Z ) is the Bayesian joint posterior predictive density, and ℓ = − log p ( z | θ ) , then θ ( g N ) is distributed according to the Bayesian posterior π ( θ | Z ) as N → ∞ , under weak conditions.
In other words, posterior uncertainty in θ is equivalent to predictive uncertainty in { z i } ∞ i = n +1 .
Fong et al.
(2021) specify more general p ( Z ′ | Z ) directly beyond the Bayesian posterior predictive, and define the (finite) martingale posterior as π N ( θ ∈
A | Z ) = (cid:82) 1 ( θ ( g N )
∈ A ) p ( dZ ′ | Z ) .
In particular, the joint predictive density can be factorized into a sequence of 1-step-ahead predictives, p ( Z ′ | Z ) =
[Equation]
Ni = n +1 p
( z i | z 1: i − 1 ) , and the sequence { p ( z i | z 1: i − 1 ) } Nn +1 is elicited directly, removing the need for the likelihood and prior.
Hyperparameters for the sequence of predictive distributions can be fitted in a data-driven way by maximizing
[Equation]
which is analogous to the log marginal likelihood.
Fong et al.
(2021) requires the sequence of predictives to be conditionally identically distributed (c.i.d.), which is a martingale condition on the sequence of predictives that ensures g N exists almost surely.
The Bayesian posterior predictive density is a special case, as exchangeability of p ( Z ′ | Z ) implies the sequence of predictives is c.i.d.
In fact, De Finetti’s theorem (De Finetti, 1937) guarantees that any exchangeable joint density implies an underlying likelihood-prior form, but specifying the predictive density directly can be advantageous.
It allows for easier computation, as we no longer require posterior approximations, and it also widens the class of available nonparametric predictives which we will see shortly.
2. E XCHANGEABLE G ENERATIVE M ODELS
To construct a martingale posterior, we can either specify a sequence of one-step predictive distributions or the joint predictive density distribution directly, as long as the c.i.d.
condition is satisfied.
Here, we opt to specify an exchangeable p ( Z ′ | Z ) directly, which then implies the required c.i.d.
predictives.
We now briefly review exchangeable generative models which can be used to specify the exchangeable joint predictive.
For a set of random variables Z = { z i } ni =1 with each z i ∈
Z = R d , we say the joint distribution p ( Z ) is exchangeable if it is invariant to the arbitrary permutation of the indices, that is, p ( Z )
= p ( π · Z )
for any permutation π of [ n ] .
A simple way to construct such exchangeable random variables is to use a permutation-equivariant mapping .
A mapping f : Z n → Z n is permutation equivariant if f ( π · Z ) =
π · f ( Z ) for any π .
Given f , we can first generate i.i.d.
random variables and apply f to construct a potentially correlated but exchangeable set of random variables
Z as follows:
[Equation]
i .
d
[Equation]
Z
[Equation]
For f , we employ the modules introduced in Lee et al.
Specifically, we use a permutation equivariant module called Induced Self-Attention Block ( ISAB ).
An ISAB mixes input sets through a learnable set of parameters called inducing points via Multihead Attention Blocks ( MAB s) (Vaswani et al., 2017; Lee et al., 2019).
ISAB ( E ) =
MAB ( E , H ) ∈ R n × d where H = MAB ( I, E ) ∈
R m × d .
Here, I ∈ R m × d is a set of m inducing points and MAB ( · , · ) computes attention between two sets.
The time-complexity of an ISAB is O ( nm ) , scales linear with input set sizes.
M ETHODS
In this section, we present a novel extension of NPF called MPNP s. The main idea is to elicit joint predictive distributions that are constructed with equivariant neural networks instead of assuming priors for θ , and let the corresponding martingale posterior describe the functional uncertainty in the
NP s. We describe how we construct a MPNP in Section 3.1 and train it in Section 3.2.
3. M ARTINGALE P OSTERIOR N EURAL P ROCESSES
Recall that the functional uncertainty in a NP is encoded in a parameter θ .
Rather than learning an approximate posterior q ( θ | Z c ) , we introduce a joint predictive p ( Z ′
|
Z c ϕ pred ) generating a pseudo context set Z ′ =
{ z ′ i } N −| c |
i = of size ( N − | c | ) ≥ 1 .
Having generated a pseudo context, we combine with the existing context Z c , and construct the empirical density as
[Equation]
Given g N , the estimate of the function parameter θ is then recovered as
[Equation] [Equation]
where in our case we simply choose ℓ ( z, θ ) :
= − log N ( y | µ θ ( x ) , σ 2 θ ( x ) I d out ) .
The uncertainty in θ ( g N ) is thus induced by the uncertainty in the generated pseudo context Z ′ .
[Figure]
Amortization The procedure of recovering θ via Eq.
11 would originally require iterative optimization process except for simple cases.
Fortunately, in our case, we can amortize this procedure, thanks to the mechanism of CNP s amortizing the inference procedure of estimating θ from the context.
Given Z c , a CNP learns an encoder producing θ that is trained to maximize the expected likelihood.
That is,
[Equation]
where g c is the empirical density of Z c .
Hence, given Z ′ and Z c , we can just input Z ′ ∪ Z c into f enc and use the output ˜ θ
( Z ′ ∪ Z c ) as a proxy for θ ( g N ) .
Compared to exactly computing θ ( g N ) , obtaining ˜ θ ( Z ′ ∪ Z c ) requires a single forward pass through f enc , which scales much better with N .
Moreover, computation for multiple Z ′ required for bagging can easily be parallelized.
Specifying the joint predictives We construct the joint predictives p ( Z ′ | Z c ; ϕ pred ) with a neural network.
Other than the requirement of Z ′ being exchangeable (and thus c.i.d.), we give no inductive bias to p ( Z ′ | Z c ; ϕ pred ) and let the model learn ϕ pred from the data.
We thus use the exchangeable generative model described in Section 2.4.
Specifically, to generate Z ′ , we first generate E = { ε i } ni =1 from some distribution (usually chosen to be a unit Gaussian N (0 , I d ) ), and pass them through an equivariant ISAB block to form Z ′ .
To model the conditioning on Z c , we set the inducing point in the ISAB as a transform of Z c .
That is, with an arbitrary feed-forward neural network h ,
ISAB (
[Equation]
(13) where h ( Z c ) = { h ( z i ) } i ∈ c .
The resulting model is an implicit generative model (Mohamed and
Lakshminarayanan, 2016) in a sense that we can draw samples from it but cannot evaluate likelihoods.
Generating Representations When z is low-dimensional, it would be moderately easy to learn the joint predictives, but in practice, we often encounter problems with high-dimensional z , for instance when the input x is a high-resolution image.
For such cases, directly generating z may be harder than the original problem, severely deteriorating the overall learning procedure of MPNP .
Instead, we propose to generate the encoded representations of z .
The encoders of the most of the
NPF s first encode an input z
i into a representation r i .
For the remaining of the forward pass, we only need r i s instead of the original input z .
Hence we can build a joint predictives p ( R ′ | R c ; ϕ pred ) generating R ′ =
{ r ′ i } N −| c | i = conditioned on R c
= { r i } i ∈ c as for generating Z ′ from Z c .
In the experiments, we compare these two versions of MPNP s (generating Z ′ and generating R ′ ), and found that the one generating R ′ works much better both in terms of data efficiency in training and predictive performances, even when the dimension of z is not particularly large.
See Fig.
1 for our method applying to CNP model (Garnelo et al., 2018a).
3. T RAINING
With the generator p ( Z ′ | Z c ; ϕ pred )
[Equation]
the marginal likelihood for a task τ
= ( Z, c ) is computed as
[Equation]
Note that p ( Z ′ | Z c ; ϕ pred ) is c.i
[Equation]
d., so there exists a corresponding martingale posterior π N such that
[Equation]
We approximate the marginal likelihood via a consistent estimator,
[Equation]
where Z ′ (1) , . . .
, Z ′ ( K ) i .
i .
d . ∼ p ( Z ′ | Z c ; ϕ pred ) .
This objective would be suffice if we are given sufficiently good ˜
θ ( Z c ∪ Z ′ ( k ) )
[Equation]
but we have to also train the encoder to properly amortize the parameter construction process Eq.
For this, we use only the given context data to optimize
[Equation]
that is, we train the parameters ( ϕ enc , ϕ dec ) using CNP objective.
Furthermore, we found that if we just maximize Eq.
16 and Eq.
17, the model can cheat by ignoring the generated pseudo contexts and use only the original context to build function estimates.
To prevent this, we further maximize the similar CNP objectives for each generated pseudo context to encourage the model to actually make use of the generated contexts.
[Equation]
Combining these, the loss function for the MPNP is then
[Equation]
R ELATED W ORKS
CNP (Garnelo et al., 2018a) is the first NPF model which consists of simple MLP layers as its encoder and decoder.
NP (Garnelo et al., 2018b) also uses MLP layers as its encoder and decoder but introduces a global latent variable to model a functional uncertainty.
Conditional Attentive Neural
Process ( CANP )
(Kim et al., 2018) and Attentive Neural Process ( ANP ) (Kim et al., 2018) are the models which apply attention modules as their encoder block in order to well summarize context information relevant to target points.
Louizos et al.
(2019) proposed NP s model which employs local latent variables instead of a global latent variable by applying a graph neural network.
By applying convolution layers as their encoder, Gordon et al.
(2020) and Foong et al.
(2020) introduced a translation equivariant CNP s and NP s model, respectively.
In addition to these works, Bootstrapping Neural
Process ( BNP ) (Lee et al., 2020) suggests modeling functional uncertainty with the bootstrap (Efron, 1992) method instead of using a single global latent variable.
E XPERIMENTS
We provide extensive experimental results to show how MPNP and Martingale Posterior Attentive
Neural Process ( MPANP ) effectively increase performance upon the following baselines: CNP , NP ,
BNP , CANP , ANP , and Bootstrapping Attentive Neural Process ( BANP ).
All models except deterministic models (i.e., CNP and CANP ) use the same number of samples; K = 5 for the image completion task and K = 10 for the others.
Refer to Appendices A and C for more detailed experimental setup including model architectures, dataset and evaluation metrics.
[Figure]
− −    −
[Table]
5. 1D R EGRESSION
In this section
[Table]
we conducted 1D regression experiments following Kim et al.
(2018) and Lee et al.
In this experiments
[Table]
the dataset curves are generated from GP with 4 different settings: i)
RBF kernels, ii) Mat ´ ern 5/2 kernels, iii) Periodic kernels
[Table]
and iv) RBF kernels with Student’s t noise.
Infinite Training Dataset Previous works (Garnelo et al., 2018b; Kim et al., 2018; Le et al., 2018) assumed that there exists a GP curve generator that can provide virtually infinite amount of tasks for training.
We first follow this setup, training all models for 100,000 steps where a new task is generated from each training step.
We compare the models by picking checkpoints achieving the lowest validation loss.
Table 1 clearly shows that our model outperforms the other models in most cases.
This results show that our model well captures the functional uncertainty compared to the other methods.
In Appendix B, we also report the comparison with the baselines with increased number of parameters to match the additional number of parameters introduced for the generator in our model, where ours still significantly outperforms the baselines.
Finite Training Dataset
We also compare the models on more realistic setting assuming a finite amount of training tasks.
Specifically, we first configured the finite training dataset consisting of { 51200 , 102400 , 256000 } examples at the start of the training, instead of generating new tasks for each training step.
We then trained all models with the same 100,000 training iterations in order to train the models with the same training budget as in the infinite training dataset situation.
Fig.
3 clearly shows that our model consistently outperforms other models in terms of the target log-likelihood even when the training dataset is finite.
This indicates that MPNP s effectively learn a predictive distribution of unseen dataset from a given dataset with small number of tasks.
Refer to Appendix B for more detailed results.
[Figure]
[Table]
5. I MAGE C OMPLETION
Next we conducted 2D image completion tasks for three different datasets, i.e., MNIST, SVHN, and CelebA. For training, we uniformly sample the number of context pixels | c | ∈ { 3 , ..., 197 } and the number of target pixels | t | ∈ { 3 , ..., 200 − | c |} from an image.
For evaluation, we uniformly sample the number of context pixels | c | ∈ { 3 , ..., 197 } and set all the remaining pixels as the targets.
Table 2 clearly demonstrates that our model outperforms the baselines over all three datasets, demonstrating the effectiveness of our method for high-dimensional image data.
See Appendix B for the visualizations of completed images along with the uncertainties in terms of predictive variances, and Appendix C for the detailed training setup.
5. B AYESIAN O PTIMIZATION
Using pre-trained models with RBF kernels in Section 5.1 Infinite Training Dataset experiments, we conducted Bayesian optimization (Brochu et al., 2010) for two benchmark functions (Gramacy and
Lee, 2012; Forrester et al., 2008).
As a performance measurement, we use best simple regret, which measures the difference between the current best value and the global optimum value.
Fig.
4 depicts the normalized regret and the cumulative normalized regret averaged over 100 trials of the Gramacy and Lee (2012) function.
Here, we also consider a GP variant with RBF kernel, tuned by pretraining (Wang et al., 2022).
It clearly demonstrates that our model shows the best performance among NP s for both the normalized regret and the cumulative normalized regret.
Appendix B. provides the results for the Forrester et al.
(2008) function and Appendix C.4 provides detailed experimental setups.
[Figure]
[Table]
[Figure]
[Table]
5. P REDATOR -P REY M ODEL
Following Lee et al.
[Table]
we conducted the predator-prey population regression experiments.
We first trained the models using the simulation datasets which are generated from a Lotka-Volterra model (Wilkinson, 2018) with the simulation settings followed by Lee et al.
Then tested on the generated simulation test dataset and real-world dataset which is called Hudson’s Bay harelynx data.
As mentioned in Lee et al.
(2020), the real-world dataset shows different tendency from generated simulation datasets, so we can treat this experiment as model-data mismatch experiments.
In Table 3, we can see the MPNP s outperform the other baselines for the test simulation datasets but underperforms in the real-world dataset compare to other baselines.
This shows that model-data mismatch is an open problem for the MPNP s.
C ONCLUSION
In this paper, we proposed a novel extension of NP s by taking a new approach to model the functional uncertainty for NP s. The proposed model MPNP utilizes the martingale posterior distribution (Fong et al., 2021), where the functional uncertainty is driven from the uncertainty of future data generated from the joint predictive.
We present a simple architecture satisfying the theoretical requirements of the martingale posterior, and propose a training scheme to properly train it.
We empirically validate
MPNP s on various tasks, where our method consistently outperforms the baselines.
Limitation As we presented in the Predator-Prey Model experiments in Section 5.4, our method did not significantly outperform baselines under model-data mismatch.
This was also higlighted in
Fong et al.
(2021): model-data mismatch under the martingale posterior framework remains an open problem.
Our method with direct input generation also performed poorly, as we found it difficult to prevent models from generating meaningless inputs that are ignored by the decoders.
We present more details on unsuccessful attempts for direct input generation in Appendix D.
Societal Impacts Our work is unlikely to bring any negative societal impacts.
Modeling functional uncertainty may be related to the discussion of safe AI within the community.
Reproducibility Statement We argued our experimental details in Appendix C which contains used libraries and hardwares.
We presented all the dataset description in Appendix C. We describes the model architecture details in Appendix A. A CKNOWLEDGEMENTS
This work was partly supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government(MSIT) (No.2019-0-00075, Artificial Intelligence Graduate School Program(KAIST)), Institute of Information & communications
Technology Planning & Evaluation (IITP) grant funded by the Korea government(MSIT) (No.20220-00713), and Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government(MSIT) (No.2021-0-02068, Artificial Intelligence Innovation
Hub).
A M ODEL A RCHITECTURES
In this section, we summarize the model architectures which we used in experiments.
Here, we only present simplified structures for each model.
To see exact computation procedures for BNP s, please refer to Lee et al.
Fig.
5 shows our method applying to CANP model (Kim et al., 2018).
A. M ODULES
Linear Layer Lin ( d in , d out ) denotes the linear transformation of the input with dimension d in into the output with dimension d out .
Multi-Layer Perceptron MLP ( n l , d in , d hid , d out ) denotes a multi-layer perceptron with the structure:
[Equation]
where ReLU denotes the element-wise Rectified Linear Unit (ReLU) activation function.
Multi-Head Attention MHA ( n head , d out )( Q, K, V ) denotes a multi-head attention (Vaswani et al., 2017) with n head heads which takes input as ( Q, K, V ) and outputs the feature with dimension d out .
The actual computation of MHA ( n head , d out )( Q, K, V ) can be written as follows:
[Equation]
Lin ( d k , d out )( K ) , n head ) ( V ′ i ) n head i =1 = split ( Lin ( d v , d out )( V ) , n head )
H = concat ([ softmax ( Q ′ i
K ′⊤ i / (cid:112)
d out ) V ′
i ] n head i =1 )
O = LN ( Q ′ + H )
MHA ( n head , d out )( Q, K, V ) =
LN ( O + ReLU ( Lin ( d out , d out )( O )))
where ( d q , d k , d v ) denotes the dimension of Q, K, V respectively, split and concat are the splitting and concatenating A in the feature dimension respectively, and LN denotes the layer normalization (Ba et al., 2016).
Self-Attention SA ( n head , d out ) denotes a self-attention module which is simply computed as
SA ( n head , d out )( X ) =
MHA ( n head , d out )( X, X, X ) .
Multi-head Attention Block MAB ( n head , d out ) denotes a multi-head attention block module (Lee et al., 2019) which is simply computed as MAB ( n head , d out )
( X, Y ) =
MHA ( n head , d out )( X, Y, Y ) .
Induced Set Attention Block ISAB ( n head , d out ) denotes a induced set attention block (Lee et al., 2019) which constructed with two stacked MAB layers.
The actual computation of
ISAB ( n head , d out )( X, Y ) can be written as follows:
[Equation]
A. CNP , NP , BNP , N EURAL B OOTSTRAPPING N EURAL P ROCESS ( NEUBNP ) AND MPNP
Encoder
The models only with a deterministic encoder ( CNP , BNP , MPNP ) use the following structure:
[Equation]
For the MPNP , f enc ( Z c ) changes into concat ([ r c , r ′ c ]) where r ′ c is the feature of the pseudo context data generated from generator in paragraph Generator .
The model also with a latent encoder ( NP ) uses:
[Equation]
where d z = d x + d y denotes the data dimension.
Data dimensions vary through tasks, d x = 1 , d y = 1 for 1D regression tasks, d x = 2 , d y = 1 for MNIST image completion task, d x = 2 , d y = 3 for
SVHN and CelebA image completion tasks, and d x = 1 , d y = 2 for Lotka Volterra task.
Adaptation Layer
BNP uses additional adaptation layer to combine bootstrapped representation and the base representation.
This can be done with a simple linear layer
[Equation]
Decoder All models use a single MLP as a decoder.
The models except NP uses the following structure:
[Equation]
and NP uses:
[Equation]
Generator
MPNP use a single ISAB module as a generator.
The ISAB uses the following structure:
[Equation]
where ϵ i s are i.i.d.
sampled from Gaussian distribution with dimension 128 and n gen denotes a number of pseudo context data.
A. CANP , ANP , BANP AND MPANP
Encoder
The models only with a deterministic encoder ( CANP , BANP , Neural Bootstrapping Attentive Neural Process ( N eu BANP ) and MPANP ) use the following structure:
[Equation]
For the MPANP , f enc ( Z c ) changes into r c = MHA ( n
[Equation]
, d [Equation] , concat [Equation]
[ r k , r ′ k ]) , concat ([ r
[Equation]
, r [Equation] )) , f enc = r c , where r ′ k and r ′ v are the key and value features of the pseudo context data generated from generator in paragraph Generator .
ANP constructed as:
[Equation]
Note that r q and r k are from the same MLP.
Adaptation Layer Like BNP , BANP also uses adaptation layer with same structure to combine bootstrapped representations.
Decoder All models use the same decoder structure as their non-attentive counterparts.
Generator
MPANP use a single ISAB module as a generator.
The ISAB uses the following structure:
[Equation]
where ϵ i s are i.i.d.
sampled from Gaussian distribution with dimension 256 and n gen denotes a number of pseudo context data.
B A DDITIONAL E XPERIMENTS B. 1D R EGRESSION
Full results for Table
We provide the full test results for 1D regression tasks including context, target, and task log-likelihood values in Table 4.
[Table]
[Figure]
‘Task’ denotes the task loglikelihood.
All values are averaged over four seeds
[Table]
Increasing the encoder size of baselines Since the generator increases the size of the encoder in
MPNP s, one can claim that the performance gain of MPNP s may come from the increased model size.
To verify this, we increased the hidden dimensions of the encoder of baselines and compared them with ours.
The results displayed in Table 5 further clarify that ours still outperforms the baselines even when the number of parameters gets in line.
B. H IGH -D R EGRESSION
We conducted additional experiments on the synthetic high-dimensional regression data (i.e., generating one-dimensional y from four-dimensional x with RBF kernel).
Here we used the same model structures with the 1D regression task except for the input layer, and the same settings for the RBF kernel with 1D regression except for l ∼ Unif (0 . 5 , 3 . 0) .
We fixed the base learning rate to 0 .for
all models throughout the high-dimensional regression experiments.
Table 6 clearly shows our MPNP s still outperform baselines for log-likelihood values we measured.
B. I MAGE C OMPLETION
MNIST We provide some completed MNIST images in Fig.
It shows that both MPNP and
MPANP successfully fill up the remaining parts of the image for a given context and capture the uncertainties as predictive variances.
CelebA We also present five examples from the CelebA dataset in Fig.
It shows that MPANP provides perceptually reasonable predictions even for complex three-channel images.
B. B AYESIAN O PTIMIZATION
We provide the results for Bayesian optimization on the Forrester et al.
(2008) function in Fig.
Our MPNP s consistently outperform baselines as discussed in Section 5.3.
We also present the visual results for Bayesian optimization in Figs.
9 and 10.
[Table]
C E XPERIMENTAL D ETAILS
We attached our code in supplementary material.
Our codes used python libraries JAX (Bradbury et al., 2018),
Flax (Heek et al., 2020) and Optax (Hessel et al., 2020).
These python libraries are available under the Apache-2.0 license 1 .
We conducted all experiments on a single NVIDIA GeForce RTX 3090 GPU, except for the image completion tasks presented in Section 5.2; we used 8 TPUv3 cores supported by TPU Research
Cloud 2 for the 2D image completion task.
For optimization, we used Adam (Kingma and Ba, 2015) optimizer with a cosine learning rate schedule.
Unless specified, we selected the base learning rate from a grid of { 5 × 10 − 4 .
50 , 5 × 10 − 4 .
25 , 5 × 10 − 4 . 00 , 5 × 10 − 3 .
75 , 5 × 10 − 3 .
50 } based on validation task log-likelihood.
C. E VALUATION METRIC
Following Le et al.
(2018), for CNP and CANP , which are deterministic models, we used the normalized predictive log-likelihood 1 n (cid:80) ni
=1 log p ( y i | x i , Z c ) .
For other models, we used a approximation of the normalized predictive log-likelihood as:
[Equation]
where θ k s are independent samples for k ∈
[ K ] .
C. 1D R EGRESSION
To generate tasks ( Z, c ) , we first sample x i .
i .
d .
∼
Unif ( − 2 , 2) and generate Y using each kernel.
We use RBF
[Equation]
(cid:17) , Matern 5 / 2 kernel k
[Equation]
x
[Equation]
) = s [Equation] (cid:16) 1 + [Equation]
+ [Equation] [Equation] 2 (cid:17) , and periodic kernel k ( x, x ′ ) =
s 2 · exp (cid:16) − 2sin 2 ( π || x − x
[Equation]
| [Equation]
(cid:17) where all kernels use s ∼ Unif (0 . 1 . 1 . 0) , ℓ ∼ Unif (0 . 1 . 0 . 6) , and p ∼ Unif (0 . 1 . 0 . 5) .
To generate t-noise dataset, we use Student- t with degree of freedom 2 .
1 to sample noise ϵ ∼ γ · T (2 . 1) where γ ∼ Unif (0 , 0 . 15) .
Then we add the noise to the curves generated from RBF kernel.
We draw index set | c | ∼ Unif (3 , 50 − 3) and n − | c | ∼ Unif (3 , 50 − | c | ) to maintain max | Z | ≤ 50 .
We use a batch size of 256 for training.
C. I MAGE C OMPLETION
We use the following datasets for image completion experiments.
[Figure]
MNIST We split MNIST (LeCun et al., 1998) train dataset into train set with 50,000 samples and validation set with 10,000 samples.
We use whole 10,000 samples in test dataset as test set.
We
[Figure] [Figure]
[Figure]
h i t
[Figure]
t
i
[Figure]
i
[Figure]
r a t i
[Figure]
i
[Figure]
t h
i t e r a t i o n
[Figure]
t h
i t e r a t i o n
[Figure]
t h
i t e r a t i o n
[Figure] [Figure]
i t e r a t i o n
[Figure]
t h
i t e r a t i o n
[Figure]
t h i t e r
[Figure]
t i o n
[Figure]
make 28 × 28 grids which both axes starting from − 0 .
5 to 0 .
5 to indicate the coordinate of pixels, and normalize pixel values into [ − 0 .
We use a batch size of 128 for training.
SVHN We split SVHN (Netzer et al., 2011) train dataset into train set with 58,600 samples and validation set with 14,657 samples.
We use whole 26,032 samples in test dataset as test set.
We make 32 × 32 grids which both axes starting from − 0 .
5 to 0 .
5 to indicate the coordinate of pixels, and normalize pixel values into [ − 0 .
We use a batch size of 128 for training.
CelebA We use splits of CelebA (Liu et al., 2015) dataset as provided (162,770 train samples, 19,867 validation samples, 19,962 test samples).
We crop 32 × 32 pixels of center of images.
We make 32 × 32 grids which both axes starting from − 0 .
5 to 0 .
5 to indicate the coordinate of pixels, and normalize pixel values into [ − 0 .
We use a batch size of 128 for training.
C. B AYESIAN O PTIMIZATION
We use the following benchmark functions for Bayesian optimization experiments.
Throughout the experiments, we adjust the function to have the domain of [ − 2 . 0 , 2 . 0] .
Gramacy and Lee (2012) function
[Equation]
(22) where x ∈ [0 . 5 , 2 . 5] and a global optimum is at x ∗ ≈ 0 . 5486 .
Forrester et al. (2008) function
[Equation]
(23) where x ∈ [0 , 1] and a global optimum is at x ∗ ≈ 0 . 7572 .
D D IRECTLY G ENERATING I NPUT M ODEL
In this section, we present our model generating pseudo contexts directly in the input space.
We will present two kinds of model structure, i) directly generating pseudo context pair ( x, y ) simultaneously by ISAB, ii) generating pseudo context data x and y , sequentially.
D. C ONSTRUCTION
Generating pseudo context pair simultaneously.
The generator of our first model which simultaneously generating pseudo context pair ( x ′ , y ′ ) , takes real context dataset Z c as input and outputs pseudo context dataset Z ′ .
Here the generator is the one layer ISAB module.
Then we concatenate
Z c and Z ′ in order to treat this concatenated set as context dataset.
Then the encoder takes this concatenated context set as input.
And the others are the same with CNP or CANP .
Sequentially generating pseudo context data x and y
In this model, the generator takes real context dataset Z c as input and outputs only x ′ s of Z ′ .
Here the generator is the one layer ISAB module with additional one linear layer.
Then we consider these x ′ s as our target dataset and find the mean and variance of y ′ for each x ′ by forwarding the model with context dataset Z c and target x ′ .
We sample y ′
from the Gaussian distribution with mean and variance from the prior step.
We again concatenate Z c with Z ′ and use them as context dataset.
Training Having directly generated a pseudo context set, we construct our empirical density as
[Equation] [Equation]
[Table]
Given g N
[Table]
we find the function parameter θ as
[Equation]
where we simply choose l ( z, θ ) :
= − log N ( y | µ θ ( x ) , σ 2 θ ( x ) I d out ) .
In order to train the directly generating input model, which well approximate θ ( g N ) , we should construct different objective function from Eq. 19 because we can compute the exact (cid:82) ℓ ( z, θ ) g N ( dz ) , unlike the feature generating model.
First, we approximate the marginal likelihood which is,
[Equation]
where Z ′ (1) , . . .
, Z ′ ( K ) i .
i .
d . ∼ p ( Z ′ | Z c ; ϕ pred ) .
Eq.
26 is the same training object with Eq. 16.
As we mentioned in Section 3.2, if we are given sufficiently well approximated ˜ θ
( Z c ∪ Z ′
( K ) )
then this objective would be suffice.
However only with Eq. 26, we cannot train the encoder to properly amortize the parameter construction process Eq.
To overcome this issue, we use (cid:82) ℓ ( z, θ ) g N ( dz ) as our second training objective which is,
[Equation]
Combining these two functions, our loss function for the direct MPNP is then
[Equation]
D. S AMPLE
In this section, we presents how the directly generating input model actually samples the pseudo context datasets.
In Fig.
11, we report generated pseudo context datasets and posterior samples from two different cases of directly generating input models for 1D regression task with RBF kernel.
Here we can see that the generator samples pseudo context datasets far from the real context dataset.
This phenomenon occurs because the generator learns to generate meaningless inputs ignored by the decoder.
In Fig.
12, we report how two different directly generating MPANP s predict posterior samples for 1D regression task with RBF kernel.
Although directly generated pseudo context dataset are a bit far from context dataset, our model still well capture the functional uncertainty in this case.
We report
[Figure]
[Figure] [Figure]
[Figure]
[Figure]
−   the test results for 1D regression tasks on RBF for two directly generating models in Table 7.
DSI and DSE indicate simultaneously generating models and sequentially generating models, respectively.
Table 7 shows that our directly generating models still outperform CNP and CANP in the perspective of log-likelihood.