On Feynman--Kac training of partial Bayesian neural networks
Recently, partial Bayesian neural networks (pBNNs), which only consider a
subset of the parameters to be stochastic, were shown to perform competitively
with full Bayesian neural networks. However, pBNNs are often multi-modal in the
latent variable space and thus challenging to approximate with parametric
models. To address this problem, we propose an efficient sampling-based
training strategy, wherein the training of a pBNN is formulated as simulating a
Feynman--Kac model. We then describe variations of sequential Monte Carlo
samplers that allow us to simultaneously estimate the parameters and the latent
posterior distribution of this model at a tractable computational cost. Using
various synthetic and real-world datasets we show that our proposed training
scheme outperforms the state of the art in terms of predictive performance.
Zheng Zhao, Sebastian Mair, Thomas B. Schön, Jens Sjölund
# On Feynman-Kac training of partial Bayesian neural networks

###### Abstract

Recently, partial Bayesian neural networks (pBNNs), which only consider a subset of the parameters to be stochastic, were shown to perform competitively with full Bayesian neural networks. However, pBNNs are often multi-modal in the latent-variable space and thus challenging to approximate with parametric models. To address this problem, we propose an efficient sampling-based training strategy, wherein the training of a pBNN is formulated as simulating a Feynman-Kac model. We then describe variations of sequential Monte Carlo samplers that allow us to simultaneously estimate the parameters and the latent posterior distribution of this model at a tractable computational cost. We show on various synthetic and real-world datasets that our proposed training scheme outperforms the state of the art in terms of predictive performance.

## 1 Introduction

Bayesian neural networks (BNNs) are an important class of machine learning models for quantifying uncertainty. However, computing BNN posterior distributions is an open challenge (Izmailov et al., 2021) since many standard statistical inference methods such as Markov chain Monte Carlo (MCMC) are poorly suited to the combination of a high-dimensional parameter space and a massive number of data points (Papamakrou et al., 2022). Many approximate solutions have been proposed to overcome this computational challenge, for example, MCMC with stochastic likelihoods (Andrieu and Roberts, 2009; Welling and Teh, 2011; Chen et al., 2014; Zhang et al., 2020), parametric approximations of the posterior distributions in the form of variational Bayes (Blundell et al., 2015; Gal and Ghahramani, 2016) or stochastic averaging methods (Izmailov et al., 2018; Maddox et al., 2019). In a similar vein, Izmailov et al. (2020) conduct variational inference in a subspace of the parameter space. Other ways include, for instance, restricting the inference only on the last layer (Ober and Rasmussen, 2019; Kristiadi et al., 2020).

In this paper, we focus on the training of partial Bayesian neural networks (pBNNs), a family of BNNs where only part of the model parameters are random variables. This is motivated by the recent work of, for example, Daxberger et al. (2021) and Sharma et al. (2023), which show that pBNNs and standard BNNs are comparable in terms of prediction performance despite the dimensionality of the random variable in the pBNNs being much smaller than standard BNNs. While both approaches make posterior inference much easier, it comes at the cost of making the pBNN a latent-variable model that requires estimation of both the deterministic parameters and posterior inference.

To be more precise, let \(x\mapsto f(x;\phi,\psi)\) be a neural network parametrised by a deterministic parameter \(\psi\in\mathbb{R}^{w}\) and a random variable \(\phi\in\mathbb{R}^{d}\) that follows a given prior distribution \(\pi(\phi)\). Suppose that we have a dataset \(\{(x_{n},y_{n})\}_{n=1}^{N}\) with an associated likelihood conditional density function \(p(y_{n}\mid\phi;\psi)\)1, and that the observation random variables are independent conditioned on \(\phi\). The goal of training is twofold. First, to learn the deterministic parameter \(\psi\) from the dataset, and second, to compute the posterior distribution \(p(\phi\mid y_{1:N};\psi)\), where \(y_{1:N}\coloneqq\{y_{1},y_{2},\ldots,y_{N}\}\). This is a classical parameter estimation problem in a latent variable model (Cappe et al., 2005). The main benefit of pBNNs is that when the dimension \(d\ll w\) is relatively small, the second goal of computing the posterior distribution with fixed \(\psi\) is generally tractable. However, the first goal of estimating the parameter \(\psi\) remains to be solved.

Footnote 1: Throughout the paper we omit the data covariate \(x_{n}\) of \(y_{n}\) to make the notation cleaner.

For practitioners, a popular approach is to jointly optimise \(\phi\) and \(\psi\) according to the maximum a posteriori
(MAP) objective \((\phi,\psi)\mapsto\log p(y_{1:N}\mid\phi;\psi)+\log p(\phi)\), see, for example, Daxberger et al. (2021) and Sharma et al. (2023). Then, the estimated \(\psi\) is used to compute the posterior distribution over \(\phi\) by any applicable Bayesian inference method (e.g., MCMC). The advantages of this method lie in the simplicity and fast computation. Moreover, there is a plethora of optimisers (e.g., stochastic gradient descent and its variants) specifically designed for the setting of large \(w\) and \(N\). However, for multi-modal distributions the MAP estimated \(\psi\) and \(\phi\) are prone to be restricted on a single mode.

In contrast, statisticians generally favour maximum likelihood estimation (MLE) for estimating \(\psi\). Compared to the aforementioned MAP approach, the MLE objective takes all possible outcomes of the random variable \(\phi\) into account by integrating out \(\phi\), and the estimator can be consistent (Cappe et al., 2005). However, the MLE method is not directly applicable to pBNNs since the marginal likelihood \(p(y_{1:N};\psi)=\int p(y_{1:N}\mid\phi;\psi)\,\pi(\phi)\,\mathrm{d}\phi\) is in general intractable. Instead, various lower bounds on the marginal likelihood are used by, for instance, brute-force Monte Carlo, expectation maximisation (EM, Bishop, 2006), and variational Bayes (VB, Blei et al., 2017; Sjolund, 2023). However, EM and VB often compromise in using parametric families of distributions, and they often incur Monte Carlo computation overheads to approximate expectations. Another common approach is to transform the (gradient of) MLE into an expectation with respect to the posterior distribution by Fisher's identity (see, e.g., Cappe et al., 2005, Chap. 10), and then use an MCMC sampler to approximate the expectation (at the cost of demanding computations).

In this paper, we study how to apply and adapt sequential Monte Carlo (SMC) methods to train pBNNs, by representing the training as a simulation of a Feynman-Kac model (Chopin and Papaspiliopoulos, 2020). More specifically, the Feynman-Kac model is composed of a sequence of potential functions given by the likelihood model and a sequence of invariant Markov kernels that so as to anneal to the target posterior distribution \(p(\phi\mid y_{1:N};\psi)\). Computing the target posterior distribution then amounts to a sequential simulation of the Feynman-Kac model, and SMC is the natural sampling framework for the model. The motivation for studying SMC samplers to train pBNNs is that SMC samplers are able to simultaneously produce consistent MLEs and to sample from the posterior distribution while retaining a tractable computational cost. Compared to most MCMC-based methods, SMC samplers are immediately parallelisable (Lee et al., 2010) by leveraging graphics processing units, and are easier to calibrate. On the other hand, parallelising MCMC is still an open problem (see, e.g., Jacob et al., 2020).

Our contributions are as follows. (i) We show and discuss how to apply and adapt SMC samplers to train pBNNs (Section 2). (ii) We propose approximate SMC samplers that are scalable in the number of data points and are therefore better suited to pBNNs (Section 3). (iii) We benchmark the proposed samplers on a variety of synthetic and real-world datasets, and the results show that the proposed training scheme is state of the art in terms of predictive performance (Section 4).

## 2 Training via Feynman-Kac

Recall that we aim to find a maximum likelihood estimate for the deterministic part of the pBNN and then to compute the posterior distribution \(p(\phi\mid y_{1:N};\psi)\) of the stochastic part of the pBNN based on the learnt parameter \(\psi\). In this section, we recap a sequential online algorithm to recursively compute the target posterior distribution via a Feynman-Kac model and how to jointly estimate the parameter \(\psi\).

To ease the exposition of the idea, let us for now suppose that the deterministic variable \(\psi\) in the pBNN is fixed. Due to the conditional independence of the observations, the posterior distribution admits a recursion

\[p(\phi\mid y_{1:n};\psi)=\frac{p(y_{n}\mid\phi;\psi)}{z_{n}(\psi)}\,p(\phi\mid y _{1:n-1};\psi) \tag{1}\]

for \(n=1,2,\ldots,N\), where \(z_{n}(\psi)\coloneqq p(y_{n}\mid y_{1:n-1};\psi)\) is the normalising constant, and the initial is defined by \(p(\phi\mid y_{1:0};\psi)\coloneqq\pi(\phi)\). If we have computed the distribution \(p(\phi\mid y_{1:n-1};\psi)\) for some \(n\), we can then compute the next \(p(\phi\mid y_{1:n};\psi)\) by Equation (1). Ultimately, we continue the iteration until \(n=N\) to reach the target posterior distribution. This recursion is the gist of sequential online learning or annealing.

As in most Bayesian inference problems, the challenge lies in the computationally intractable normalising constant \(z_{n}(\psi)\). Hence, in practice, we often use a tractable sequence of approximations \(Q_{N}\coloneqq\{q_{n}(\phi\mid y_{1:n};\psi)\colon n=1,2,\ldots,N\}\) such that \(q_{n}(\phi\mid y_{1:n};\psi)\propto q_{n-1}(\phi\mid y_{1:n-1};\psi)\) is computable for all \(n\)'s. A convenient choice is a Gaussian \(q_{n}\) and then to use, for example, Taylor expansions, variational Bayes (Opper, 1999), or Gauss quadrature methods (Golub and Meurant, 2010) to run the sequence. However, in the context of BNNs, Gaussian approximations can result in large errors (Foong et al., 2020) particularly when the true posterior distribution is non-Gaussian (e.g., multi-modal ones). This motivates us to use a Monte Carlo (MC)-based method to come up with such an approximate sequence.
### Sequential Monte Carlo sampling

Sequential Monte Carlo (SMC) samplers (Del Moral et al., 2006; Chopin and Papaspiliopoulos, 2020) are natural Monte Carlo (MC) methods for approximating the target posterior distribution in the sequential learning framework. Specifically, we choose \(Q_{N}=\{S_{n}^{J}\colon n=1,2,\ldots,N\}\), where \(S_{n}^{J}\coloneqq\{(w_{n,j},\phi_{n,j})\colon j=1,2,\ldots,J\}\) are \(J\) weighted Monte Carlo samples that approximately represent \(p(\phi\mid y_{1:n};\psi)\). Suppose that we are able to draw the initial samples \(S_{0}^{J}\sim\pi(\phi)\) from the given prior, then we can recursively compute \(S_{n}^{J}\) for any \(n=1,2,\ldots,N\) via

\[\phi_{n,j} =\phi_{n-1,j}, \tag{2}\] \[\overline{w}_{n,j} =w_{n-1,j}\,p(y_{n}\mid\phi_{n-1,j};\psi),\]

and normalisation \(w_{n,j}=\overline{w}_{n,j}\,/\,\sum_{i=1}^{J}\overline{w}_{n,i}\) for all samples \(j=1,2,\ldots,J\). This approximation is consistent in the sense that the resulting Dirac measure given by \(S_{n}^{J}\) converges weakly to that of \(p(\phi\mid y_{1:n};\psi)\) as \(J\to\infty\) for all \(n\)(Chopin and Papaspiliopoulos, 2020). Furthermore, the method is particularly favourable in the pBNN context, as the dimension \(d\) for the stochastic part of the pBNN is usually not large which in turn allows us to use significantly more MC samples compared to full BNNs. This method dates back to Neal (2001) who simulates static target distributions with a likelihood tempering (see also Chopin, 2002).

However, the SMC chain using Equation (2) rarely works in reality, since the samples \(\{\phi_{n,j}\}_{j=1}^{J}\) will become less informative and the weights \(\{w_{n,j}\}_{j=1}^{J}\) will degenerate as \(n\) increases (Del Moral et al., 2006). In practice, we additionally introduce a Markov transition kernel \(h_{n}(\cdot\mid\phi_{n-1})\) between each step \((n-1,n)\), so that the samples \(\{\phi_{n,j}\}_{j=1}^{J}\) are perturbed according to the kernel. Specifically, the update of samples in Equation (2) modifies to

\[\phi_{n,j}\mid\phi_{n-1,j}\sim h_{n}(\phi_{n,j}\mid\phi_{n,j-1}). \tag{3}\]

Then, we arrive at the so-called (marginal) Feynman-Kac model

\[q_{N}(\phi_{N}\mid y_{1:N};\psi) \tag{4}\] \[\coloneqq\frac{1}{l_{N}(\psi)}\int\prod_{n=1}^{N}p(y_{n}\mid\phi _{n};\psi)\prod_{n=1}^{N}h_{n}(\phi_{n}\mid\phi_{n-1})\] \[\qquad\qquad\qquad\times\pi(\phi_{0})\,\mathrm{d}\phi_{0:N-1},\]

where \(l_{N}(\psi)\) is the normalising constant (see the definition of Feynman-Kac in Chopin and Papaspiliopoulos, 2020, Chap. 5). Now, to guarantee that the terminal \(q_{N}(\phi_{N}\mid y_{1:N};\psi)\) exactly hits the target posterior distribution \(p(\phi\mid y_{1:N};\psi)\), we need to choose the Markov kernel in a way that \(h_{n}\) leaves the previous posterior distribution \(p(\phi\mid y_{1:n-1};\psi)\) invariant, viz., \(\int h_{n}(\phi\mid\phi_{n-1})\,\rho(\phi_{n-1})\,\mathrm{d}\phi_{n-1}=p(\phi \mid y_{1:n-1};\psi)\) for any operating distribution \(\rho\). Note that \(h_{n}\) indeed depends on \(y_{1:n-1}\) and \(\psi\), but we omit them for clean notation. This choice of Markov kernel also ensures that the marginal likelihood is given by \(l_{N}(\psi)=\prod_{n=1}^{N}z_{n}(\psi)=p(y_{1:N};\psi)\).

It is now clear that computing the target posterior distribution \(p(\phi\mid y_{1:N};\psi)\) amounts to solving the Feynman-Kac model in Equation (4). The SMC sampler simulates the Feynman-Kac model with the weighted samples \(S_{n}^{J}\) for \(n=0,1,\ldots,N\) using Equations (2) and (3). Moreover, recall that we need to estimate the parameter \(\psi\) via MLE

\[\operatorname*{arg\,min}_{\psi\in\mathbb{R}^{w}}-\log l_{N}(\psi).\]

It turns out that the SMC sampler produces a consistent estimate \(z_{n}(\psi)\approx\sum_{j=1}^{J}\overline{w}_{n,j}\), so that we can estimate the log-likelihood via \(\log l_{N}(\psi)\approx\sum_{n=1}^{N}\log\bigl{(}\sum_{j=1}^{J}\overline{w}_{ n,j}\bigr{)}\) as a by-product of the computation of the posterior distributions. Then, we can solve the optimisation problem above by using any optimiser for training neural networks. For a detailed exposition of SMC parameter estimators, see, for example, Johansen et al. (2008), Schon et al. (2011), and Kantas et al. (2015). We summarise the SMC training of pBNNs in Algorithm 1 together with a gradient descent-based optimisation of the MLE objective. Within the algorithm, we use the shorthand \(\ell_{n}(\cdot)\) for the approximation to the marginal log-likelihood \(\log l_{n}(\cdot)\).

```
0: Training data \(\{(x_{i},y_{i})\}_{i=1}^{N}\), number of samples \(J\), initial parameter \(\psi_{0}\), learning rate function \(r\)
0: The MLE estimate \(\psi_{i}\) and weighted posterior samples \(\{(w_{N,j},\phi_{N,j})\}_{j=1}^{J}\sim p(\phi\mid y_{1:N};\psi_{i})\) for\(i=1,2,\ldots\) untilconvergentdo
0: Draw \(\{\phi_{0,j}\}_{j=1}^{J}\sim\pi(\phi)\) \(w_{0,j}=1\,/\,J\) for all \(j=1,2,\ldots,J\) \(\ell_{0}(\psi_{i-1})=0\) for\(n=1\)to\(N\)do// Parallelise\(j\) Resample \(\{(w_{n-1,j},\phi_{n-1,j})\}_{j=1}^{J}\) if needed Draw \(\phi_{n,j}\mid\phi_{n-1,j}\sim h_{n}(\phi_{n,j}\mid\phi_{n,j-1})\) \(\overline{w}_{n,j}=w_{n-1,j}\,p(y_{n}\mid\phi_{n,j};\psi_{i-1})\) \(\ell_{n}(\psi_{i-1})=\ell_{n-1}(\psi_{i-1})-\log\bigl{(}\sum_{j=1}^{J} \overline{w}_{n,j}\bigr{)}\) \(w_{n,j}=\overline{w}_{n,j}\,/\,\sum_{k=1}^{J}\overline{w}_{n,k}\)  end for
11:\(\psi_{i}=\psi_{i-1}-r(i)\,\nabla\ell_{N}(\psi_{i-1})\)
```

**Algorithm 1**SMC sampler for pBNN
Scalable sequential Monte Carlo samplers

We still have a few critical challenges left to solve in order to apply the SMC sampler in Algorithm 1 for training pBNNs. First, the gradient computation \(\nabla\ell_{N}\) might be biased due to the non-differentiability of the resampling and Markov transition steps. Second, it is not easy to design the Markov kernel \(h_{n}\) and also to compute the Markov moves. Third, the algorithm does not scale well in the number of data points \(N\). In particular, for each gradient descent step to compute \(\nabla\ell_{N}\), we need a complete SMC loop over the entire dataset. In what follows, we detail these issues and show how to tackle them.

### Compensating gradient biases

The non-differentiability is a notorious issue for sequential Monte Carlo samplers, since most resampling techniques (e.g., systematic and stratified) and MCMC algorithms incur discrete randomness. This in turns induces biases in the gradient computation \(\nabla\ell_{N}\). There are recent developments to tackle the differentiability issue by, for example, optimal transport-based smooth resampling (Corenflos et al., 2021) and coupled MCMC (Arya et al., 2022, 2023), but they come at the cost of introducing additional computation costs and calibrations. However, we can in fact avoid differentiating the SMC algorithm by invoking Fisher's identity (see, e.g., Cappe et al., 2005, Chap. 10), which states that \(\nabla\mathrm{log}\,l_{N}(\psi)=\int\nabla\mathrm{log}\,p(y_{1:N},\phi;\psi) \,p(\phi\mid y_{1:N};\psi)\,\mathrm{d}\phi\). Moreover, unlike particle filtering in the system identification context (see, e.g., Poyiadjis et al., 2011, Sec. 2.2), the model design facilitates \(\nabla\mathrm{log}\,p(y_{1:N},\phi;\psi)=\nabla\mathrm{log}\,p(y_{1:N}\mid \phi;\psi)\). Hence, we can modify the gradient computation in Algorithm 1 to

\[\nabla\ell_{N}(\psi)=\sum_{j=1}^{J}w_{N,j}\,\nabla\mathrm{log}\,p(y_{1:N}\mid \phi_{N,j};\psi), \tag{5}\]

which only requires differentiating the likelihood.

### Choosing the Markov kernel

Recall the definition of the Markov kernel \(h_{n}\): It is invariant with respect to the posterior distribution \(p(\phi\mid y_{1:n-1};\psi)\). Since the energy \(p(y_{1:n-1}\mid\phi;\psi)\,\pi(\phi)\) is usually analytically available, it is common to use any MCMC chain for \(h_{n}\)(Del Moral et al., 2006). In particular, when the problem dimension \(d\) is low, using a standard random walk Metropolis-Rosenbluth-Teller-Hasting MCMC suffices in practice. If the dimension is relatively high, we can also leverage the gradient information of the energy to define the kernel with, for instance, a Langevin dynamic. However, it is worth remarking that choosing the number of MCMC steps is tricky, and this still remains an open challenge in the SMC community, see, for example, the discussion in Chopin and Papaspiliopoulos (2020, Sec. 17.2) or a recent development in Dau and Chopin (2021).

The computational problem of using MCMC for the kernels is that at each \(n\), we need to load all the data before \(n\) to evaluate the energy function. For gradient-based MCMCs, we also need to compute \(\nabla_{\phi}\,\mathrm{log}\,p(\phi\mid y_{1:n-1};\psi)\). This computation becomes even more demanding when \(n\) approaches \(N\) (which is large), and hence the computational cost grows at least quadratically in \(n\). A straightforward remedy to this problem is to use pseudo-marginal MCMC samplers (see, e.g., Andrieu and Roberts, 2009; Welling and Teh, 2011; Chen et al., 2014; Zhang et al., 2020). However, these mini-batching MCMC samplers usually take long mixing steps and are hard to calibrate, which contradicts the gist of using MCMC kernels in SMC: An ideal SMC sampler should need as few mixing steps as possible. Furthermore, the variance of the estimator increases in \(n\), and it is likely that we need to adaptively increase the batch size to control the error. However, it turns out that we can rectify this problem by moving the stochastic mini-batching approximation outside of the Markov kernel to that of the gradient step as explained in the next section.

### Stochastic gradient sequential Monte Carlo

To make Algorithm 1 scalable in the number of data points \(N\), it is natural to approximate the marginal log-likelihood by a stochastic approximation, as in stochastic gradient descent. Let \(1\leq M\leq N\) denote the batch size and let \(S_{M}:=\{S_{M}(1),S_{M}(2),\ldots,S_{M}(M)\}\) be a sequence of independent random integers (uniformly distributed in \([1,N]\)) that represent the batch indices. We may approximate the marginal log-likelihood by \(\mathrm{log}\,p(y_{1:N};\psi)\approx N/\,M\,\mathrm{log}\,p(y_{S_{M}};\psi)\), where \(y_{S_{M}}:=\{y_{S_{M}(1)},y_{S_{M}(2)},\ldots,y_{S_{M}(M)}\}\) represents the corresponding subdataset. The same also goes for the gradient computation. More specifically,

\[\begin{split}&\nabla\mathrm{log}\,p(y_{1:N};\psi)\\ &\approx\frac{N}{M}\,\mathbb{E}\big{[}\nabla\mathrm{log}\,p(y_{S_{M }};\psi)\big{]}\\ &=\frac{N}{M}\int\mathbb{E}\big{[}\nabla\mathrm{log}\,p(\phi,y_{ S_{M}};\psi)\,p(\phi\mid y_{S_{M}};\psi)\big{]}\,\mathrm{d}\phi,\end{split} \tag{6}\]

where the expectation is taken over \(S_{M}\). However, due to the latent variable \(\phi\), the approximation \(N/\,M\,\mathrm{log}\,p(y_{S_{M}};\psi)\) is biased. Consequently, the gradient is also biased, and the bias scales in the difference
between \(p(\phi\mid y_{S_{M}})\) and \(p(\phi\mid y_{1:N})\).

Using the approximation in Equation (6) amounts to running Algorithm 1 on the subdataset \(y_{S_{M}}\) and then using the posterior samples of \(p(\phi\mid y_{S_{M}};\psi)\) to compute the gradient \(\int\nabla\!\log(\phi,y_{S_{M}};\psi)\,\mathrm{d}\phi\). We then arrive at the stochastic gradient SMC sampler summarised in Algorithm 2.

```
Inputs: The same as in Algorithm 1, and batch size \(M\) Outputs: The MLE estimate \(\psi_{i}\)
1for\(i=1,2,\ldots\) untilconvergentdo
2 Draw \(\{\phi_{0,j}\}_{j=1}^{J}\sim\pi(\phi)\)
3\(w_{0,j}=1\,/\,J\) for all \(j=1,2,\ldots,J\)
4 Draw subdataset \(y_{S_{M}}\subseteq y_{1:N}\)
5for\(n=1\)to\(M\)do// Parallelise \(j\)
6 Resample \(\{(w_{n-1,j},\phi_{n-1,j})\}_{j=1}^{J}\) if needed
7 Draw \(\phi_{n,j}\mid\phi_{n-1,j}\sim h_{n}(\phi_{n,j}\mid\phi_{n,j-1})\)
8\(\overline{w}_{n,j}=w_{n-1,j}\,p(y_{S_{M}}(n)\mid\phi_{n,j};\psi_{i-1})\)
9\(w_{n,j}=\overline{w}_{n,j}\,/\,\sum_{k=1}^{J}\overline{w}_{n,k}\)
10 end for
11\(g(\psi_{i-1})=\frac{N}{M}\sum_{j=1}^{J}w_{M,j}\,\nabla\!\log p(y_{S_{M}}\mid \phi_{M,j};\psi)\)
12\(\psi_{i}=\psi_{i-1}+r(i)\,g(\psi_{i-1})\)
13 end for
```

**Algorithm 2**Stochastic gradient SMC sampler for pBNN

Compared to Algorithm 1, the stochastic gradient version in Algorithm 2 does not have to load the entire dataset to compute the gradient. More importantly, this also facilitates the Markov kernel design dilemma in Section 3.2 in the following way. The kernel \(h_{n}\) is now chosen to be invariant with respect to \(p(\phi\mid y_{S_{M}(1)},\ldots,y_{S_{M}(n-1)})\) which is far easier to compute than the original when \(M\ll N\). Essentially, the algorithm is a direct application of the stochastic gradient method on a latent variable model by using SMC samplers to approximate the gradient.

### Open-horizon sequential Monte Carlo

Algorithm 2 defines an approximate flow of gradient that can be computed efficiently by applying SMC samplers in a closed data horizon. However, the algorithm does not directly output the target posterior distribution \(p(\phi\mid y_{1:N};\psi)\) unlike Algorithm 1. To compute the posterior distribution, we may need to run another Bayesian inference (e.g., by SMC or MCMC) based upon the estimated \(\psi\). Moreover, for every optimisation step, the SMC estimators are independent and cold-start from the prior \(\pi\). This causes a waste of computation, since Algorithm 2 does compute the posterior distributions \(p(\phi\mid y_{S_{M}};\psi_{i})\) at subdatasets which are approximations to the target posterior distribution. In light of this observation, we can make Algorithm 2 even more efficient by linking the posterior distribution estimates in conjunction with the gradient. More specifically, we modify Algorithm 2 by warm-starting each SMC sampler from the previous posterior distribution estimate and perform the gradient update in conjunction with the SMC sampler. We then arrive at an SMC sampler that simultaneously estimates the posterior distribution and the parameter, summarised in Algorithm 3.

```
Inputs: Same as in Algorithm 2 Outputs: Same as in Algorithm 1
1 Draw \(\{\phi_{0,j}\}_{j=1}^{J}\sim\pi(\phi)\)
2\(w_{0,j}=1\,/\,J\) for all \(j=1,2,\ldots,J\)
3for\(i=1,2,\ldots\) untilconvergentdo// Parallelise \(j\)
4 Draw subdataset \(y_{S_{M}^{i}}\subseteq y_{1:N}\)
5 Resample \(\{(w_{i-1,j},\phi_{i-1,j})\}_{j=1}^{J}\) if needed
6 Draw \(\phi_{i,j}\mid\phi_{i-1,j}\sim\hat{h}_{i}(\phi_{i,j}\mid\phi_{i,j-1})\)
7\(\overline{w}_{i,j}=w_{i-1,j}\,p(y_{S_{M}^{i}}\mid\phi_{i,j};\psi_{i-1})\)
8\(w_{i,j}=\overline{w}_{i,j}\,/\,\sum_{s=1}^{J}\overline{w}_{i,s}\)
9\(g(\psi_{i-1})=\frac{N}{M}\sum_{j=1}^{J}w_{i,j}\nabla\!\log p(y_{S_{M}^{i}} \mid\phi_{i,j};\psi_{i-1})\)
10\(\psi_{i}=\psi_{i-1}+r(i)\,g(\psi_{i-1})\)
11 end for
```

**Algorithm 3**Open-horizon SMC sampler

In Algorithm 3, we start from samples drawn from the given prior \(\pi\). Then, at each iteration \(i\), we randomly draw a subdataset \(y_{S_{M}^{i}}\) of \(y_{1:N}\) and compute the posterior distribution and gradient on the subdataset by using the previous estimate. Equivalently, the algorithm is an SMC sampler applied on a growing dataset with open horizon instead of \(y_{1:N}\) which has a fixed size (cf. Kantas et al., 2015, Eq. 5.6). Compared to Algorithms 1 and 2, this open-horizon SMC sampler is computationally more efficient, especially when \(N\) is large.

However, Algorithm 3 no longer targets the posterior distribution \(p(\phi\mid y_{1:N};\psi)\). To see what the algorithm does, let us fix the parameter \(\psi\), to find that the algorithm computes the following Feynman-Kac model

\[\begin{split}\frac{1}{\hat{l}_{P}(\psi)}\int\prod_{i=1}^{P}& p(y_{S_{M}^{i}}\mid\phi_{i};\psi)\prod_{i=1}^{P}\hat{h}_{i}(\phi_{i}\mid \phi_{i-1})\\ &\times\pi(\phi_{0})\,\mathrm{d}\phi_{0:P-1},\end{split} \tag{7}\]

where \(P\) is the number of iterations for the algorithm, and \(\hat{l}_{P}\) is the normalising constant. Evidently, Equation (7) is neither equal to the original Feynman-Kac model in Equation (4), nor to its expectation. It is, however, possible to use Poisson estimators (Beskos et al., 2006; Fearnhead et al., 2008; Jacob and Thiery,
2015) to make Equation (7) an unbiased estimator of the original Feynman-Kac model. To do so, we need to let the number of iterations \(P\sim\text{Poisson}(\lambda)\) be a Poisson random variable and modify the weight update in Algorithm 3 to

\[\overline{w}_{i,j}=w_{i-1,j}\,\Big{(}\frac{N}{M}\log p(y_{S_{M}^{i}}\mid\phi_{i, j};\psi_{i-1})+\log c\Big{)},\]

where \(c\) is a constant that guarantees the positivity of the weights almost surely. Then, we can show by the Poisson estimator that

\[\mathbb{E}\bigg{[}\frac{\mathrm{e}^{\lambda}}{c}\prod_{i=1}^{P} \frac{1}{\lambda}\,\Big{(}\frac{N}{M}\log p(y_{S_{M}^{i}}\mid\phi;\psi)+\log c \Big{)}\bigg{]}\] \[=\prod_{i=1}^{N}p(y_{i}\mid\phi;\psi),\]

where the expectation is taken on both \(P\) and \(S_{M}\). However, finding such a positive constant \(c\) is still an open but stimulating challenge in the statistic community (see, e.g., a recent progress in Jin et al., 2022).

For training pBNNs, it may not be essential to debias Algorithm 3. In practice, we use BNNs/pBNNs as powerful predictive models to quantify uncertainty. Hence, computing the exact posterior distribution is often unnecessary from a practical point of view (Wilson and Izmailov, 2020; Izmailov et al., 2021). As a consequence, the prior is no longer fixed but is a flexible model free to choose so as to aim for better predictive performance. To this end, we can further relax the invariance property of the Markov kernel \(\hat{h}\). This relaxation bridges a connection to the classical stochastic filtering approaches for optimisations (Bell, 1994; Gerber and Douc, 2021). As an example, the extended Kalman filter approximations to the algorithm are stochastic natural gradient descents with information matrices determined by the Markov kernel (Olivier, 2018; Martens, 2020). As another example, if we choose the Markov kernel to be that of a Brownian motion, then the resulting algorithm is akin to using stochastic filters for training neural networks with uncertainties (Singhal and Wu, 1988; De Freitas et al., 2000; Chang et al., 2022).

It is worth remarking a minor practical advantage of Algorithm 3 over Algorithm 2. The Markov kernel in Algorithm 2 is a function that has dynamic input sizes. This means that implementing the algorithm in commonly used automatic differentiation libraries which assume static input shapes (e.g., JAX and Tensorflow) is not trivial. On the other hand, Algorithm 3 has no such an issue.

## 4 Experiments

In this section, we evaluate our proposed Algorithms 2 and 3 which we call SGSMC and OHSMC, respectively, in several ways. First, we test the methods on a synthetic model to see if they can recover the true parameters and posterior distribution. Then, we train pBNNs using the methods for regression and classification tasks on synthetic, UCI, and MNIST datasets. Our implementation is publicly available at [https://github.com/zgbkdlm/pbnn](https://github.com/zgbkdlm/pbnn).

BaselinesWe compare against (i) the maximum a posteriori (MAP) method for estimating the parameter and use Hamiltonian Monte Carlo (MAP-HMC) to compute the posterior distribution based on the learnt parameter (Sharma et al., 2023, Sec. 6), (ii) stochastic weight averaging Gaussian (SWAG, Maddox et al., 2019) with the MAP objective function, and (iii) stochastic mean-field Gaussian variational Bayes (VB, Hoffman et al., 2013). In addition, SGSMC-HMC refers to sampling the posterior distribution by HMC and estimating the parameters by SGSMC.

SettingsWe use Adam for all methods with a learning rate of 0.01 unless stated otherwise. As for the prior, we consistently employ a standard Gaussian distribution. Whenever sampling from the posterior distribution is needed, we use \(J=1,000\) samples. We use the same amount for evaluations. For SGSMC, we use an MCMC random walk kernel, whereas for OHSMC, we use a random walk kernel (with variance 0.01). VB uses 100 MC samples to approximate the evidence lower bound. Further details (e.g., batch sizes, number of epochs, and pBNN structures) for all the experiments are found in the supplementary materials.

### Synthetic parameter estimation

Consider a model

\[\begin{bmatrix}\phi_{0}\\ \phi_{1}\end{bmatrix} \sim\mathrm{N}\bigg{(}\begin{bmatrix}0\\ 0\end{bmatrix},\begin{bmatrix}2&0\\ 0&1\end{bmatrix}\bigg{)}, \tag{8}\] \[y_{n}\mid\phi \sim\mathrm{N}\Big{(}\frac{\phi_{1}}{\psi}+\frac{1}{2}\,(\phi_{0} ^{2}+\psi^{2}),1\Big{)},\]

where we generate a sequence of 100 independent data points \(y_{1:100}\) under the true parameter \(\psi=1\) and the realisation \(\phi=\begin{bmatrix}0&0\end{bmatrix}\). The goal is to estimate the parameter \(\psi\) and also to compute the posterior distribution which has a non-Gaussian crescent shape. For this model, we can conveniently compute a tight lower bound of the MLE objective by brute-force Monte Carlo with 10,000 samples, which we call MC. However, for training pBNNs, the MC approach is in general not applicable.
In Figure 1, we see that the MC estimate is closest to the truth followed by OHSMC, VB, and SGSMC. The MAP estimator on the other hand, significantly diverges from the truth. Moreover, from Figure 2 we see that the OHSMC samples are close to the true density, although the samples do not explore the tails of the true distribution. The VB estimate is correct in the mean approximation, but the Gaussian approximation does not fit the shape of the true distribution. The SWAG method, which uses the results based on the MAP estimate, produces a covariance matrix whose diagonal is numerically zero.

### Synthetic regression and classification

RegressionNext, we benchmark the algorithms on a synthetic regression problem where we have access to the true underlying function. The training, validation, and test data are generated as per \(y_{n}=f(x_{n})+\xi_{n}\), for \(n=1,2,\ldots,100\), where \(f(x)\coloneqq x\sin(x\tanh(x))\), and i.i.d. noise \(\xi_{n}\sim\mathrm{N}(0,1)\). We use a three-layer pBNN (with output sizes 20, 10, and 1), where the stochastic part is on the second hidden layer. We repeat the experiment 100 times and report the averaged negative log predictive density (NLPD) on the test data, and root mean-square error (RMSE) on the true function including their standard deviations.

Table 1 shows that the proposed OHSMC method performs best in both metrics. Furthermore, Figure 3 shows that the predictive samples drawn from the pBNNs learnt by OHSMC best capture the local optima of the true function. Moreover, OHSMC extrapolates the regression problem best.

ClassificationAs for classification, we test the methods on a synthetic two-moon dataset. Apart from NLPD, we also report the expected calibration error (ECE, Guo et al., 2017) and accuracy. The utilised neural network has four dense layers (with output sizes 100, 20, 5, and 1), where the stochastic part is the third layer.

Table 2 shows that our OHSMC either performs better than its peers or on par in terms of NLPD. According to the ECE and accuracy metrics, OHSMC is best.

### UCI regression and classification

We now move from synthetic experiments to real-world UCI data (Kelly et al., Accessed 2023). Throughout, we use a neural network with four dense layers (with output sizes 50, 20, 5, and \(C\), where \(C\) is the number of labels), and place the stochastic part on the third layer. All experiments are repeated ten times and we report averaged metrics.

Table 3 shows the results on two regression and two classification UCI datasets. As before, our proposed OHSMC method either outperforms or performs on par with other methods in terms of NLPD for all datasets. As for regression, the RMSE is significantly lower for OHSMC. Moreover, our methods are marginally better than the others when evaluating ECE for the classification tasks. Results on additional datasets can be found in the supplementary materials.

### MNIST classification

We also consider classification on MNIST data (LeCun et al., Accessed 2023) by using a neural network with two convolutional layers followed by two dense layers. The stochastic part is on the first convolution layer. Due to their high computational cost, HMC methods do not apply here. Table 4 shows that our OHSMC

\begin{table}
\begin{tabular}{l c c c} \hline \hline Method & NLPD (std.) & ECE (std.) & Acc. (std.) \\ \hline MAP-HMC & **0.28** (0.06) & 0.07 (0.01) & 0.87 (0.02) \\ OHSMC & **0.28** (0.07) & **0.06** (0.01) & **0.88** (0.02) \\ SGSMC-HMC & 0.32 (0.08) & 0.08 (0.01) & 0.86 (0.03) \\ SWAG & 0.31 (0.06) & 0.07 (0.01) & 0.86 (0.03) \\ VB & 0.29 (0.05) & 0.08 (0.01) & 0.86 (0.03) \\ \hline \hline \end{tabular}
\end{table}
Table 2: Classification results in two-moon data.
method significantly outperforms other methods according to all the three metrics.

In this MNIST task, OHSMC (with 100 samples), SWAG, and VB (with 100 samples) take around 0.80, 0.26, and 0.33 seconds, respectively, for running 1,000 iterations with batch size 20. Note that the time for OHSMC includes _both_, the parameter estimation and posterior sampling, while for SWAG the time is only for parameter estimation. The times are profiled on an NVIDIA A100 40GB GPU.

## 5 Conclusions

In this paper, we have shown a sequential Monte Carlo (SMC) routine to efficiently train partial Bayesian neural networks (pBNNs). Specifically, we have proposed two approximate SMC samplers that are suitable for estimating the parameters and posterior distributions of pBNNs. The proposed training scheme either outperforms, or performs on par with state-of-the-art methods on a variety of datasets.

Limitations and future workThe proposed training scheme is immediately parallelisable on GPUs but at the cost of a high memory usage. An interesting future work is to theoretically analyse the convergence

\begin{table}
\begin{tabular}{c c c c} \hline \hline Method & NLPD (std.) & ECE (std.) & Acc. (std.) \\ \hline OHSMC & **2.87** (0.24) & **0.35** (0.06) & **99.02\%** (0.07) \\ SWAG & 4.27 (0.83) & 1.17 (0.43) & 97.76\% (0.71) \\ VB & 4.53 (0.28) & 0.44 (0.08) & 98.51\% (0.09) \\ \hline \hline \end{tabular}
\end{table}
Table 4: Classification results on MNIST. Within the table, NLPD and ECE are scaled by \(10^{-2}\).

Figure 4: Visualisation of the two-moons classifications. The scatter points are the test data with hollow/solid representing the label. The grey lines represent the classification hyperplanes sampled from the trained pBNNs.

\begin{table}
\begin{tabular}{r c c c c c c c c} \hline \hline \multirow{2}{*}{Method} & \multicolumn{2}{c}{yacht} & \multicolumn{2}{c}{energy} & \multicolumn{2}{c}{satellite} & \multicolumn{2}{c}{ionosphere} \\ \cline{2-9}  & NLPD & RMSE & NLPD & RMSE & NLPD & ECE & NLPD & ECE \\ \hline MAP-HMC & 0.95 (0.01) & 0.27 (0.04) & 0.94 (0.00) & 0.24 (0.02) & 0.28 (0.03) & 0.04 (0.00) & **0.21** (0.09) & 0.09 (0.03) \\ OHSMC & **0.92** (0.00) & **0.11** (0.04) & **0.92** (0.00) & **0.09** (0.00) & **0.27** (0.03) & **0.03** (0.00) & **0.21** (0.13) & **0.08** (0.03) \\ SGSMC-HMC & 0.93 (0.00) & 0.17 (0.02) & **0.92** (0.00) & 0.14 (0.02) & 0.84 (0.70) & 0.10 (0.10) & 0.24 (0.19) & **0.08** (0.04) \\ SWAG & 0.95 (0.02) & 0.25 (0.08) & 0.95 (0.00) & 0.25 (0.02) & 0.35 (0.07) & 0.07 (0.03) & 0.36 (0.12) & 0.14 (0.03) \\ VB & 0.94 (0.00) & 0.20 (0.02) & 0.93 (0.00) & 0.19 (0.03) & 0.28 (0.04) & **0.03** (0.00) & 0.28 (0.16) & 0.10 (0.04) \\ \hline \hline \end{tabular}
\end{table}
Table 3: Results on two regression and two classification UCI datasets. Results on more UCI datasets are provided in the supplementary material.

Figure 3: Visualisation of the synthetic regression. The scatter points represent the test data and the dashed line depicts the true function. The grey lines are predictive samples from their learnt pBNNs.
of Algorithm 3 as a general method for inference in latent variable models, since the algorithm carries a joint flow of both gradient and distribution.

## Acknowledgements

The authors would like to thank Adrien Corenflos for his comments. This work was partially supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg Foundation, and by the Kjell och Marta Beijer Foundation. The computations/data handling were enabled by the supercomputing resource Berzelius provided by National Supercomputer Centre at Linkoping University and the Knut and Alice Wallenberg foundation.
