# Bayesian Neural Network Inference via Implicit Models and the Posterior
Predictive Distribution

Joel Janek Dabrowski
Data61, CSIRO,
Australia
<EMAIL_ADDRESS>&Daniel Edward Pagendam
Data61, CSIRO,
Australia
<EMAIL_ADDRESS>
###### Abstract

We propose a novel approach to perform approximate Bayesian inference in
complex models such as Bayesian neural networks. The approach is more scalable
to large data than Markov Chain Monte Carlo, it embraces more expressive
models than Variational Inference, and it does not rely on adversarial
training (or density ratio estimation). We adopt the recent approach of
constructing two models: (1) a primary model, tasked with performing
regression or classification; and (2) a secondary, expressive (e.g. implicit)
model that defines an approximate posterior distribution over the parameters
of the primary model. However, we optimise the parameters of the posterior
model via gradient descent according to a Monte Carlo estimate of the
posterior predictive distribution – which is our only approximation (other
than the posterior model). Only a likelihood needs to be specified, which can
take various forms such as loss functions and synthetic likelihoods, thus
providing a form of a likelihood-free approach. Furthermore, we formulate the
approach such that the posterior samples can either be independent of, or
conditionally dependent upon the inputs to the primary model. The latter
approach is shown to be capable of increasing the apparent complexity of the
primary model. We see this being useful in applications such as surrogate and
physics-based models. To promote how the Bayesian paradigm offers more than
just uncertainty quantification, we demonstrate: uncertainty quantification,
multi-modality, as well as an application with a recent deep forecasting
neural network architecture.

##  1 Introduction

The Bayesian statistical paradigm offers a principled approach to representing
uncertainty in both a model and its predictions by treating the model
parameters as latent random variables. Inference of the latent parameters is,
however, generally intractable in complex and nonlinear models such as
Bayesian Neural Networks (BNNs). Approximate inference approaches are
employed, with Variational Inference (VI) and Markov Chain Monte Carlo (MCMC)
being the most popular [1, 2]. The challenges with these approaches involve
balancing complexity and expressiveness.

MCMC methods, such as Hamiltonian Monte Carlo (HMC) [3], provide accurate
samples from the posterior and are considered the gold standard approach for
inference [4]. This however comes at the cost of computational complexity and
poor scalability due to burn-in periods, sample rejection, and the requirement
to process an entire dataset in order to make a new proposal. Stochastic
gradient MCMC methods (e.g. [5, 6]) address this by rather performing MCMC
through “noisy backpropagation”. Such approaches have however been shown to
compromise the scalability of HMC [7].

VI approaches rely on optimisation to provide a tractable and scalable
approach. They have enjoyed extensive application in deep learning models such
as the Variational Auto-Encoder (VAE) [8] and the BNN [9]. However, the
accuracy of VI is compromised when the variational family is restricted to
simple parametric distributions as to promote tractability. These parametric
distributions generally lack the expressiveness required for BNNs.
Furthermore, VI can tend to underestimate the variance [10, 11]; especially
under the commonly adopted Mean Field VI (MFVI) assumption.

Recently, several studies have considered replacing parametric distributions
with highly expressive implicit distributions [12, 13, 14, 15]. Implicit
distributions include the generators in Generative Adversarial Networks [16]
and have relation to Normalising Flows [17]. Direct application of VI to
implicit distributions is however challenging as these distributions cannot be
evaluated, but only sampled from. To overcome this, the problem is reframed
into an adversarial problem where a density ratio within the Evidence Lower
Bound (ELBO) is approximated with a discriminator [18]. The discriminator is
only equal to the density ratio if it is a Bayes-optimal classifier, thus
requiring a highly expressive discriminator and increased computational
resources [19]. Furthermore, density ratio estimation is also challenging in
high-dimensional spaces [20].

To avoid density ratio estimation, alternative approaches have been proposed
by defining semi-implicit variational approximations [19, 21, 22]. These
approaches use a hierarchical framework where the variational distribution is
explicit, but an implicit mixing distribution is imposed on the parameters of
the variational distribution. The approach however does not directly optimise
using the ELBO, but a surrogate thereof [19, 21]. To directly optimise
according to the ELBO, Titsias and Ruiz [21] use and MCMC approach, which is
however computationally expensive.

The commonality between the discussed approaches is that they all strive to
conform to the MCMC or VI approach. We aim to devise an alternative which is
both expressive and scalable to large data.

$\boldsymbol{x}$$\boldsymbol{y}$$\boldsymbol{z}$$\boldsymbol{\theta}$$\boldsymbol{\phi}$Primary
model
$f_{\boldsymbol{\theta}}(\boldsymbol{x})$$g_{\boldsymbol{\phi}}(\boldsymbol{z})$Implicit
posterior model (a) Unconditioned posterior model.

$\boldsymbol{x}$$\boldsymbol{y}$$\boldsymbol{x}$$\boldsymbol{z}$$\boldsymbol{\theta}$$\boldsymbol{\phi}$Primary
model
$f_{\boldsymbol{\theta}}(\boldsymbol{x})$$g_{\boldsymbol{\phi}}(\boldsymbol{z},\boldsymbol{x})$Implicit
posterior model (b) Conditional posterior model.

Figure 1: Architecture of the hierarchical Bayesian neural model, with input
vectors $\boldsymbol{x}$ (predictor variables) and $\boldsymbol{z}$ (i.i.d.
uniform random variates). Subplot (a) shows the architecture where the
implicit posterior is parametrically static and modelled as independent of
$\boldsymbol{x}$ (the inputs to the primary model). Subplot (b) shows the more
flexible solution, where realisations from the implicit posterior model are
conditioned on the input vector $\boldsymbol{x}$ with $\boldsymbol{\theta}$
providing a local (rather than global) parameterisation of the primary model.
In both cases, the vector $\boldsymbol{y}$ is a vector that parameterises a
likelihood function. The parameters $\boldsymbol{\phi}$ of the overall model
are optimised using gradient descent in order to maximise the posterior
predictive distribution.

As illustrated in Figure 1, we propose the use of two models: (1) a primary
predictive model (e.g. a regressor or classifier), $f_{\boldsymbol{\theta}}$,
defined by parameter vector $\boldsymbol{\theta}$, and (2) a posterior model
$g_{\boldsymbol{\phi}}$ that generates random samples of the parameter vector
$\boldsymbol{\theta}$. A likelihood function is defined along with the primary
model, and the posterior model approximates the (implicit) posterior
distribution over the parameters of the primary model. The parameters of the
primary model are indirectly trained by optimising (via gradient descent) the
parameters, $\boldsymbol{\phi}$, of $g_{\boldsymbol{\phi}}$ in order to
maximise the posterior predictive distribution.

The key contributions of this work are:

  1. 1.

We propose a simple and natural inference approach which optimises a posterior
model according to the posterior predictive distribution. Other than the
posterior model, our only approximation is a Monte Carlo estimate of the
posterior predictive distribution which has an affinity to gradient descent
optimisation, making it scalable to large datasets. Only a likelihood function
(which is not limited to probability distributions) and the posterior model
need be specified. The expressive implicit posterior model is the default,
however parametric distributions are not excluded.

  2. 2.

We show that by modelling the conditional posterior (conditional on the
primary model inputs), the apparent complexity of the primary model can be
greatly increased. This increase is owing to the posterior model producing a
unique set of parameters for any given input and therefore results in models
with locally adaptive behaviour.

  3. 3.

We demonstrate that the implicit model used for the posterior distribution is
highly flexible and capable of inducing multi-modality in the predictive
distribution where appropriate.

We see this work having impact in applications where uncertainty
quantification is required such as medical, robotics, finance, and forecasting
applications. The approach may also be useful in applications concerned with
out-of-distribution data and problems relating to transfer learning. We see
the conditional posterior model being highly applicable in problems relating
to physics-based modelling, emulator/surrogate models, and anomaly detection.

##  2 Methods

###  2.1 The Posterior Predictive Distribution Formulation

The aim in a regression or classification problem is to train a model
$f_{\boldsymbol{\theta}}$ parametrised by $\boldsymbol{\theta}$, using a
dataset $\mathcal{D}$ =
$\\{(\boldsymbol{x}_{1},\boldsymbol{y}_{1}),\dots,(\boldsymbol{x}_{N},\boldsymbol{y}_{N})\\}$
of input-output pairs, such that, given a new input
$\acute{\boldsymbol{x}}\notin\mathcal{D}$, the model is able to accurately
predict the target $\boldsymbol{y}$ associated with $\acute{\boldsymbol{x}}$.
In the Bayesian statistical approach, $\boldsymbol{\theta}$ is treated as a
latent random variable and following the assimilation of observational data,
our beliefs about which parameter values are acceptable can be updated (i.e.
the posterior distribution). As such, a neural network modelled in this way is
referred to as a Bayesian Neural Network (BNN) and can generate an ensemble of
predictions $\acute{\boldsymbol{y}}_{1},\dots,\acute{\boldsymbol{y}}_{L}$ for
a given $\acute{\boldsymbol{x}}$ through the posterior predictive
distribution. The overall challenge when using the Bayesian statistical
approach is to find, estimate, or approximate the posterior distribution over
$\boldsymbol{\theta}$, which is given by Bayes’ theorem
$p(\boldsymbol{\theta}|\mathcal{D})=p(\mathcal{D}|\boldsymbol{\theta})p(\boldsymbol{\theta})/p(\mathcal{D})$.
The marginal distribution in the denominator is generally intractable and
approximate inference approaches are frequently required to infer the
posterior.

In this study, we propose a simple approach whereby we attempt to directly
infer the posterior predictive distribution, which is the object of inference
for making probabilistic predictions. Given a new input-output pair
$(\acute{\boldsymbol{x}},\acute{\boldsymbol{y}})\notin\mathcal{D}$, the
posterior predictive distribution is given by [23] $\displaystyle
p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\mathcal{D})$
$\displaystyle=\int
p(\acute{\boldsymbol{y}},\boldsymbol{\theta}|\acute{\boldsymbol{x}},\mathcal{D})d\boldsymbol{\theta}\hskip
199.16928pt$ $\displaystyle\hskip 29.02193pt=$ $\displaystyle\int
p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\boldsymbol{\theta})p(\boldsymbol{\theta}|\mathcal{D})d\boldsymbol{\theta},$
$\boldsymbol{\theta}$ conditionally independent of $\acute{\boldsymbol{x}}$,
(1a) $\displaystyle\hskip 29.02193pt=$ $\displaystyle\int
p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\boldsymbol{\theta})p(\boldsymbol{\theta}|\acute{\boldsymbol{x}},\mathcal{D})d\boldsymbol{\theta},$
$\boldsymbol{\theta}$ conditionally dependent on $\acute{\boldsymbol{x}}$ (1b)
The annotations of (1a) and (1b) correspond to Figure 1(a) and Figure 1(b) and
correspond to the situation where the posterior distribution is conditionally
independent and dependent111(1b) may appear like amortisation [8], however the
intention is to condition the posterior on $\acute{\boldsymbol{x}}$. on
$\acute{\boldsymbol{x}}$ respectively. We refer to these two models as (a) the
unconditioned posterior model and (b), the conditional posterior model. See
Appendix A for graphical models and a derivation.

The first factor in (1a) and (1b) is the likelihood of observing
$\acute{\boldsymbol{y}}$ given $\acute{\boldsymbol{x}}$ and
$\boldsymbol{\theta}$. This likelihood can be evaluated via the elements that
are output from the primary model
$f_{\boldsymbol{\theta}}(\acute{\boldsymbol{x}})$ and the targets
$\acute{\boldsymbol{y}}$ from $\mathcal{D}$. The second factor is the
posterior distribution given by Bayes theorem, which can be analytically
intractable, or often only known up to a constant of proportionality. We
propose to approximate the posterior with the distribution
$q_{\boldsymbol{\phi}|\mathcal{D}}(\boldsymbol{\theta})$, which is
parametrised by $\boldsymbol{\phi}$ and optimised with $\mathcal{D}$. The goal
is then to find the parameters $\boldsymbol{\phi}$ that maximise the posterior
predictive distribution:

$\displaystyle\boldsymbol{\phi}^{*}$
$\displaystyle\triangleq\arg\max_{\boldsymbol{\phi}}\int
p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\boldsymbol{\theta})q_{\boldsymbol{\phi}|\mathcal{D}}(\boldsymbol{\theta})d\boldsymbol{\theta}.$
(2)

Suppose we draw $L$ samples
$\boldsymbol{\theta}^{(1)},\dots.\boldsymbol{\theta}^{(L)}$ from
$q_{\boldsymbol{\phi}|\mathcal{D}}(\boldsymbol{\theta})$. The Monte Carlo
estimate of the posterior predictive distribution is

$\displaystyle p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\mathcal{D})$
$\displaystyle\simeq\frac{1}{L}\sum_{l=1}^{L}p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\boldsymbol{\theta}^{(l)})$
(3)

To optimise the model, the negative of (3) can be used as a loss function.
When we define the posterior model
$g_{\boldsymbol{\phi}}(\boldsymbol{z},\acute{\boldsymbol{x}})\sim
q_{\boldsymbol{\phi}|\mathcal{D}}(\boldsymbol{\theta})$ (where
$\acute{\boldsymbol{x}}$ is included with (1b)), this loss function is given
by

$\displaystyle\mathcal{L}$
$\displaystyle=-\frac{1}{L}\sum_{l=1}^{L}p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},g_{\boldsymbol{\phi}}(\boldsymbol{z}^{(l)},\acute{\boldsymbol{x}}))$
(4)

In practice, it may be useful to consider the negative log posterior estimate
and make use of the log-sum-exp formulation for numerical stability. The
gradient of the loss function is

$\displaystyle\frac{\partial\mathcal{L}}{\partial\boldsymbol{\phi}}$
$\displaystyle=-\frac{1}{L}\sum_{l=1}^{L}\frac{\partial
p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},g_{\boldsymbol{\phi}}(\boldsymbol{z}^{(l)},\acute{\boldsymbol{x}}))}{\partial\boldsymbol{\phi}}\frac{\partial
g_{\boldsymbol{\phi}}(\boldsymbol{z}^{(l)},\acute{\boldsymbol{x}})}{\partial\boldsymbol{\phi}}$
(5)

Note that, VI has the challenge of finding an unbiased, low variance Monte
Carlo estimate of the ELBO’s gradient [24], whereas, in our approach, the
target is a Monte Carlo estimate whose gradient is tractable. Given this
gradient, the gradient-descent optimisation objective is then

$\displaystyle\boldsymbol{\phi}^{*}\triangleq\arg\min_{\boldsymbol{\phi}}\left[-\frac{1}{L}\sum_{l=1}^{L}p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},g_{\boldsymbol{\phi}}(\boldsymbol{z}^{(l)},\acute{\boldsymbol{x}}))\right]$
(6)

A discussion on convergence in gradient descent optimisation for this form is
provided in Appendix B.

Following Figure 1 and equation (6), $f_{\boldsymbol{\theta}}$ is the primary
model, intended to specify the parameters of a likelihood (or approximate
likelihood) function
$p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\boldsymbol{\theta})$. For a
sample of parameter vectors
$\boldsymbol{\theta}^{(1)},\dots,\boldsymbol{\theta}^{(L)}$, the loss function
in (4) can be evaluated. The sampled vectors $\boldsymbol{\theta}^{(l)}$ are
generated by the approximate posterior model
$g_{\boldsymbol{\phi}}(\boldsymbol{z},\acute{\boldsymbol{x}})$ (parametrised
by $\boldsymbol{\phi}$) that is intended to target
$p(\boldsymbol{\theta}|\mathcal{D})$.

###  2.2 Primary Model and Likelihood

The primary model $f_{\boldsymbol{\theta}}$, that is responsible for
parametrising the likelihood function, can conceivably be any model with
parameter vector $\boldsymbol{\theta}$. Candidate primary models are
parametrised functions that are differentiable with respect to their
parameters. These include linear models, deep neural networks, and various
physics-based models.

With an implicit posterior model, the likelihood associated with the primary
model is not constrained to probability distributions. The likelihood operates
as a loss or distance function, whose only requirements are those associated
with gradient descent optimisation, such as convexity and differentiability
[25]. This is especially useful when a given loss function does not have a
known associated density (or mass), such as heuristic loss functions, scaled
errors in time-series forecasting, and the hinge losses in classification.
Likelihood-free approaches such as approximate Bayesian Computation and
synthetic likelihood approaches could also be adopted.

###  2.3 Posterior Model

The implicit posterior distribution is a generative process that transforms
random samples $\boldsymbol{z}^{(l)}$ from some parametric distribution
$\boldsymbol{z}$ with density $s(\cdot)$ (e.g. a uniform or standard normal)
to samples from the posterior distribution $\boldsymbol{\theta}^{(l)}$
according to a transform $g_{\boldsymbol{\phi}}(\boldsymbol{z})$. If the
transform $g_{\boldsymbol{\phi}}$ is a Borel-measurable function,
$q_{\boldsymbol{\phi}|\mathcal{D}}(\boldsymbol{\theta})$ is a valid density
given by [26, 27, 28]

$\displaystyle
q_{\boldsymbol{\phi}|\mathcal{D}}(\boldsymbol{\theta})=\frac{\partial}{\partial\theta_{1}}\cdots\frac{\partial}{\partial\theta_{m}}\int_{\\{\boldsymbol{z}:g_{\boldsymbol{\phi}}(\boldsymbol{z})\leq\boldsymbol{\theta}\\}}s(\boldsymbol{z})d\boldsymbol{z}$
(7)

The posterior model $g_{\boldsymbol{\phi}}(\boldsymbol{z})$ can take various
forms such as a neural network (which is sometimes referred to as a
HyperNetwork [29]). As the posterior model models an implicit distribution, it
makes no assumptions about the form of the posterior distribution, and only
provides a means for sampling from it.

The posterior model described in Figure 1(a) operates by transforming samples
from $\boldsymbol{z}$ to samples from the posterior. These generated samples
can directly be applied in equation (4). In the conditional posterior model
described in Figure 1(b), $\acute{\boldsymbol{x}}$ is included as an input to
the posterior model. Conceptually, this allows the posterior model to produce
the best primary model parameters for the given input. The result is the
apparent complexity of the primary model is increased. Conditioning on
$\acute{\boldsymbol{x}}$ is similar to the idea of the conditional variational
autoencoder (VAE) [30] and the conditional generative adversarial network
(GAN) [31].

The posterior model does not exclude parametric distributions. A parametric
distribution could be directly applied as an unconditioned posterior model, or
a linear implicit model could be used. A mixture density network (MDN) [32]
can be used to provide conditional posterior model. Similar to the
architecture of the variational autoencoder’s encoder [8], the MDN transforms
$\acute{\boldsymbol{x}}$ through a neural network to a posterior parametric
mixture density, where samples can be drawn.

###  2.4 Incorporating a Prior

The prior $p(\boldsymbol{\theta})$ lurks within the posterior represented in
(1a) and (1b). Without explicitly specifying it, we assume an improper,
uninformative prior distribution that is constant across the domain of the
parameters. This may be considered particularly convenient in context of
black-box models such as neural networks when there is frequently no auxiliary
information that can be drawn upon to formulate a prior. Informative priors
may also be difficult to represent given that such models often have high-
dimensional parameter spaces and may exhibit complex patterns of dependence.

In some contexts, it may however be advantageous to include a prior. A weakly
informative prior can be incorporated through regularisation. In fact, the
objective of a prior is to regularize the posterior distribution, constraining
it to contain the bulk of the probability mass within reasonable bounds [23,
33]. Weakly informative priors may thus be enforced by (1) regularising the
primary model with methods such as dropout (2) limiting the output of
$g_{\boldsymbol{\phi}}$ to regularize $\boldsymbol{\theta}$ to a specific
range using an activation function, (3) specifying parametric distributions on
the posterior and likelihood to force the prior into a family of distributions
by conjugacy, (4) restricting the dimension of $\boldsymbol{z}$ to form a
degenerate distribution on a union of manifolds within the domain of
$\boldsymbol{\theta}$ [34], and (5) restricting the complexity of
$g_{\boldsymbol{\phi}}$ to regularize the posterior distribution complexity.

Such regularisation approaches can however encourage degeneracy in the
posterior as a degenerate distribution may be the simplest representation for
the implicit model. In the extreme case, this may be a Dirac density (e.g. see
Appendix C).

###  2.5 Model Training

To train the model, the parameters of $g_{\boldsymbol{\phi}}$ are optimised
via gradient descent such that $g_{\boldsymbol{\phi}}$ produces a distribution
over $\boldsymbol{\theta}$ that maximises the posterior predictive
distribution. The parameters of the primary model $\boldsymbol{\theta}$ are
thus indirectly optimised by optimising $\boldsymbol{\phi}$. The algorithm for
training the model is provided in Algorithm 1. The use of minibatches is
discussed in Appendix D.

primary model $f_{\boldsymbol{\theta}}$, posterior model
$g_{\boldsymbol{\phi}}$, training data $\mathcal{D}$, minibatch size $B$,
number of MC samples $L$, number of epochs $K$.

Initialise the posterior model parameters $\boldsymbol{\phi}$ randomly.

for $i\in(1,K)$ do

for minibatch $n\in\mathcal{D}$ do

for $l\in(1,L)$ do

Sample $\boldsymbol{z}$ (depending on the architecture of
$g_{\boldsymbol{\phi}}$).

Sample $\boldsymbol{\theta}^{(l)}$ from
$\boldsymbol{g}_{\boldsymbol{\phi}}(\boldsymbol{z})$.

Pass $\boldsymbol{x}_{n}$ through $f(\boldsymbol{\theta}^{(l)})$ to produce
$\hat{\boldsymbol{y}}_{n}$.

Evaluate the likelihood
$p(\boldsymbol{y}_{n}|\boldsymbol{x}_{n},\boldsymbol{\theta}^{(l)})$ given
$\hat{\boldsymbol{y}}_{n}$.

Compute the loss
$\mathcal{L}_{n}=-\frac{1}{L}\sum_{l=1}^{L}p(\boldsymbol{y}_{n}|\boldsymbol{x}_{n},\boldsymbol{\theta}^{(l)})$.

Compute the gradients
$\frac{\partial\mathcal{L}_{n}}{\partial\boldsymbol{\phi}}$.

Backpropagate and update $\boldsymbol{\phi}$.

return $\boldsymbol{\phi}$

Algorithm 1 Model training.

Given that the Monte Carlo estimate of the posterior predictive distribution
in (3) is based on the likelihood function and that an uninformed prior is
assumed, it is easy to see that the approach has the potential to converge
towards a degenerate distribution at the maximum likelihood (see Appendix C).
To avoid this, the use of early stopping is essential and is used in all
experiments.

Note that the posterior model has been described as a model over the complete
set of $\boldsymbol{\theta}$. However, in practice it may be convenient to
define a separate posterior model for subsets of $\boldsymbol{\theta}$. For
example, if $f_{\boldsymbol{\theta}}$ is a neural network, a separate
posterior model can be defined for each layer in $f_{\boldsymbol{\theta}}$.
This reduces memory capacity as it assumes independence of the parameters
across layers and it is more convenient to implement in deep learning
frameworks.

##  3 Results

We demonstrate initial results relating to our contributions described in
Section 1. The aim is to demonstrate our approach and its features. For this,
a pedagogical comparison is taken to highlight differences between various
approaches and provoke thought. A detailed performance comparison (e.g.
uncertainty coverage) across multiple datasets and models is reserved for
future work.

To provide effective demonstrations, synthetic datasets are constructed with
specific properties. We also test our approach in a forecasting problem with a
larger dataset and more complex primary model in Section 3.4.

Additional results are provided in Appendices E and F.

###  3.1 Bayesian Inference with Uncertainty Estimates

We consider a synthetic dataset of containing 32 samples from the function
$y=x\sin(x)$ and additional high-variance samples {(7,-7), (8.5,7), (10,-7),
(11.5,7)} [35]. The additional high-variance samples (originally intended to
emulate a heteroscedastic process) are used to demonstrate how the different
models adapt to structural variations in the data. The test dataset contains
input samples at a finer resolution than the training dataset (1024 samples).
Furthermore, the range of the test inputs extends beyond the range of the
training set to provide of out-of-distribution samples.

We demonstrate the proposed approach and compare it with a VI approach (Bayes
by Backprop) [9], the HMC approach [3], and the prior-contrastive adversarial
VI approach [18, 36]. For the proposed approach, we consider conditional,
unconditioned, and MDN posterior models. In all cases, a Multilayer-Perceptron
(MLP) serves as a regression model and the likelihood is assumed to be a
Gaussian with a known standard deviation of 0.01. The models are trained with
10 Monte Carlo samples and 30 Monte Carlo samples are drawn from the posterior
predictive distribution during testing. Model configurations are described in
Table 1 and the results are illustrated in Figure 2.

Table 1: Model architecture details. Model |  Architecture
---|---
Primary model |  MLP with a hidden layer comprising 512 ReLU hidden units and a linear output.
Conditional Posterior |  MLP with 4 uniform random inputs with the input $\boldsymbol{x}$ and 16 ReLU hidden units.
Unconditioned Posterior |  MLP with 4 uniform random inputs and 16 ReLU hidden units.
MDN Posterior |  MLP with the input $\boldsymbol{x}$, 16 ReLU hidden units, and unimodal Gaussians at the output.
VI |  Mean field variational assumption with a spike and slab priors as in [9]
HMC |  Factorized Gaussian prior. The algorithm is run over 4000 steps with 25 leapfrog steps with step size 0.001.
Prior-contrastive |  Generative model equivalent to the conditional posterior model and a MLP discriminator with 512 ReLU hidden units.
Figure 2: Results for the MLP primary model. Blue markers are the targets, the
grey curves are 30 MC samples from the posterior predictive distribution, and
the green curve is the mean of these samples. Note that these figures plot
standardised data.

All models treat the high-variance samples differently. The conditional
posterior model seems to embrace the change in the structure of the data as
its variance expands to include these samples. The unconditioned model appears
to have discovered two possible paths through the high-variance data samples,
suggesting a bimodal posterior. The MDN and prior-contrastive adversarial
models tend to provide a close fit the high-variance samples, whereas VI seems
to have treated these samples as outliers. HMC seems to partially fit to these
samples, but its variance remains relatively low.

The variance of the models all expand over the out-of-distribution regions to
some degree. HMC has a rapidly expanding variance in the out-of-distribution
region to the left of the data. The MDN model also has a rapidly expanding
variance, however its mean is significantly different to that of HMC. The
variance of the conditional posterior model and the prior-contrastive
adversarial model increase in the out-of-distribution regions, but not as
significantly as HMC. VI generally seems to underestimate the variance overall
as expected [11].

###  3.2 Increased Model Complexity

We consider testing the posterior models with a simple linear regression
primary model (a straight line) on the non-linear $y=x\sin(x)$ dataset
described in Section 3.1. This allows us to assess the affect of the
conditional posterior model on the primary model. The results are presented in
Figure 3.

Figure 3: Results with a linear regression model for the primary model. Blue
markers are the targets, the grey curves are 30 MC samples from the posterior
predictive distribution, and the green curve is the mean of these samples.
Note that these figures plot standardised data. Figure 4: A plot of a sample
of primary model parameters and the prediction with these parameters. With a
linear primary model, the parameters correspond to a gradient and offset of a
straight line. The key observation is that the parameters vary over
$\boldsymbol{x}$ to produce a non-linear prediction.

What is of interest is that the conditional posterior approach fits the non-
linear data with a linear model. In comparison, the other models fit the data
with straight lines as expected. The conditional posterior model has learned
to localise the parameters of the primary model and vary them according to the
input $\acute{\boldsymbol{x}}$. The parameters $\boldsymbol{\theta}$ are
varied over the continuous space of $\acute{\boldsymbol{x}}$ as illustrated in
Figure 4. This result is similar to local regression methods [37, 38], but
over a continuous range. A reasonable consequence is that a simpler primary
model can be used, reducing overall complexity of the approach.

###  3.3 Likelihood Flexibility

The posterior distribution is only constrained by the likelihood function and
could, for example, be multimodal. A multimodal posterior would suggest
several variations of parameters corresponding to several output options for a
given input. For example, this can occur in tracking applications where
targets can split and merge, or inverse problems with one-to-many mappings.

To test this, a synthetic dataset based on $x\sin(x)$ (with scaled inputs) is
generated according to

$\displaystyle y=\begin{cases}a(x)\sin(a(x))+\epsilon&x\in(0,0.6)\\\
a(x)\sin(a(x))+1+\epsilon&x\in(0.3,1)\end{cases}$ (8)

where $\epsilon\sim\mathcal{N}(0,0.01)$ and $a(x)=10x-5$ scales the inputs to
the range (0,1). The key feature of this function is that there are two
parallel curves over the domain $(0.3,0.6)$. A total of 128 samples are drawn
from this function ensuring that each curve has the same number of samples.
Ideally, the model should predict the lower curve over $(0,0.3)$, either
curves over $(0.3,0.6)$, and the upper curve over $(0.6,1)$.

To achieve varying parameters over $\boldsymbol{x}$, the conditional posterior
model is used. To encourage a bimodal posterior distribution, we consider
using the $L1$-norm as a likelihood function. Note that the $L1$-norm is not a
density function, highlighting the flexibility of the likelihood
specification. To support non-smooth transitions between curves, the primary
model should use non-smooth activation functions, such as the ReLU. The
results are illustrated in left panel of Figure 5. The MC samples from the
model (grey curves), illustrate how the model is able to switch between the
different curves over the domain $(0.3,0.6)$. According to the $L1$-norm, both
parallel curves have the same likelihood and the model is permitted to switch
between them at any time. Owing to the conditioning on
$\acute{\boldsymbol{x}}$, the model reverts to a unimodal posterior and
predicts one of the appropriate curves outside of the domain $(0.3,0.6)$.

Figure 5: Model predictions for a multimodal dataset. Left: using the
$L1$-norm as a likelihood function. Right: including sample labels as inputs
to the posterior model. Blue markers are the targets and the grey curves plot
30 Monte Carlo prediction samples. The posterior model is bimodal over the
domain $(0.3,0.6)$ where it can predict either curve, and is unimodal
otherwise.

If there are labels indicating which curve the targets originate from, these
labels could be included as inputs to the posterior model along with
$\acute{\boldsymbol{x}}$. These labels provide an indication to the posterior
model of a potential multimodal posterior distribution. To demonstrate this, a
label is provided as an input to the conditional posterior model and a
unimodal Gaussian likelihood is used. The labels are generated by labelling
even indexed samples with zeros and odd indexed samples with ones in the
domain $(0.3,0.6)$. During testing, the labels are randomly assigned. The
results are illustrated in the right panel of Figure 5. The model learns a
clearer distinction between the curves compared with the $L1$-norm-based
approach. Note that a bimodal posterior is learned despite a unimodal
likelihood.

###  3.4 Forecasting Demonstration

The proposed approach is applied to the N-BEATS model [39] in a multi-step-
ahead forecasting task. N-BEATS is used as the primary model, and we compare
results with the unconditioned posterior model, the conditional posterior
model, and standard maximum likelihood training. The dataset used is the
monthly average temperature in England over the period of January 1723 to
December 1970 from Rob Hyndman’s Time-Series Data Library222Rob Hyndman’s
Time-Series Data Library: https://pkg.yangzhuoranyang.com/tsdl/. LIC: GPL-3..
The model is configured to forecast 3 months ahead given the last 6 months of
data. Input-output pairs are created with a 6-sample input window and its
adjacent 3-sample output window, which are slid across the dataset. A total of
2960 input-output pairs are formed, which are split into a training and
testing set comprising 2075 and 885 samples respectively. The model
architecture details are provided in Table 2. The models are trained over 100
epochs using ADAM with a learning rate of 0.01 and a batch size of 128.

Table 2: Model architecture details for the forecasting results. Model |  Architecture
---|---
N-BEATS |  Input sequence length: 24, output sequence length: 12, stacks: 1, blocks: 3, block type: generic $\theta$-dimension: 32, FC stack layer dimension: 64, block weights: shared.
Posterior models |  MLP with 4 uniform random inputs (and the input $\boldsymbol{x}$ for the conditional model model) and 16 hidden neurons with ReLU activation functions.

The forecasting results are provided in Table 3 in terms of Root Mean Squared
Error (RMSE), Mean Absolute Percentage Error (MAPE), and the epoch run time.
As a baseline results for Naïve forecasting method are also provided. The
results demonstrate that the standard N-BEATS model produces the highest
errors and the N-BEATS model with a conditional posterior model produces the
lowest errors. N-BEATS is able to produce lower errors owing to the increased
modelling capacity through the conditional posterior distribution. With the
unconditioned posterior model, N-BEATS is able to generalise better than the
standard training owing to less overfitting, where Bayesian models are more
inclined to find “flat minima” that provide better generalisation [40].

The epoch times demonstrate an increased computational complexity of the
proposed approach compared to the standard maximum likelihood training. This
increase is due to the additional complexity of the posterior model
architecture and the multiple MC samples. This cost however comes with the
benefits of model uncertainty and improved generalisation. Furthermore, memory
requirements could be reduced to linear scaling by using weight normalization
reparameterisation and only outputting the scaling factors from the posterior
model [12].

Table 3: Average errors over the test dataset. Epoch times are provided for a 2.6 GHz 6-Core Intel Core i7 with 16GB RAM. Note that the relatively high epoch time for the conditional posterior model is due to the minibatch strategy (see Appendix D). Model | RMSE | RMSE std dev. | MAPE | MAPE std dev. | Epoch run time (s)
---|---|---|---|---|---
Naive | 4.47 | 2.13 | 72.5 | 135.0 | N/A
N-BEATS | 2.87 | 1.11 | 55.6 | 156.9 | 0.22
N-BEATS with unconditioned posterior | 1.86 | 0.76 | 35.7 | 101.9 | 1.31
N-BEATS with conditional posterior | 1.35 | 0.64 | 27.8 | 94.9 | 4.2

##  4 Conclusion

We have proposed a novel approach for Bayesian inference in complex models and
demonstrated key features of the approach. We show how the approach is able to
provide uncertainty estimates in high variance samples and out-of-distribution
regions, model multi-modal posterior distributions, improve capacity of a
primary model, and provide better generalisation in a forecasting task.

The approach does not rely on MCMC or VI and provides flexibility in the
likelihood specification. The model targets the posterior predictive
distribution in training, which is natural as this is generally the primary
distribution of interest in practice. Training is performed via gradient
descent, providing scalability to large datasets.

Future work involves applying the approach to the domains described in Section
1, especially the physics-based and surrogate modelling domains where we see
the conditional posterior approach having impact.

## Acknowledgements

We would like to thank Edwin Bonilla for the various discussions around this
topic of this paper. This work was supported by the CSIRO MLAI Future Science
Platform.

## References

  * Jospin et al. [2022] Laurent Valentin Jospin, Hamid Laga, Farid Boussaid, Wray Buntine, and Mohammed Bennamoun.  Hands-on bayesian neural networks—a tutorial for deep learning users.  _IEEE Computational Intelligence Magazine_ , 17(2):29–48, 2022.  doi: 10.1109/MCI.2022.3155327.
  * Abdar et al. [2021] Moloud Abdar, Farhad Pourpanah, Sadiq Hussain, Dana Rezazadegan, Li Liu, Mohammad Ghavamzadeh, Paul Fieguth, Xiaochun Cao, Abbas Khosravi, U. Rajendra Acharya, Vladimir Makarenkov, and Saeid Nahavandi.  A review of uncertainty quantification in deep learning: Techniques, applications and challenges.  _Information Fusion_ , 76:243–297, 2021.  ISSN 1566-2535.  doi: https://doi.org/10.1016/j.inffus.2021.05.008.
  * Neal [2011] Radford M Neal.  Mcmc using hamiltonian dynamics.  _Handbook of markov chain monte carlo_ , 2(11):2, 2011.
  * Goan and Fookes [2020] Ethan Goan and Clinton Fookes.  Bayesian neural networks: An introduction and survey.  In _Case Studies in Applied Bayesian Data Science_ , pages 45–87. Springer, 2020.
  * Welling and Teh [2011] Max Welling and Yee W Teh.  Bayesian learning via stochastic gradient langevin dynamics.  In _Proceedings of the 28th International Conference on Machine Learning (ICML-11)_ , pages 681–688, 2011.
  * Chen et al. [2014] Tianqi Chen, Emily Fox, and Carlos Guestrin.  Stochastic gradient hamiltonian monte carlo.  In _International conference on machine learning_ , pages 1683–1691. PMLR, 2014.
  * Betancourt [2015] Michael Betancourt.  The fundamental incompatibility of scalable hamiltonian monte carlo and naive data subsampling.  In _International Conference on Machine Learning_ , pages 533–540. PMLR, 2015.
  * Kingma and Welling [2014] Diederik P. Kingma and Max Welling.  Auto-encoding variational bayes.  In _Proceedings of the Second International Conference on Learning Representations (ICLR 2014)_ , April 2014.
  * Blundell et al. [2015] Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra.  Weight uncertainty in neural network.  In Francis Bach and David Blei, editors, _Proceedings of the 32nd International Conference on Machine Learning_ , volume 37 of _Proceedings of Machine Learning Research_ , pages 1613–1622, Lille, France, 07–09 Jul 2015. PMLR.
  * Blei et al. [2017] David M Blei, Alp Kucukelbir, and Jon D McAuliffe.  Variational inference: A review for statisticians.  _Journal of the American statistical Association_ , 112(518):859–877, 2017.
  * Bishop [2006] C.M. Bishop.  _Pattern Recognition and Machine Learning_.  Information Science and Statistics. Springer, 2006.  ISBN 9780387310732.
  * Krueger et al. [2017] David Krueger, Chin-Wei Huang, Riashat Islam, Ryan Turner, Alexandre Lacoste, and Aaron Courville.  Bayesian hypernetworks.  In _2017 NeurIPS Workshop on Bayesian Deep Learning_. Neural Information Processing Systems (NeurIPS), 2017.
  * Pawlowski et al. [2017] Nick Pawlowski, Andrew Brock, Matthew CH Lee, Martin Rajchl, and Ben Glocker.  Implicit weight uncertainty in neural networks.  In _2017 NeurIPS Workshop on Bayesian Deep Learning_. Neural Information Processing Systems (NeurIPS), 2017.
  * Wang et al. [2018] Kuan-Chieh Wang, Paul Vicol, James Lucas, Li Gu, Roger Grosse, and Richard Zemel.  Adversarial distillation of bayesian neural network posteriors.  In _International Conference on Machine Learning_ , pages 5190–5199. PMLR, 2018.
  * Henning et al. [2018] Christian Henning, Johannes von Oswald, João Sacramento, Simone C Surace, Jean-Pascal Pfister, and Benjamin F Grewe.  Approximating the predictive distribution via adversarially-trained hypernetworks.  In _2018 NeurIPS Workshop on Bayesian Deep Learning_ , page 121. Neural Information Processing Systems (NeurIPS), 2018.
  * Goodfellow et al. [2014] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial nets.  In Z. Ghahramani, M. Welling, C. Cortes, N. D. Lawrence, and K. Q. Weinberger, editors, _Advances in Neural Information Processing Systems 27_ , pages 2672–2680. Curran Associates, Inc., 2014.
  * Papamakarios et al. [2021] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan.  Normalizing flows for probabilistic modeling and inference.  _Journal of Machine Learning Research_ , 22(57):1–64, 2021.
  * Huszár [2017] Ferenc Huszár.  Variational inference using implicit distributions.  _arXiv:1702.08235_ , 2017.
  * Yin and Zhou [2018] Mingzhang Yin and Mingyuan Zhou.  Semi-implicit variational inference.  In _International Conference on Machine Learning_ , pages 5660–5669. PMLR, 2018.
  * Sugiyama et al. [2012] Masashi Sugiyama, Taiji Suzuki, and Takafumi Kanamori.  _Density Ratio Estimation in Machine Learning_.  Cambridge University Press, 2012.  doi: 10.1017/CBO9781139035613.
  * Titsias and Ruiz [2019] Michalis K Titsias and Francisco Ruiz.  Unbiased implicit variational inference.  In _The 22nd International Conference on Artificial Intelligence and Statistics_ , pages 167–176. PMLR, 2019.
  * Moens et al. [2021] Vincent Moens, Hang Ren, Alexandre Maraval, Rasul Tutunov, Jun Wang, and Haitham Ammar.  Efficient semi-implicit variational inference.  _arXiv:2101.06070_ , 2021.
  * Gelman et al. [2013] A. Gelman, J.B. Carlin, H.S. Stern, D.B. Dunson, A. Vehtari, and D.B. Rubin.  _Bayesian Data Analysis, Third Edition_.  Chapman & Hall/CRC Texts in Statistical Science. Taylor & Francis, 2013\.  ISBN 9781439840955.
  * Mohamed et al. [2020] Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih.  Monte carlo gradient estimation in machine learning.  _J. Mach. Learn. Res._ , 21(132):1–62, 2020.
  * Boyd and Vandenberghe [2004] S. Boyd and L. Vandenberghe.  _Convex Optimization_.  Cambridge University Press, 2004.  ISBN 9781107394001.
  * Mohamed and Lakshminarayanan [2016] Shakir Mohamed and Balaji Lakshminarayanan.  Learning in implicit generative models.  _arXiv:1610.03483_ , 2016.
  * Tran et al. [2017] Dustin Tran, Rajesh Ranganath, and David Blei.  Hierarchical implicit models and likelihood-free variational inference.  In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.
  * Wasserman [2013] L. Wasserman.  _All of Statistics: A Concise Course in Statistical Inference_.  Springer Texts in Statistics. Springer New York, 2013.  ISBN 9780387217369.
  * Ha et al. [2016] David Ha, Andrew Dai, and Quoc V Le.  Hypernetworks.  _arXiv:1609.09106_ , 2016.
  * Sohn et al. [2015] Kihyuk Sohn, Honglak Lee, and Xinchen Yan.  Learning structured output representation using deep conditional generative models.  _Advances in neural information processing systems_ , 28:3483–3491, 2015.
  * Mirza and Osindero [2014] Mehdi Mirza and Simon Osindero.  Conditional generative adversarial nets.  _arXiv:1411.1784_ , 2014.
  * Bishop [1994] Christopher M Bishop.  Mixture density networks.  Technical Report NCRG94004, Aston University, 1994.
  * Jaynes et al. [2003] E.T. Jaynes, E.T.J. Jaynes, G.L. Bretthorst, and Cambridge University Press.  _Probability Theory: The Logic of Science_.  Cambridge University Press, 2003.  ISBN 9780521592710.
  * Arjovsky and Bottou [2017] Martin Arjovsky and Léon Bottou.  Towards principled methods for training generative adversarial networks.  _arXiv:1701.04862_ , 2017.
  * Gal [2016] Yarin Gal.  _Uncertainty in deep learning_.  PhD thesis, University of Cambridge, 2016.
  * Mescheder et al. [2017] Lars Mescheder, Sebastian Nowozin, and Andreas Geiger.  Adversarial variational bayes: Unifying variational autoencoders and generative adversarial networks.  In _International Conference on Machine Learning_ , pages 2391–2400. PMLR, 2017.
  * Loader [2006] C. Loader.  _Local Regression and Likelihood_.  Statistics and Computing. Springer New York, 2006.  ISBN 9780387227320.
  * Hastie et al. [2013] T. Hastie, R. Tibshirani, and J. Friedman.  _The Elements of Statistical Learning: Data Mining, Inference, and Prediction_.  Springer Series in Statistics. Springer New York, 2013.  ISBN 9780387216065.
  * Oreshkin et al. [2020] Boris N. Oreshkin, Dmitri Carpov, Nicolas Chapados, and Yoshua Bengio.  N-beats: Neural basis expansion analysis for interpretable time series forecasting.  In _International Conference on Learning Representations_ , 2020.
  * Wilson and Izmailov [2020] Andrew G Wilson and Pavel Izmailov.  Bayesian deep learning and a probabilistic perspective of generalization.  _Advances in neural information processing systems_ , 33:4697–4708, 2020.
  * Bubeck [2015] Sébastien Bubeck.  _Convex Optimization: Algorithms and Complexity_.  Foundations and trends in machine learning. Now Publishers, 2015.  ISBN 9781601988607.

##  Appendix A Posterior Predictive Distribution Proof

The proof of (1a) and (1b) is provided using $d$-separation:

###### Proof.

Consider the dataset input-output pairs given by $\mathcal{D}$ =
$\\{(\boldsymbol{x}_{i},\boldsymbol{y}_{i})\\}_{i=1}^{N}$ and a primary model
$f_{\boldsymbol{\theta}}(\boldsymbol{x})$ parameterised by
$\boldsymbol{\theta}$. The conditional independence assumptions are
illustrated in the graphical model presented in Figure 6. The posterior
predictive distribution for a new input-output pair
$(\acute{\boldsymbol{x}},\acute{\boldsymbol{y}})\notin\mathcal{D}$ is given by

$\displaystyle p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\mathcal{D})$
$\displaystyle=\int
p(\acute{\boldsymbol{y}},\boldsymbol{\theta}|\acute{\boldsymbol{x}},\mathcal{D})d\boldsymbol{\theta}$

By the definition of conditional probability

$\displaystyle p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\mathcal{D})$
$\displaystyle=\int
p(\acute{\boldsymbol{y}}|\boldsymbol{\theta},\acute{\boldsymbol{x}},\mathcal{D})p(\boldsymbol{\theta}|\acute{\boldsymbol{x}},\mathcal{D})d\boldsymbol{\theta}$
(9)

In the first factor, $\acute{\boldsymbol{y}}$ are independent of $\mathcal{D}$
given that the inputs are i.i.d and given that the outputs are conditionally
independent given $\boldsymbol{\theta}$, which produces

$\displaystyle p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\mathcal{D})$
$\displaystyle=\int
p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\boldsymbol{\theta})p(\boldsymbol{\theta}|\acute{\boldsymbol{x}},\mathcal{D})d\boldsymbol{\theta}$
(10)

This result corresponds to the conditional posterior model described in
equation (1b) and Figure 1(b). In the the second factor, $\boldsymbol{\theta}$
is conditioned on $\acute{\boldsymbol{x}}$ as indicated by the dashed link in
Figure 6. If this conditioning were removed, $\boldsymbol{\theta}$ is
conditionally independent on $\acute{\boldsymbol{x}}$ given that
$\acute{\boldsymbol{y}}$ is not observed. Thus

$\displaystyle p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\mathcal{D})$
$\displaystyle=\int
p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\boldsymbol{\theta})p(\boldsymbol{\theta}|\mathcal{D})d\boldsymbol{\theta}$
(11)

This result corresponds to the unconditional posterior model described in
equation (1a) and Figure 1(a) (and is the common representation of the
posterior predictive distribution in the literature). ∎

$N$$\boldsymbol{x}_{i}$$\boldsymbol{y}_{i}$$\boldsymbol{\theta}$$\acute{\boldsymbol{x}}$$\acute{\boldsymbol{y}}$
Figure 6: Graphical model for general Bayesian modelling. The data
$\mathcal{D}$ =
$\\{(\boldsymbol{x}_{1},\boldsymbol{y}_{1}),\dots,(\boldsymbol{x}_{N},\boldsymbol{y}_{N})\\}$
are assumed i.i.d. and represented using plate notation, the inputs
$\boldsymbol{x}$ are not treated as random variables, and the outputs
$\boldsymbol{y}$ are assumed conditionally independent given the model
parameters $\boldsymbol{\theta}$. The dashed line corresponds to the
conditional posterior model in Figure 1(b), where $\boldsymbol{\theta}$ is
conditioned on $\acute{\boldsymbol{x}}$.

##  Appendix B Convergence

Our approach is to encapsulate the Bayesian inference problem within a
frequentist optimisation problem. In this frequentist approach, we treat the
inputs $\boldsymbol{z}$ and $\boldsymbol{x}$, and the targets $\boldsymbol{y}$
as random, and we use them to find a point estimate of $\boldsymbol{\phi}$. As
described in (2), this optimisation results in inferring the posterior
distribution of $\boldsymbol{\theta}$ given $\mathcal{D}$.

In this context, our approach conforms to the standard stochastic gradient
descent (SGD) problem and thus converges accordingly. That is, random samples
(or batches) of inputs ($\boldsymbol{z}$ and $\boldsymbol{x}$) and the targets
($\boldsymbol{y}$) are used to maximise $\boldsymbol{\phi}$ according to a
“likelihood” given by (4). Various proofs and discussions on the convergence
of SGD are available in the literature, such as in Bubeck [41, Ch. 6].

Note that, although (4) is treated as a likelihood in the encapsulating
frequentist problem, it is the MC estimate of the posterior predictive
distribution which directs the Bayesian inference problem. By optimising
$\boldsymbol{\phi}$ with maximum likelihood, we indirectly infer the posterior
distribution of $\boldsymbol{\theta}$ given $\mathcal{D}$.

The maximum likelihood optimisation approach provides insight into the
requirements for the likelihood function
$p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\boldsymbol{\theta}^{(l)})$.
In general, convergence of SGD requires that the objective $\mathcal{L}$ is
convex and differentiable. Suppose
$-p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\boldsymbol{\theta}^{(l)})$
is given by some distance or loss function
$c(\acute{\boldsymbol{y}},\hat{\boldsymbol{y}}^{(l)})$ (e.g. mean squared
error), where
$\hat{\boldsymbol{y}}^{(l)}=f(\acute{\boldsymbol{x}};\boldsymbol{\theta}^{(l)})$
and $\boldsymbol{\theta}^{(l)}=g_{\boldsymbol{\phi}}(\boldsymbol{z}^{(l)})$.
The loss function in (4) is then given by

$\displaystyle\mathcal{L}$
$\displaystyle=-\frac{1}{L}\sum_{l=1}^{L}p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},g_{\boldsymbol{\phi}}(\boldsymbol{z}^{(l)}))$
$\displaystyle=\frac{1}{L}\sum_{l=1}^{L}c(\acute{\boldsymbol{y}},f(\acute{\boldsymbol{x}};g_{\boldsymbol{\phi}}(\boldsymbol{z}^{(l)})))$
(12)

Given that the sum of convex functions is also convex, if $c$ is convex then
$\mathcal{L}$ is convex (e.g. if $c$ is convex,
$\nabla_{\boldsymbol{\phi}}^{2}c_{l}\geq
0\implies\sum_{l=1}^{L}\nabla_{\boldsymbol{\phi}}^{2}c_{l}\geq
0\implies\sum_{l=1}^{L}c_{l}$ is convex). As a result, $c$ can generally be
any distance or cost function that is used in machine learning literature,
where a key requirement is differentiability. Note however that, as with
training standard neural networks, if $f$ and $g$ are non-linear. $c$ may be
non-convex and the solution is only locally optimal.

##  Appendix C Degenerate Implicit Posterior

A proof by construction (or proof by example) approach is taken to show that
the implicit posterior can produce a degenerate distribution which is located
at the maximum likelihood, as described in sections 2.4 and 2.5. (Note that
the selected example makes use of the unconditioned posterior model for
notational convenience. The example is however trivially extended to the
conditional posterior model case, which includes both $\boldsymbol{z}$ and
$\boldsymbol{x}$ as inputs.)

###### Proof.

Consider an unconditioned posterior model in the form of an MLP with $K$
layers, weights $\boldsymbol{W}_{1:M}$ biases $\boldsymbol{b}_{1:M}$, and
activation functions $\psi()$:

$\displaystyle\boldsymbol{\theta}=(\psi_{\boldsymbol{W}_{K},\boldsymbol{b}_{K}}\circ\psi_{\boldsymbol{W}_{K-1},\boldsymbol{b}_{K-1}}\circ\cdots\circ\psi_{\boldsymbol{W}_{2},\boldsymbol{b}_{2}}\circ\psi_{\boldsymbol{W}_{1},\boldsymbol{b}_{1}})(\boldsymbol{z})$
(13)

If $\boldsymbol{W}_{1}=\mathbf{0}$

$\displaystyle\psi_{\boldsymbol{W}_{1},\boldsymbol{b}_{1}}(\boldsymbol{z})=\psi(\boldsymbol{W}_{1}\boldsymbol{z}+\boldsymbol{b}_{1})=\psi(\mathbf{0}\boldsymbol{z}+b_{1})=\psi(\boldsymbol{b}_{1})$
(14)

The inputs $\boldsymbol{z}$ are ignored by the MLP resulting in a constant
output with respect to $\boldsymbol{z}$ (the same result is achieved in the
conditional posterior model when the weights associated with both
$\boldsymbol{z}$ and $\boldsymbol{x}$ are zero). The parameters
$\boldsymbol{\phi}$ however can still be optimised to provide an optimal, but
fixed output for $\boldsymbol{\theta}$, denoted by
$\tilde{\boldsymbol{\theta}}$.

Consider the objective function in (4)

$\displaystyle\mathcal{L}$
$\displaystyle=-\frac{1}{L}\sum_{l=1}^{L}p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},g_{\boldsymbol{\phi}}(\boldsymbol{z}^{(l)}))$
(15)

If $g_{\boldsymbol{\phi}}$ is deterministic and fixed

$\displaystyle\tilde{\boldsymbol{\theta}}=g_{\boldsymbol{\phi}}(\boldsymbol{z}^{(i)})=g_{\boldsymbol{\phi}}(\boldsymbol{z}^{(j)})~{}\forall~{}i,j\in(1,\dots,L)$
(16)

and thus

$\displaystyle
p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\tilde{\boldsymbol{\theta}})=p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},g_{\boldsymbol{\phi}}(\boldsymbol{z}^{(i)}))=p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},g_{\boldsymbol{\phi}}(\boldsymbol{z}^{(j)}))~{}\forall~{}i,j\in(1,\dots,L)$
(17)

The MC estimate of the posterior predictive distribution is then

$\displaystyle\mathcal{L}=-\frac{1}{L}\sum_{l=1}^{L}p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\tilde{\boldsymbol{\theta}}))=-p(\acute{\boldsymbol{y}}|\acute{\boldsymbol{x}},\tilde{\boldsymbol{\theta}})$
(18)

which is the negative likelihood of $\acute{\boldsymbol{y}}$ given
$\tilde{\boldsymbol{\theta}}$. Minimising $g_{\boldsymbol{\phi}}$ produces the
optimal parameters for $f_{\boldsymbol{\theta}}$ according to this likelihood,
which produces a maximum likelihood estimate. ∎

From the Bayesian perspective, the proved result is a degenerate Dirac
distribution that is located at the maximum likelihood solution. Note that
with an uninformed prior, the maximum likelihood estimate is equivalent to the
maximum a posteriori estimate.

##  Appendix D Minibatches

The standard approach to using minibatches in stochastic gradient descent can
be used in the unconditioned posterior model. However, in the conditional
model, if $\acute{\boldsymbol{x}}$ is provided in batch form, the output of
the posterior model will be a batch of primary model parameters
$\boldsymbol{\theta}$: a set of primary parameters associated with each
$\acute{\boldsymbol{x}}$ in the batch. As the batched inputs are passed
through the primary model, each element in the batch be associated with the
corresponding element in parameter batch. This is easily achieved via matrix
multiplication, however, the tensors containing the variables should be
carefully arranged. An illustration comparing standard batch processing and
the modified batch processing is provided in Figure 7. Note that the modified
batch approach is equivalent to stochastic gradient descent with a single
sample. The result is that the only limit on batch size is memory capacity.

+WeightMatrixInputBatchBias

+WeightMatrixInputBatchBias

Figure 7: Illustration of the linear operations for batch processing. On the
left is the standard batch processing operation with a linear product of a
$\mathbb{R}^{6\times 5}$ weight matrix, an $\mathbb{R}^{5\times 3}$ input with
a batch size of 3, and a $\mathbb{R}^{6\times 1}$ bias vector duplicated for
each batch. On the right is the modified approach. The three
$\mathbb{R}^{6\times 5}$ weight matrices corresponding to a batch size of 3
form a $\mathbb{R}^{6\times 5\times 3}$ tensor, the three bias vectors
corresponding to a batch size of 3 form a $\mathbb{R}^{6\times 1\times 3}$
tensor, and the input batch matrix is rearranged into a $\mathbb{R}^{5\times
1\times 3}$ tensor. The matrix operations are standard matrix/vector
operations over the first and second dimensions of the tensors. (For
illustrative purposes, the batch dimension is last, but in practice it is more
appropriate to have it first.)

##  Appendix E Posterior Distribution Model Configurations

We consider varying the architecture and number of MC samples used in training
for the posterior model. Results are plotted in Figure 8 and Figure 9. The
columns specify the number of Monte Carlo samples $L$ used when training the
model. The rows specify a configuration of the posterior model using a list
containing two or three integers. In the case of two integers, the posterior
model contains an input layer, an output layer, and no hidden layer. As such,
the posterior model provides a linear translation of its inputs
$\boldsymbol{z}$. In the case where three integers are provided in the list, a
MLP with one hidden layer is used, where the number of units in the hidden
layer are provided by the second integer in the list. When a hidden layer is
used, the ReLU activation function applies.

Considering the number of MC samples, a single MC sample does not seem to
provide sufficient capacity to model the posterior distribution. However, for
this dataset, there is no significant difference between 10 and 100 MC
samples, suggesting that 10 samples are sufficient. In general, this is
confirmed by Bishop [11], who suggests that MC estimate of an integral only
requires a few samples if the samples are i.i.d. In the case of the posterior
models, $\boldsymbol{z}$ are indeed i.i.d. randomly generated samples.

The first three rows in the plots use a linear transfer function where no
hidden layer in the posterior model is used. Especially for the unconditioned
models, the out-of-distribution variation to the left of the data is larger
for the linear models than for the MLP models. The higher complexity seems to
encourage more certainty around that region. By conditioning on the inputs
$\acute{\boldsymbol{x}}$, it does seem to improve the uncertainty.
Furthermore, owing to the conditioning on $\acute{\boldsymbol{x}}$, the
conditional models generally produce more complex curves around the high
variance samples.

When the posterior model maps from a lower dimension to a higher dimension,
the posterior distribution will lie on a union of manifolds within the domain
of $\boldsymbol{\theta}$ [34]. With the highest model complexity $[N,N,N]$,
the posterior is able to span the domain of $\boldsymbol{\theta}$. The
similarity of the results from models with a small number of dimensions (e.g.
the $[1]$ and $[5,16,N]$ models) with the $[N,N,N]$ posterior model suggest
that these models are still able to provide reasonable uncertainty estimates.
A smaller posterior model can significantly reduce the overall computational
complexity.

Figure 8: Unconditioned posterior model. Columns are plots for 1, 10, and 100
MC samples used during training. Rows plot results with variations in the
posterior model architecture, where the brackets denote the number of units in
the layers of the network. All models have $N$ outputs. Blue markers are the
targets, the grey curves are 30 MC samples from the posterior predictive
distribution, and the green curve is the mean of these samples. Note that
these figures plot standardised data. Figure 9: Conditional posterior model.
Columns are plots for 1, 10, and 100 MC samples. Rows plot results with 1
input ([1]), 32 inputs ([32]), 5 inputs with 16 hidden neurons ([5,16]), $N$
inputs ([NN]), and $N$ inputs with $N$ hidden layers ([NNN]). All models have
$N$ outputs. Blue markers are the targets, the grey curves are 30 MC samples
from the posterior predictive distribution, and the green curve is the mean of
these samples. Note that these figures plot standardised data.

##  Appendix F Prior Inclusion Results

We demonstrate including a prior by adding the sum-of-squares regulariser in
the loss function presented in (4). If the likelihood function is chosen as
the sum of squared errors the loss function is given by

$\displaystyle\mathcal{L}$
$\displaystyle=-\frac{1}{L}\sum_{l=1}^{L}\bigg{(}\underbrace{-\sum_{i=1}^{N}(y_{i}-y_{i}^{\prime})^{2}}_{\text{likelihood}}-\underbrace{\vphantom{\sum_{i=1}^{N}}\lambda\left(\boldsymbol{\theta}^{(l)}\right)^{T}\boldsymbol{\theta}^{(l)}}_{\text{regulariser}}\bigg{)}$
(19)

Note the resemblance of contents within the parentheses to the maximum a
posterior (MAP) estimate where
$\lambda(\boldsymbol{\theta}^{(l)})^{T}\boldsymbol{\theta}^{(l)}$ relates to a
zero mean isotropic Gaussian with precision $2\lambda$. The key difference is
that this regularized loss function is computed within a MC sum over
parameters sampled from the posterior distribution.

The results are presented in Figure 10. Comparing these results to the results
with no regularisation presented in Figure 2, the left panel shows slight
under-fitting of the data owing to the regularization. The variance increases
significantly around the high variance samples and the out-of-distribution
regions. Note the similarity of this result to the result presented for the
linear model in Figure 3. This illustrates that the regularization is
penalising the complexity of the primary model as desired. When the
regularisation constant $\lambda$ is increased to an extreme value, the model
produces a horizontal line for all MC samples as illustrated in the right
panel. Here the model has fitted the data with an offset, which corresponds to
the mean of the targets $\boldsymbol{y}$ (which is zero). With such high
regularisation, the model be forced into a restricted region of
$\boldsymbol{\theta}$ resulting in a degenerate Dirac distribution for the
posterior predictive distribution as described in Section 2.4 and Appendix C.

Figure 10: Results with a regulariser introduced for prior information. Blue
markers are the targets, the grey curves are 30 MC samples from the posterior
predictive distribution, and the green curve is the mean of these samples.
Note that these figures plot standardised data.