Affine Invariant Ensemble Transform Methods to Improve Predictive
  Uncertainty in Neural Networks
We consider the problem of performing Bayesian inference for logistic
regression using appropriate extensions of the ensemble Kalman filter. Two
interacting particle systems are proposed that sample from an approximate
posterior and prove quantitative convergence rates of these interacting
particle systems to their mean-field limit as the number of particles tends to
infinity. Furthermore, we apply these techniques and examine their
effectiveness as methods of Bayesian approximation for quantifying predictive
uncertainty in neural networks.
Diksha Bhandari, Jakiw Pidstrigach, Sebastian Reich
# Affine Invariant Ensemble Transform Methods to Improve Predictive Uncertainty in ReLU Networks

# Affine Invariant Ensemble Transform Methods to Improve Predictive Uncertainty in ReLU Networks

Diksha Bhandari, Jakiw Pidstrigach, Sebastian Reich

**Abstract** We consider the problem of performing Bayesian inference for logistic regression using appropriate extensions of the ensemble Kalman filter. Two interacting particle systems are proposed that sample from an approximate posterior and prove quantitative convergence rates of these interacting particle systems to their mean-field limit as the number of particles tends to infinity. Furthermore, we apply these techniques and examine their effectiveness as methods of Bayesian approximation for quantifying predictive uncertainty in ReLU networks.

## 1 Introduction

The task in inverse problems is the inference of an unknown parameter \(\theta\in\mathbb{R}^{D}\) from noisy observations \(d\in\mathbb{R}^{N}\), which are generated through

\[d=G(\theta)+\eta, \tag{1}\]

where \(G\) denotes a forward map from model parameters to observable output data \(d\) and \(\eta\) denotes observational noise, which is commonly assumed to be Gaussian, that is, \(\eta\sim\mathcal{N}(0,P_{\eta})\). A Bayesian inverse problem can be formulated as producing samples from a random variable \(\theta\) conditioned on \(d\). Given the prior \(\pi_{\text{prior}}(\theta)\) on \(\theta\), the inverse problem is formulated as finding the posterior \(\pi_{\text{post}}(\theta)\) on \(\theta\) given \(d\). By Bayes theorem, the posterior distribution can be written in terms of the prior density \(\pi_{\text{prior}}\) and the negative log-likelihood or loss function \(\Psi:\mathbb{R}^{D}\to\mathbb{R}\) as

\[\mathrm{d}\pi_{\text{post}}(\theta)\propto\exp(-\Psi(\theta))\mathrm{d}\pi_{ \text{prior}}(\theta). \tag{2}\]

The problem of sampling from a target distribution is fundamental to Bayesian statistics, machine learning, and data assimilation. Considerable research has been done on sampling methods for Bayesian inverse problems, mainly in the case of \(l_{2}\)-loss functions [8, 21, 23, 24]. In this work, we will focus on \(\Psi\) being the cross-entropy loss instead, which is used in logistic regression problems.

We will introduce two methods, closely related to the algorithms introduced in [34] to approximate the posterior distribution for Bayesian logistic regression. We will prove that the methods are well-defined and provide quantitative bounds on their convergence in the large-ensemble (mean-field) limit.

We will then apply these methods to Bayesian neural networks. This gives us the possibility to quantify uncertainty in the model weights and the network outputs. Our numerical experiments show that our methods outperform the state-of-the-art methods for this problem.
### Classification and logistic regression

In this work, we focus on classification problems, that is, the problems arising from classifying objects into classes \(C_{1},\ldots,C_{k}\). For notational convenience, we will focus on the case of binary classification, i.e., we only have classes \(C_{1}\) and \(C_{2}\). However, all methods can be generalized to \(k\) classes, see [34, Section 7.2].

Consider the data set

\[\mathcal{D}=\{(\phi^{n},d^{n})\}_{n=1}^{N},\]

where \(d^{n}\in\{0,1\}\) are targets for the binary classification case with input features \(\phi^{n}\in\mathbb{R}^{D}\) for \(n=1,...,N\). The \(\phi^{n}\)'s can either be the representation of the data in any numerical form, or stem from a feature map. In Section 5, we will train a neural network for the use as feature map. We also introduce the shorthand notation

\[\Phi=(\phi^{1},...,\phi^{N})\in\mathbb{R}^{D\times N} \tag{3}\]

and define a parametric family of models as follows. Given a parameter \(\theta\), the probability of an example belonging to class \(C_{1}\) will be given by

\[\mathbb{P}_{\theta}[\phi\in C_{1}]=\sigma(\langle\theta,\phi\rangle),\]

where \(\sigma\) is the sigmoid function

\[\sigma(z)=\frac{1}{1+\exp{(-z)}}. \tag{4}\]

The negative log-likelihood of the given dataset \(\mathcal{D}\) under our parametric model is given by the cross-entropy loss function

\[\Psi(\theta)=-\sum_{n=1}^{N}\{d^{n}\log(y_{n}(\theta))+(1-d^{n})\log(1-y_{n}( \theta))\}, \tag{5}\]

where

\[y_{n}(\theta)=\sigma(\langle\theta,\phi^{n}\rangle)=\mathbb{P}_{\theta}[\phi^ {n}\in C_{1}]. \tag{6}\]

For \(n=1,\ldots,N\), we introduce the vector \(y(\theta)\in\mathbb{R}^{N}\) as

\[y(\theta)=(y_{1}(\theta),\ldots,y_{n}(\theta))^{\mathrm{T}}, \tag{7}\]

and the vector of target data labels \(d\in\mathbb{R}^{N}\) as

\[d=(d^{1},\ldots,d^{n})^{\mathrm{T}}. \tag{8}\]

We investigate sampling methods based on the ensemble Kalman filter (EnKF) [11] and its extension to Bayesian logistic regression as already considered in [34]. As discussed in the previous section, the Bayesian inverse problem now consists of sampling from \(\pi_{\mathrm{post}}\) given by (2).

Note that the likelihood in this section stems from linear logistic regression, since \(y_{n}=\sigma(\langle\theta,\phi^{n}\rangle)\). However, the methods can be generalized to nonlinear logistic regression, i.e., \(y_{n}=\sigma(f^{n}(\theta))\), with \(f^{n}(\theta)\) being nonlinear, see [34, Section 7.1]. Furthermore, the method proposed in Section 2.2 can be implemented in a derivative-free manner, as discussed in [34, Section 7.1].
### Literature review

Since its introduction in [10], the EnKF has been a popular choice for performing data assimilation tasks due to its robustness and wide applicability. In particular, the EnKF has shown promise as a derivative-free Bayesian inference technique [11, 34, 23]. More recently, the EnKF has been combined with sampling methods for Bayesian inference [8, 9] in order to transform random samples at \(s=0\) into samples from the posterior as \(s\to\infty\). Studying the theoretical aspects of homotopy methods based on the EnKF has also been an active area of research [38, 39, 34, 35, 4]. Furthermore, it should also be noted that both the EnKF and ensemble Kalman inversion (EKI) [24] can be cast within an interactive particle setting. However, most of the work is done in the case of a quadratic likelihood \(\Psi(\theta)\) or a perturbation of a quadratic likelihood. The work [24, 34] introduces multiple EnKF-type methods to study Bayesian logistic regression, i.e., \(\Psi\) being the negative log-likelihood of a logistic regression model. In this paper, we further develop two of the methods proposed in [34] by taking inspiration from [17, 23, 15, 37] and deploy them for uncertainty quantification in Bayesian logistic regression.

As neural networks have become widely used in critical applications, the need for accurate uncertainty quantification has grown [13]. To address this, a popular technique is using Bayesian inference [36], which allows machine learning algorithms to assign low confidence to test points that deviate significantly from the training data or prior information. Bayesian neural networks are commonly used for this purpose [43, 1, 30, 3, 44, 14, 33, 29]. Popular approximate Bayesian approaches include traditional variational inference [3, 19], Laplace approximation [28], Sampling-based approaches like Hamiltonian Monte Carlo (HMC) [32] and stochastic gradient Langevin dynamics [42]. More recently, in an attempt to make these Bayesian approaches more computationally feasible, many methods of partial stochastic networks have been considered [40, 25, 6, 41]. Moreover, many non-Bayesian methods [26, 22, 27] for uncertainty quantification have also been proposed. These methods aim to provide usable estimates of predictive uncertainty even in the presence of limited data.

### Outline and our contribution

The fundamental idea underlying this study is to develop efficient Bayesian inference methods for logistic regression based on the EnKF. Furthermore, we apply these methods to uncertainty quantification in ReLU networks and compare them to the state-of-the-art. To that end we will derive two interacting particle systems (IPSs) which sample from an approximate posterior. Moreover, we prove quantitative convergence rates of the IPSs to their mean-field-limit as the number of particles tends to infinity. We demonstrate the efficacy of the proposed methods for estimating uncertainty in ReLU networks through numerical experiments for binary classification in Section 6.

The remainder of this paper has been organized as follows. The dynamical system formulations for ensemble transform methods based on the homotopy approach and the infinite time second-order dynamical sampler for Bayesian inference in logistic regression are described in Section 2. Section 3 analyzes the proposed ensemble transform methods and derives their mean-field limits. Therein, we quantitatively bound the distance of the \(J\)-particle system to its infinite-particle limit in Theorem 1. Efficient methods of time-stepping of interacting particle systems for robust implementations and pseudocode describing the algorithms that we employ in this paper are discussed in Section 4. Section 5 introduces the framework for application of proposed methods to Bayesian logistic regression for uncertainty quantification in ReLU networks. Experimental results on predictive uncertainty of ReLU networks using proposed Bayesian methods for inference are shown in Section 6. These experiments demonstrate how reliable the uncertainty estimates are for different methods considered.
## 2 Dynamical system formulations

We denote by \(\theta_{s}^{1},\ldots\theta_{s}^{J}\) an ensemble of \(J\) particles, indexed by a time parameter \(s\geq 0\). We denote by \(m_{\theta_{s}}\) the empirical mean of the ensemble,

\[m_{\theta_{s}}=\frac{1}{J}\sum_{j=1}^{J}\theta_{s}^{j}, \tag{9}\]

and by \(P_{\theta_{s}}\) the empirical covariance matrix,

\[P_{\theta_{s}}=\frac{1}{J}\sum_{j=1}^{J}(\theta_{s}^{j}-m_{s})(\theta_{s}^{j}- m_{s})^{\mathrm{T}}. \tag{10}\]

We introduce the matrix of ensemble deviations \(\Theta_{s}\in\mathbb{R}^{D\times J}\) as

\[\Theta_{s}=\left(\theta_{s}^{1}-m_{\theta_{s}},\theta_{s}^{2}-m_{\theta_{s}}, \ldots,\theta_{s}^{J}-m_{\theta_{s}}\right). \tag{11}\]

We will adapt the convention that upper indices stand for the ensemble index, while lower indices stand for the time index or the \(i\)th entry of an \(N\)-dimensional vector. Associated to the particle system is the empirical measure

\[\mu_{\theta_{s}}=\frac{1}{J}\sum_{j=1}^{J}\delta_{\theta_{s}^{j}}, \tag{12}\]

where \(\delta_{\theta}\) stands for the Dirac delta at \(\theta\). We denote the expectation with respect to this measure by \(\mu_{\theta_{s}}[f]\), i.e. for a function \(f\),

\[\mu_{\theta_{s}}[f]=\frac{1}{J}\sum_{j=1}^{J}f(\theta_{s}^{j}).\]

Finally, we assume that the prior measure \(\pi_{\mathrm{prior}}\) is a Gaussian and given as

\[\pi_{\mathrm{prior}}=\mathcal{N}(m_{\mathrm{prior}},P_{\mathrm{prior}}).\]

In each of the following two subsections, we introduce an IPS to sample from the posterior in a logistic regression problem.

### Homotopy using moment matching

In homotopy approaches to Bayesian inference, we assume that the initial ensemble \(\{\theta_{0}^{j}\}\) is distributed according to the prior \(\pi_{\mathrm{prior}}\). One then evolves the ensemble such that at some fixed terminal time \(s\), often \(s=1\), the ensemble is distributed according to the posterior \(\pi_{\mathrm{post}}\). To derive such an evolution, one starts by defining a path \(\pi_{s}\) in the space of measures, which starts at a prior distribution \(\pi_{0}=\pi_{\mathrm{prior}}\) and ends at \(\pi_{1}=\pi_{\mathrm{post}}\)[5, 35].

We will study a homotopy approach introduced in [34]. In this section, we will shortly summarize the resulting differential equations. See [34, Section 4.2] for more details and the derivation of the equations.

We will need the gradient and Hessian of the negative log-likelihood \(\Psi\), defined in (5). The gradient is given by

\[\nabla_{\theta}\Psi(\theta)=\sum_{n=1}^{N}(y_{n}(\theta)-d^{n})\phi^{n}=\Phi(y (\theta)-d), \tag{13}\]
while the Hessian is

\[D_{\theta}^{2}\Psi(\theta)=\Phi R(\theta)\Phi^{\mathrm{T}}. \tag{14}\]

Here \(R(\theta)\in\mathbb{R}^{N\times N}\) is a diagonal matrix with diagonal entries

\[r_{nn}=y_{n}(\theta)(1-y_{n}(\theta)). \tag{15}\]

For \(s\in[0,1]\), the dynamical system to transform the prior is given by

\[\begin{split}\frac{\mathrm{d}}{\mathrm{d}s}\theta_{s}^{j}& =-\frac{1}{2}P_{\theta_{s}}\left(\mu_{\theta_{s}}[D_{\theta}^{2} \Psi(\theta)](\theta_{s}^{j}-m_{\theta_{s}})+2\mu_{\theta_{s}}[\nabla_{\theta }\Psi(\theta)]\right)\\ &=-\frac{1}{2}P_{\theta_{s}}\Phi\Big{(}\mu_{\theta_{s}}[R]\Phi^{ \mathrm{T}}(\theta_{s}^{j}-m_{\theta_{s}})+2(\mu_{\theta_{s}}[y]-d)\Big{)} \end{split} \tag{16}\]

with random initial conditions \(\theta_{0}^{j}\sim\mathcal{N}(m_{\mathrm{prior}},P_{\mathrm{prior}})\) i.i.d. for \(j=1,\ldots,J\).

One can also evolve the mean \(m_{\theta_{s}}\) and ensemble deviations \(\Theta_{s}\) instead of the ensemble \(\{\theta_{s}^{j}\}_{j=1}^{J}\). Their evolution equations are given by

\[\frac{\mathrm{d}}{\mathrm{d}s}m_{\theta_{s}}=-P_{\theta_{s}}\Phi\Big{(}\mu_{ \theta_{s}}[y]-d\Big{)}, \tag{17}\]

and

\[\frac{\mathrm{d}}{\mathrm{d}s}\Theta_{s}=-\frac{1}{2}P_{\theta_{s}}\Phi\mu_{ \theta_{s}}[R]\Phi^{\mathrm{T}}\Theta_{s} \tag{18}\]

respectively, where \(P_{\theta_{s}}\) is the empirical covariance matrix given in (10). However, the full ensemble is still required to compute the empirical expectation value of \(R(\theta)\) and \(y(\theta)\).

### Deterministic second-order dynamical sampler

Alternatively to transporting from the prior to the posterior in a fixed time interval, one can also construct systems that sample the target distribution as \(s\to\infty\). Markov Chain Monte Carlo algorithms are the most famous family of algorithms with this property. However, they normally work on a single sample trajectory \(\theta_{s}\) instead of an ensemble.

The algorithm introduced in [34, Section 5] combines the homotopy approaches with overdamped Langevin dynamics to motivate an IPS that approximates the posterior as \(s\to\infty\). The system of equations is given by

\[\begin{split}\frac{\mathrm{d}\theta_{s}^{j}}{\mathrm{d}s}& =-\frac{1}{2}P_{\theta_{s}}\Bigg{(}\Phi\Big{(}\mu_{\theta_{s}}[R] \Phi^{\mathrm{T}}(\theta_{s}^{j}-m_{\theta_{s}})+2(\mu_{\theta_{s}}[y]-d) \Big{)}\\ &\qquad+P_{\mathrm{prior}}^{-1}(\theta_{s}^{j}+m_{\theta_{s}}-2m _{\mathrm{prior}})\Bigg{)}+P_{\theta_{s}}^{1/2}\mathrm{d}W_{s}^{j},\end{split} \tag{19}\]

where \(W_{s}^{j}\) denotes \(D\)-dimensional standard Brownian motion. For details on the derivation see [34, Section 5]. Note, that similar systems are also introduced in [17, 23, 18].

We now modify (19) by replacing the stochastic driving term \(P_{\theta_{s}}^{1/2}\mathrm{d}W_{s}^{j}\) by \(\frac{1}{2}(\theta_{s}^{j}-m_{\theta_{s}})\mathrm{d}s\), rendering the system deterministic except for the choice of the initial ensemble \(\{\theta_{0}^{j}\}_{j=1}^{J}\). The advantage of making it deterministic is that we can perfectly assess convergence of the algorithm, since the particles will stop moving when they reach equilibrium. The replacement of \(P_{\theta_{s}}^{1/2}\) by \(\frac{1}{2}(\theta_{s}-m_{s})\) is motivated by the fact, that in the mean-field case, for Gaussian densities, they have the same
distributional effect, which we prove in Proposition 6 (for more details, see Section D). Furthermore, we will prove that the Gaussian assumption is well-founded in Section 3.

Therefore, the IPS ODE we will study is given by

\[\begin{split}\frac{\mathrm{d}\theta_{s}^{j}}{\mathrm{d}s}& =-\frac{1}{2}P_{\theta_{s}}\Bigg{(}\Phi\Big{(}\mu_{\theta_{s}}[R] \Phi^{\mathrm{T}}(\theta_{s}^{j}-m_{\theta_{s}})+2(\mu_{\theta_{s}}[y]-d) \Big{)}\\ &\qquad+P_{\mathrm{prior}}^{-1}(\theta_{s}^{j}+m_{\theta_{s}}-2m _{\mathrm{prior}})\Bigg{)}+\frac{1}{2}(\theta_{s}^{j}-m_{\theta_{s}})\end{split} \tag{20}\]

with random initial conditions \(\theta_{0}^{j}\sim\mathcal{N}(m_{0},P_{0})\) i.i.d. for \(j=1,\ldots,J\).

As already utilized in Section 2.1, to propagate the ensemble \(\{\theta_{s}^{j}\}_{j=1}^{J}\) in the interacting particle system (20), we can equivalently evolve the mean \(m_{\theta_{s}}\) and ensemble deviations \(\Theta_{s}\) by

\[\frac{\mathrm{d}}{\mathrm{d}s}m_{\theta_{s}}=-P_{\theta_{s}}\left(\Phi(\mu_{ \theta_{s}}[y]-d)+P_{\mathrm{prior}}^{-1}(m_{\theta_{s}}-m_{\mathrm{prior}}) \right), \tag{21}\]

and

\[\frac{\mathrm{d}}{\mathrm{d}s}\Theta_{s}=-\frac{1}{2}P_{\theta_{s}}\left(\Phi \mu_{\theta_{s}}[R]\Phi^{\mathrm{T}}\Theta_{s}+P_{\mathrm{prior}}^{-1}\Theta_ {s}\right)+\frac{1}{2}\Theta_{s} \tag{22}\]

respectively, where \(P_{\theta_{s}}\) is the empirical covariance. The method reaches an equilibrium solution as \(s\to\infty\), which gives the deterministic dynamical sampler a robustness that the homotopy methods might not possess.

## 3 Theoretical results on mean-field limits

In this section, we study what happens to the equations (16) and (20) when we let the number of particles go to infinity. As we will see, the interaction between the particles decreases, and in the limit \(J\to\infty\) all particles follow a deterministic mean-field ODE, independently of the other particles.

To that end, we introduce some notation. Both of the IPS in Section 2 only depend on the other particles through their empirical measure. Therefore, we can rewrite (16) and (20) as

\[\mathrm{d}\theta_{s}^{j}=b(\mu_{\theta_{s}})(\theta^{j}), \tag{23}\]

where \(b(\mu_{\theta})\) is defined implicitly by (16) and (20). By the continuity equation, we know that if the particles are evolved using the drift (23), then their empirical measure is a weak solution to the partial differential equation (PDE)

\[\partial_{s}\mu_{\theta_{s}}(\theta) = -\operatorname{div}(\mu_{\theta_{s}}b(\mu_{\theta_{s}})(\theta)), \qquad\mu_{\theta_{0}}=\frac{1}{J}\sum_{j=1}^{J}\delta_{\theta_{i}},\]

where \(\delta_{\theta}\) is the Dirac-delta distribution at \(\theta\). Due to the dependence of \(b\) on \(\mu_{\theta_{s}}\), this PDE is typically nonlinear. However, since \(b\) only depends on the other particles through the empirical measure, the above PDE forms a closed system. Therefore, we abstract the PDE

\[\partial_{s}\mu_{s}(\theta)=-\operatorname{div}(\mu_{s}b(\mu_{s})(\theta)) \tag{24}\]
which, given an initial condition \(\mu_{0}\), can be solved at the level of measures or densities directly. Given such a solution \((\mu_{s})_{s}\) for a fixed initial \(\mu_{0}\), we plug it into (23) and get the differential equation

\[\frac{\mathrm{d}}{\mathrm{d}t}\eta_{s}=b(\mu_{s})(\eta_{s}).\]

Note that this equation does not constitute an IPS anymore but a mean-field ODE instead; i.e., the particle evolution depends on its own distribution \(\mu_{s}\), which we obtain as the solution to (24). Furthermore, we find from (24) that \(\eta_{0}\sim\mu_{0}\) implies \(\eta_{s}\sim\mu_{s}\) for all \(s>0\).

The proofs in the following subsections work now as follows. We fix an initial condition

\[\mu_{0}=\mathcal{N}(m_{0},P_{0}),\]

for which we obtain the solution \(\mu_{s}\) to (24). Then, we define an intermediate mean-field particle system:

\[\frac{\mathrm{d}}{\mathrm{d}t}\eta_{s}^{j} = b(\mu_{s})(\eta_{s}^{j}),\quad\eta_{0}^{j}\sim\mu_{0}. \tag{25}\]

Note that the \(\eta_{s}^{j}\), \(j=1,\ldots,J\), are all independent. We then couple the IPS (16) or (20) to (25) by choosing the same initial conditions, i.e., \(\theta_{0}^{j}=\eta_{0}^{j}\) for all \(j=1,\ldots,J\). We then prove that the expected Wasserstein-distance of the empirical measures

\[\mathbb{E}[\mathcal{W}(\mu_{\theta_{s}},\mu_{\eta_{s}})]\to 0 \tag{26}\]

converges to \(0\) as \(J\to\infty\). Let us briefly discuss the expectation value in (26). Note that although (23) and (25) are deterministic, \(\mu_{\theta_{s}}\) as well as \(\mu_{\eta_{s}}\) are random probability measures. However, the randomness comes solely from the initial conditions. Therefore, the expectation in (26) is with respect to i.i.d. initial conditions \(\eta_{0}^{j}=\theta_{0}^{j}\sim\pi_{0}\).

We not only prove (26) but are able to obtain a quantitative convergence rate. Since \(\mu_{\eta_{s}}\) is the \(J\)-fold product of \(\mu_{s}\) with itself, this shows that for large \(J\) the \(\theta_{s}^{j}\) approximate independent samples from \(\mu_{s}\). One can make the last statement precise by using known rates at which \(\mu_{\eta_{s}}\) converges to \(\mu_{s}\), see [12].

### Analysing the mean-field systems

As already discussed, we will prove that the empirical measure \(\mu_{\theta_{s}}\) approximates the solution of the mean-field PDE \(\mu_{s}\) for large ensemble sizes. Therefore, it is instructive to briefly study \(\mu_{s}\). Since \(\mu_{0}\) is Gaussian in our case and the drift terms \(b(\mu_{\theta_{s}})(\theta)\) are linear in \(\theta\), \(\mu_{s}\) will remain Gaussian for all times \(s\geq 0\). We denote its mean and covariance by \(m_{s}\) and \(P_{s}\):

\[\mu_{s}=\mathcal{N}(m_{s},P_{s}).\]

Therefore, solving the mean-field PDE (24) corresponds to solving an ODE for the mean and covariance of the Gaussian distribution. We will next work out these mean-field ODEs for our two IPS.

#### 3.1.1 Homotopy using moment matching

The mean-field limit PDE, corresponding to (24), is given by

\[\partial_{s}\mu_{s}(\eta) = \frac{1}{2}\operatorname{div}\left(\mu_{s}\left(P_{s}(\Phi\mu_{s }[R]\Phi^{\mathrm{T}}(\eta-m_{s})+2\Phi(\mu_{s}[y]-d)\right)\right).\]
In the case of Gaussian initial conditions, the above PDE is equivalent to solving the following ODEs for the mean and covariance:

\[\begin{array}{rcl}\frac{\mathrm{d}}{\mathrm{d}s}m_{s}&=&-P_{s}\Phi(\mu_{s}[y]-d )=-P_{s}\mu_{s}[\nabla_{\theta}\Psi(\theta)],\\ \frac{\mathrm{d}}{\mathrm{d}s}P_{s}&=&-P_{s}\Phi\mu_{s}[R]\Phi^{\mathrm{T}}P_{ s}=-P_{s}\mu_{s}[D_{\theta}^{2}\Psi(\theta)]P_{s},\end{array} \tag{27}\]

where \(\nabla_{\theta}\Psi\) and \(D_{\theta}^{2}\Psi\) are the gradient and Hessian of the negative log-likelihood function as defined in (13) and (14), respectively.

#### 3.1.2 Deterministic second-order dynamical sampler

In this case, the mean-field limit PDE is given by

\[\begin{array}{rcl}\partial_{s}\mu_{s}(\eta)=&\frac{1}{2}\operatorname{div} \left(\mu_{s}\left(P_{s}(\Phi\mu_{s}[R]\Phi^{\mathrm{T}}(\eta-m_{s})+2\Phi( \mu_{s}[y]-d)\right)\right)\\ &\quad+\frac{1}{2}\operatorname{div}\left(\mu_{s}\left(P_{\mathrm{ prior}}^{-1}(\eta+m_{s}-2m_{\mathrm{prior}})\right)+(\eta-m_{s})\right) \right),\end{array} \tag{28}\]

where \(\mu_{0}=\mathcal{N}(m_{0},P_{0})\). As in the previous subsection, we again derive the mean-field ODEs for the mean and covariance:

\[\begin{array}{rcl}\frac{\mathrm{d}}{\mathrm{d}s}m_{s}&=&-P_{s}\Phi(\mu_{s}[y ]-d)-P_{s}P_{\mathrm{prior}}^{-1}(m_{s}-m_{\mathrm{prior}}),\\ \frac{\mathrm{d}}{\mathrm{d}s}P_{s}&=&-P_{s}\Phi\mu_{s}[R]\Phi^{\mathrm{T}}P_{ s}-P_{s}P_{\mathrm{prior}}^{-1}P_{s}+P_{s}.\end{array} \tag{29}\]

The mean-field ODE (25) becomes

\[\begin{array}{rcl}\frac{\mathrm{d}}{\mathrm{d}s}\eta_{s}^{j}&=&-\frac{1}{2} P_{s}\left(\Phi(\mu_{s}[R]\Phi^{\mathrm{T}}(\eta^{j}-m_{s}))+2\Phi(\mu_{s}[y]-d)\right) \\ &+&\frac{1}{2}P_{s}P_{\mathrm{prior}}^{-1}(\eta^{j}+m_{s}-2m_{\mathrm{prior}} ))+\frac{1}{2}(\eta^{j}-m_{s}),\end{array} \tag{30}\]

where \(\eta_{0}^{j}=\theta_{0}^{j}\sim\mathcal{N}(m_{0},P_{0})\). When the system (29) stops evolving, we have reached the equilibrium distribution

\[\mu_{*}=\mathcal{N}(m_{*},P_{*}).\]

To derive equations for \(m_{*}\) and \(P_{*}\), we set the right-hand sides of (29) to zero and obtain:

\[m_{*} =m_{\mathrm{prior}}-P_{\mathrm{prior}}\Phi(\mu_{*}[y]-d)=m_{ \mathrm{prior}}-P_{\mathrm{prior}}\mu_{*}[\nabla\Psi(\theta)], \tag{31}\] \[P_{*} =(\Phi\mu_{*}[R]\Phi^{\mathrm{T}}+P_{\mathrm{prior}}^{-1})^{-1}= (\mu_{*}[D_{\theta}^{2}\Psi(\theta)]+P_{\mathrm{prior}}^{-1})^{-1}. \tag{32}\]

These are implicit equations in \(m_{*}\) and \(P_{*}\) and the evolution equations (32) can be seen as means to find \(m_{*}\) and \(P_{*}\). Therefore, in the many-particle and large time limit, we are approximating a Gaussian with mean and covariance given by (31) and (32) respectively.

Approximating a distribution by a Gaussian is also an important topic in variational inference [16]. However, in contrast to popular methods for Gaussian variational inference (see the discussion in [16]), which are based on taking gradients with respect to the mean and covariance of the approximating Gaussian, we do not need to invert the state space (\(D\times D\)) covariance matrix for our method to work.

**Remark 1**: _The work [16] also proposes a deterministic IPS for Gaussian variational inference. While their evolution equations differ from ours, the equilibrium state agrees with (31)-(32). Furthermore, in contrast to our formulation, the IPS proposed in [16] is not affine-invariant. See [34] for a discussion of affine invariance._
### Statement of results

Since the IPS (16) and (20) are similar, differing only in additional terms for (20), the proofs of the following results are quite similar too. Due to (20) having more terms, in particular also terms that increase the ensemble spread, the proofs are more technical. We concentrate on that case. The analogous results for (16) follow by performing very similar, often nearly identical, calculations, but for fewer terms.

First of all, we prove in the following proposition that the objects of interest, \(\theta_{s}^{j}\), \(\mu_{\theta_{s}}\), \(\eta_{s}^{j}\), and \(\mu_{s}\) are well-defined.

**Proposition 1**: _The mean-field PDE (28) has a unique global solution \(\mu_{s}\). Furthermore, the IPS (20) and the mean-field IPS (30) also posses unique global solutions._

The proposition is proven in Section C.1, in Proposition 3 and Proposition 4. We are now in a position to state our main theorem.

**Theorem 1**: _Let \(\{\theta_{s}^{j}\}_{j=1}^{J}\) be the solution to (20) with associated empirical measure \(\mu_{\theta_{s}}\) and let \(\mu_{s}\) be the solution to (28). For any \(\varepsilon\), there exists a constant \(C\), depending only on \(m_{0},P_{0},\Phi,d,s\) and \(\varepsilon\), such that_

\[\mathbb{E}[\mathcal{W}_{2}(\mu_{\theta_{s}},\mu_{s}^{N})]\leqslant C_{ \varepsilon,s}J^{-\frac{1}{2}+\varepsilon},\]

_where \(\mu_{s}^{N}\) is the \(N\)-fold product measure of \(\mu_{s}\) with itself._

**Proof** (Sketch) We introduce an artificial mean-field particle system \(\eta^{j}\) as described in Section 3. The precise mean-field ODEs can be found in (30). We couple the \(\theta_{s}^{j}\) to the \(\eta_{s}^{j}\) by choosing the same initial conditions, i.e., \(\eta_{0}^{j}=\theta_{0}^{j}\). The \(\eta_{s}^{j}\) are samples from \(\mu_{s}^{N}\), and therefore the Wasserstein distance, being the infimum over all couplings, can be upper bounded by this specific coupling between \(\theta_{s}^{j}\) and \(\eta_{s}^{j}\), i.e.,

\[\mathbb{E}[\mathcal{W}_{2}(\mu_{\theta_{s}},(\mu_{s})^{N})]\leqslant\mathbb{ E}\left[\left(\frac{1}{J}\sum_{j=1}^{J}|\theta_{s}^{j}-\eta_{s}^{j}|^{2}\right)^{1 /2}\right]. \tag{33}\]

We fix a \(T\) and assume that \(s\leqslant T\). To bound (33), we define

\[\Delta_{s}=\left(\frac{1}{J}\sum_{j=1}^{J}|\theta_{s}^{j}-\eta_{s}^{j}|^{2} \right)^{1/2}\]

and upper bound its growth. In Proposition 2 we will show that we can bound any moment of \(\mathbb{E}[\Delta_{s}]\) independently of the ensemble size \(J\), i.e.,

\[\mathbb{E}\left[|\Delta_{s}|^{p}\right]^{1/p}\lesssim C_{p}, \tag{34}\]

where \(x\lesssim y\) symbolizes that the inequality \(x\leqslant ay\) holds with a constant \(a\) depending only on \(T,m_{0},P_{0},\Phi\) and \(d\).

We then use a bootstrapping technique inspired by [8]. The main idea is to show that, if

\[\mathbb{E}[\Delta_{s}]\lesssim J^{-\gamma} \tag{35}\]

for some \(\gamma\geq 0\), then one can actually improve that \(\gamma\) value to a better \(\gamma^{\prime}\), i.e.,

\[\mathbb{E}[\Delta_{s}]\lesssim J^{-\gamma^{\prime}} \tag{36}\]
holds for some \(\gamma^{\prime}>\gamma\). Plugging \(p=0\) into (34), we see that (35) holds for \(\gamma=0\). We then iteratively improve \(\gamma=0\) to any \(\gamma=\frac{1}{2}-\epsilon\).

We now go into a bit more detail on how to improve the current estimate of \(\gamma\). We denote by \(H_{\alpha}\) a random variable, which can change from occurrence to occurrence, such that \(\mathbb{E}[H_{\alpha}^{p}]^{1/p}\leqslant C_{p}J^{-\alpha}\) for all \(p\geq 0\), where \(C_{p}\) only depends on \(T,m_{0},P_{0},\Phi,d\) and \(p\). Considering the ODE for \(\Delta_{s}\) and bounding all terms of its right-hand side, we end up with

\[\frac{\mathrm{d}}{\mathrm{d}s}\mathbb{E}[\Delta_{s}] \leqslant C(\mathbb{E}[\Delta_{s}]+\mathbb{E}[\Delta_{s}H_{1/4}]+\mathbb{ E}[H_{1/2}])\] \[\leqslant C(\mathbb{E}[\Delta_{s}]+\mathbb{E}[\Delta_{s}^{\varepsilon} \Delta_{s}^{1-\varepsilon}H_{1/4}]+\mathbb{E}[H_{1/2}])\] \[\leqslant C(\mathbb{E}[\Delta_{s}]+\mathbb{E}[\Delta_{s}]^{1-\varepsilon }\mathbb{E}[\Delta_{s}(H_{1/4})^{1/\varepsilon}]^{\varepsilon}+\mathbb{E}[H_ {1/2}])\] \[\leqslant C(\mathbb{E}[\Delta_{s}]+\mathbb{E}[\Delta_{s}]^{1-\varepsilon }\mathbb{E}[\Delta_{s}^{2}]^{\varepsilon/2}\mathbb{E}[(H_{1/4})^{2/\varepsilon }]^{\varepsilon/2}+\mathbb{E}[H_{1/2}])).\]

Here we repeatedly used Holder's inequality. Now, all moments of \(H_{\alpha}\) can be bounded by \(J^{-\alpha}\)

\[\frac{\mathrm{d}}{\mathrm{d}s}\mathbb{E}[\Delta_{s}] \lesssim \mathbb{E}[\Delta_{s}]+\mathbb{E}[\Delta_{s}]^{1-\varepsilon}J^{ -1/4}+J^{-1/2}\] \[\leqslant \mathbb{E}[\Delta_{s}]+J^{-\gamma(1-\varepsilon)-1/4}+J^{-1/2}.\]

In the second inequality we used the a priori assumption that \(\mathbb{E}[\Delta_{s}]\leqslant J^{-\gamma}\). We now apply Groenwall to obtain

\[\mathbb{E}[\Delta_{s}]\lesssim J^{-\gamma^{\prime}},\]

where \(\gamma^{\prime}=\min\left(\frac{1}{2},\frac{1}{4}+\gamma(1-\varepsilon)\right)\). Plugging in \(\gamma=0\), we obtain a rate of \(J^{-1/4}\) by applying the above argument once. Then, by picking a small enough \(\varepsilon\), we can achieve any rate below \(1/2\) by applying the bootstrapping technique a second time. The full details of the proof can be found in Appendix C.3. \(\Box\)

The following a priori bound is crucial for the proof of Theorem 1.

**Proposition 2**: _For \(s\in[0,T]\),_

\[\mathbb{E}\left[\Big{|}\frac{1}{J}\sum|\theta_{s}^{j}-\eta_{s}^{j}|^{2}|^{p} \right]^{1/p}\leqslant C_{p},\]

_i.e., the \(p\)-norm can be bounded independently of \(J\) for a fixed \(T\). The constant \(C_{p}\) only depends on \(m_{0},P_{0},\Phi,d\) and \(T\)._

The proof relies on the fact that \(\sum_{j=1}^{J}|\theta_{0}^{j}-\eta_{0}^{j}|\) is a martingale and uses martingale inequalities. It can be found in Section C.2.

## 4 Algorithmic implementation

A typical way of time-stepping the interacting particle systems presented in this paper is the forward-Euler method. However, due to its restricted domain of stability, using this method can lead to restrictions on the step-size \(\Delta s\). In this section we describe tamed discretizations for the homotopy based moment matching formulation (16) and the deterministic second-order dynamical sampler (20). We introduce a step-size \(\Delta s\geq 0\) and discrete times \(s_{k}=k\Delta s\). Further, we use the shorthand \(\theta_{s_{k}}\approx\theta_{k},m_{\theta_{s_{k}}}\approx m_{k},P_{\theta_{s_ {k}}}\approx P_{k},\Theta_{s_{k}}\approx\Theta_{k},\mu_{\theta_{s_{k}}} \approx\mu_{k}\) in the forthcoming subsections.
### Homotopy using moment matching

We employ modifications to the time stepping for moment-matching method (16) by using the following tamed discretizations

\[\theta_{k+1}^{j}=\theta_{k}^{j}-\frac{\Delta s}{2}P_{k}\Phi\left(M_{k}\Phi^{ \mathrm{T}}(\theta_{k}^{j}-m_{k})+2(\mu_{k}[y]-d)\right) \tag{37}\]

where

\[P_{k}=\frac{1}{J}\Theta_{k}\Theta_{k}^{\mathrm{T}}\]

and

\[M_{k}=\left(\Delta s\Phi^{\mathrm{T}}P_{k}\Phi+\mu_{k}[R]\right)^{-1} \tag{38}\]

for \(j=1,\ldots,J\). As discussed in Section 2.1, we propagate \(\theta_{k}^{j}\) forward by evolving the associated empirical mean and ensemble deviations using (17)-(18). The resulting time-stepping of (17)-(18) is of the form

\[m_{k+1}=m_{k}-\Delta sP_{k}\Phi\left(\mu_{k}[y]-d\right) \tag{39}\]

and

\[\Theta_{k+1}=\Theta_{k}-\frac{\Delta s}{2}P_{k}\Phi M_{k}\Phi^{\mathrm{T}} \Theta_{k} \tag{40}\]

for the ensemble mean and ensemble deviations, respectively.

**Remark 2**: _Inverting the \(N\times N\) matrix (38) can prove prohibitive for large data sets. However, since \(R(\theta)\) is diagonal, taking the full inverse in (38) can be replaced by inverting the diagonal entries of \(\Delta s\Phi^{\mathrm{T}}P_{k}\Phi+\mu_{k}[R]\) only, as proposed in [2]. This inexpensive approximation still provides improved stability compared to an explicit Euler discretization of (18)._

We provide pseudo-code summarizing the second-order moment matching method described in Algorithm 1.

```
Inputs: Data set \(\{(\phi^{n},d^{n})\}_{n=1}^{N}\); feature map \(\Phi=\{\phi^{n}\}_{n=1}^{N}\); initial ensemble \(\{\theta_{0}^{j}\}_{j=1}^{J}\) drawn from a Gaussian distribution; step-size \(\Delta s\) and \(K\) such that \(\Delta sK=1\). for\(k=0\)to\(K-1\)do: 1 Evaluate ensemble mean \(m_{k}\) (9), ensemble deviations \(\Theta_{k}\) (11), covariance matrix \(P_{k}\) (10), \(y\) (6), \(\mu_{k}[y]\), \(\mu_{k}[R]\) and \(M_{k}\) (38). 2 Evolve \(m_{k}\) and \(\Theta_{k}\) using (39)- (40). 3 Determine \(\{\theta_{k}^{j}\}_{j=1}^{J}\) from \(m_{k+1}\) and \(\Theta_{k+1}\). endfor Output: Final ensemble \(\{\theta_{K}^{j}\}_{j=1}^{J}\).
```

**Algorithm 1**Homotopy using moment matching method for Bayesian inference
### Deterministic second-order dynamical sampler

We employ the idea of Trotter splitting to solve (20) numerically. The required splitting is provided by the evolution equation (16), already encountered in the homotopy approach, and the remainder

\[\frac{\mathrm{d}\theta_{s}^{j}}{\mathrm{d}s}=-\frac{1}{2}P_{\theta_{s}}P_{\text{ prior}}^{-1}(\theta_{s}^{j}+m_{\theta_{s}}-2m_{\text{prior}})+\frac{1}{2}( \theta_{s}^{j}-m_{\theta_{s}}) \tag{41}\]

Therefore, for every time step \(k\), given \(\{\theta_{k}^{j}\}_{j=1}^{J}\), we first compute \(\theta_{k+1/2}^{j}\) using the second order moment matching method (37) rewritten as

\[\theta_{k+1/2}^{j}=\theta_{k}^{j}-\frac{\Delta s}{2}P_{k}\Phi\left(M_{k}\Phi^{ \mathrm{T}}(\theta_{k}^{j}-m_{k})+2(\mu_{k}[y]-d)\right) \tag{42}\]

for \(j=1,2,...,J\). Equivalently, we can obtain \(\theta_{k+1/2}^{j}\) by evaluating \(m_{k+1/2}\) and \(\Theta_{k+1/2}\) as stated in (39)-(40) with subscript \(k+1\) replaced by \(k+1/2\). In the second half step, we approximate (41) using the following scheme:

\[\begin{split}\theta_{k+1}^{j}&=\theta_{k+1/2}^{j}- \frac{\Delta s}{2}P_{k+1/2}\left(\Delta sP_{k+1/2}+P_{\text{prior}}\right)^{-1} \left(\theta_{k+1/2}^{j}+m_{k+1/2}-2m_{\text{prior}}\right)\\ &\quad+\frac{\Delta s}{2}(\theta_{k+1/2}^{j}-m_{k+1/2})\end{split} \tag{43}\]

Again, if \(P_{\text{prior}}\) is diagonal, full matrix inversion in (43) can be replaced by inverting the diagonal only.

We provide pseudo-code describing the algorithm for deterministic second-order dynamical sampler in Algorithm 2.

```
Inputs: Data set \(\{(\phi^{n},d^{n})\}_{n=1}^{N}\); feature map \(\Phi=\{\phi^{n}\}_{n=1}^{N}\); initial ensemble \(\{\theta_{0}^{j}\}_{j=1}^{J}\) drawn from a Gaussian distribution; step-size \(\Delta s\); threshold value \(\epsilon>0\). for\(k\geq 0\)do: 1 Numerically solve ODE (20) by Trotter splitting: 1. Evaluate ensemble mean \(m_{k}\) (9), ensemble deviations \(\Theta_{k}\) (11), covariance matrix \(P_{k}\) (10), \(y\) (6), \(\mu_{k}[y]\), \(\mu_{k}[R]\) and \(M_{k}\) (38) 2. Determine \(\{\theta_{k+1/2}^{j}\}_{j=1}^{J}\) by evolving \(m_{k}\) and \(\Theta_{k}\) using the time-stepping (39)-(40). 3. Evaluate covariance matrix \(P_{k+1/2}\) (10). 4. Determine \(\{\theta_{k+1}^{j}\}_{j=1}^{J}\) using the time-stepping (43). 2. Update the ensemble \(\{\theta_{k}^{j}\}_{j=1}^{J}\rightarrow\{\theta_{k+1}^{j}\}_{j=1}^{J}\). if\(\frac{\|P_{k+1}-P_{k}\|_{2}}{\|P_{k}\|_{2}}<\epsilon\); break endfor Output: Final ensemble \(\{\theta_{k}^{j}\}_{j=1}^{J}\).
```

**Algorithm 2**Deterministic second-order dynamical sampler for Bayesian inference
Application to Bayesian logistic regression in ReLU networks

Neural networks using ReLU activation functions are possibly the most widely used neural network architectures for classification. However, it has been proven that these networks exhibit arbitrarily high confidence far away from the training data when fitted by minimizing the negative log-likelihood \(\psi\) or, equivalently, by approximating the maximum likelihood estimator (MLE), denoted by \(\tilde{\theta}_{\text{MLE}}\), as demonstrated in [22, 20]. Thus, this architecture along with a MLE training scheme is not robust and does not provide any measure of uncertainty in the model's predictions.

One way of obtaining predictive uncertainty is to place distributions over the weights of a ReLU network, which leads to Bayesian neural networks. The idea of replacing the MLE \(\tilde{\theta}_{\text{MLE}}\) by a posterior measure \(\tilde{\pi}_{\text{post}}\) over parameters \(\tilde{\theta}\) therefore enables us to make better informed predictions and to know when our model predictions are not to be trusted. To understand how uncertainty might be expressed in this setting, we put a prior \(\tilde{\pi}_{\text{prior}}\) on the parameters \(\tilde{\theta}\) which can be thought of as incorporating some prior knowledge and then refining it based on data to learn a posterior distribution. That is, after observing \(\mathcal{D}\), we can get the posterior measure through Bayes formula (2).

This Bayesian approach to neural networks introduces a certain degree of computational complexity as we now need to sample from \(\tilde{\pi}_{\text{post}}\) instead of minimizing \(\Psi\) which can also be computationally expensive. Computational approximations to \(\tilde{\pi}_{\text{post}}\) have been enabled by advances in Markov chain Monte Carlo (MCMC) methods (see [31]). However, even today's most sophisticated MCMC methods are rendered impractical for deep neural network architectures. Therefore, to make the Bayesian approach tractable, we focus on a last layer Bayesian approximation that places a prior distribution only on the output layer's parameters. Thus, we decompose an \(l\)-layered ReLU network \(G_{\tilde{\theta}}:\mathbb{R}^{n}\rightarrow\mathbb{R}\) into a feature map \(\phi_{\tilde{\theta}}:\mathbb{R}^{n}\rightarrow\mathbb{R}^{d}\) consisting of the first \(l-1\) layers of the ReLU network and the output layer, i.e.,

\[G_{\tilde{\theta}}(x)=\sigma(\langle\theta,\phi_{\tilde{\theta}}(x)\rangle)\]

with \(\tilde{\theta}=(\hat{\theta}^{\text{T}},\theta^{\text{T}})^{\text{T}}\). One now first trains the complete network using a (regularised) MLE approach, which provides \(\tilde{\theta}_{\text{MLE}}\) and the associated trained feature map

\[\phi(x):=\phi_{\tilde{\theta}_{\text{MLE}}}(x). \tag{44}\]

Furthermore, upon defining the input features \(\phi^{n}=\phi(x^{n})\), \(n=1,\ldots,N\), over the data set \(\{(x^{n},d^{n})\}_{n=1}^{D}\), this architecture is now equivalent to a Bayesian logistic regression problem in the output layer parameters \(\theta\) as discussed in detail in Section 1.1.

In this paper, we analyze the performance of ReLU networks in the case of binary classification. While the work in [25] focuses on Laplace approximations for Bayesian inference, we employ algorithms for Bayesian logistic regression based on the methods proposed in Section 2. We demonstrate experimentally that our methods in conjunction with pre-trained deterministic ReLU networks provide desirable uncertainty estimates. However, it should be noted that the methods proposed in this paper are not limited to a 'last layer' use, but can be easily extended to multiple layers or the entire network.

We use a 3-layer ReLU network with 20 hidden units at each layer and \(D=50\) units at the output layer in the subsequent numerical experiments. The data set is constructed by generating a 2D binary classification data using scikit-learn. The MLE estimator \(\tilde{\theta}_{\text{MLE}}\) is obtained using the PyTorch library. Figure 1 depicts the above generated data set.
## 6 Numerical experiments

In this section, we consider a numerical experiment for binary classification problem as an illustration for uncertainty quantification in ReLU networks. We employ the already described 3-layer neural network architecture \(G_{\tilde{\theta}}(x)\). The training data set \(\{(x^{n},d^{n})\}_{n=1}^{N}\) consists of inputs and the associated labels. We train the ReLU network using stochastic gradient descent (SGD) with 0.9 momentum, learning rate of \(3\times 10^{-4}\) and weight-decay for training over the 2D binary classification data set 1, using \(N=30\) test points for toy binary classification problem. SGD minimizes the cross entropy loss quantity (5) across the training data set, and we denote the computed minimizer by \(\tilde{\theta}_{\text{MLE}}\).

The computed set of parameters (\(\hat{\theta}_{\text{MLE}}\)) is then used to provide the feature maps (44) which is used for Bayesian logistic regression. The chosen prior is Gaussian with mean \(m_{\text{prior}}=0\) and covariance matrix \(P_{\text{prior}}=2I\). Using the ensemble of particles distributed according to the posterior \(\tilde{\pi}_{\text{post}}\), approximated with the homotopy based moment-matching method (39)-(40) and the second-order dynamical sampler (43), the predictive distribution is estimated as

\[\pi(d=1|x,\mathcal{D})=\frac{1}{J}\sum_{j=1}^{J}\sigma(\langle\theta_{*}^{j}, \phi_{\hat{\theta}_{\text{MLE}}}(x)\rangle), \tag{45}\]

where \(\{\theta_{*}^{j}\}_{j=1}^{J}\) is the final ensemble of particles obtained using Algorithm 1 and Algorithm 2. The associated Bayesian posterior distribution now translates uncertainty in weights to uncertainty in model predictions.

Results for both methods introduced in this paper are reported using an ensemble size of \(J=200\). Algorithm 1 is implemented with step-size \(\Delta s=10^{-3}\) while Algorithm 2 is implemented with \(\Delta s=10^{-1}\) and \(\Delta sK=20\).

The results from Algorithms 1 & 2 are compared to the following two alternative methods.

Figure 1: 2D binary classification data set
**Laplace approximations.** We compare the uncertainty estimates provided by the proposed EnKF based methods for Bayesian approximation over the output layer parameters to the last-layer Laplace approximation (LLLA) for inference as introduced in [25]. In this case, we perform a Laplace approximation to get the posterior of the weights of the last layer, assuming the previous layers to be fixed at MAP estimates. So, for unknown output layer parameters \(\theta\), we infer

\[\pi(\theta|\mathcal{D})=\mathcal{N}(\theta|\theta_{\mathrm{MLE}},H^{-1}) \tag{46}\]

where \(H\) is the Hessian of the negative log-posterior with respect to \(\theta\) at \(\theta_{\mathrm{MLE}}\). The predictive distribution in the case of binary classification is thus given by

\[\pi(d=1|x,\mathcal{D})=\int\sigma(\langle\theta,\phi(x)\rangle\pi(\theta| \mathcal{D})\mathrm{d}\theta, \tag{47}\]

where \(\pi(\theta|\mathcal{D})\) is approximated with (46) and the integral (47) is computed using a prob approximation as described in [25, 30].

**Ensemble learning.** As a non-Bayesian method, we investigate the uncertainty estimates from ensemble learning (or deep ensembles), introduced in [26, 33]. The technique uses an ensemble of deterministic networks, meaning that each network in the ensemble produces a point estimate rather than a distribution. We train an ensemble of \(M=5\) ReLU networks independently, using the entire training data set to train each network. Given an input point \(x_{n}\), target label \(d_{n}\) and cross-entropy loss \(\Psi(\tilde{\theta},x_{n},d_{n})\), we generate an adversarial sample

\[x_{n}^{*}=x_{n}+\xi\ \mathrm{sign}(\nabla_{\theta}\Psi(\tilde{\theta},x_{n},d_{ n})),\]

where \(\xi\sim\mathcal{N}(0,0.01)\). As described in [26], using adversarial samples for training by adding perturbations in the direction the network is likely to increase the loss provides a 'better random' direction for smoothing predictive distributions. For \(M\) models with MLE estimated parameters \(\{\tilde{\theta}_{m}\}_{m=1}^{M}\), we evaluate the ensemble predictions as

\[\pi(d=1\mid x) =M^{-1}\sum_{m=1}^{M}\pi(d=1\mid x,\tilde{\theta}_{m}), \tag{48}\] \[=M^{-1}\sum_{m=1}^{M}G_{\tilde{\theta}_{m}}(x),\]

where \(\{\tilde{\theta}_{m}\}_{m=1}^{M}\) denotes the parameters of the \(m^{\mathrm{th}}\) model in the ensemble. For classification, this corresponds to averaging the predicted probabilities.

We obtain results for the model's confidence in its prediction on and away from any input data \(x\), where confidence is defined as the maximum predictive probability. In the case of a binary classification problem, confidence in prediction can be expressed as \(\max_{i\in\{0,1\}}p(d=i\mid x)\). Ideally, one would want the model to convey low confidence with such inputs. We next report on the results from these two experimental settings.

### Predictive uncertainty in ReLU networks

As shown in Figure 2 and Figure 3, MLE predictions have very high confidence everywhere besides the region close to the decision boundary. Using deep ensembles of ReLU networks improves the accuracy of the prediction, which results in a better decision boundary and shows increased uncertainty estimates near the decision boundary. However, both MLE and ensemble learning predictions do not
express uncertainty far away from the training data. On the other hand, last-layer Bayesian approximations, implemented with either Laplace, the homotopy moment matching method or second-order dynamical sampler, assign relatively high confidence close to the training data and are uncertain otherwise. In other words, the region of high confidence is identified much closer to the training data. All Bayesian approximations closely follow the decision boundary obtained using the MLE estimates and thus do not negatively impact the network's predictive accuracy. Furthermore, the second-order dynamical sampler assigns higher confidence closer to training data than any other algorithm and allows one to use a larger step size \(\Delta s\), and can therefore be recommended for further use. It can be noted that last-layer Laplace approximations also improve predictive uncertainty but assigns lower confidence than our methods on and near the training data set.
Figure 2: Binary classification on a toy dataset using (a) MAP estimates, (b) ensemble of neural networks, last-layer Gaussian approximations over the weights obtained via (c) Laplace approximation, (d) moment matching method, (e) deterministic second-order dynamical sampler. Background colour depicts the confidence in classification while black line represents the decision boundary obtained for the toy classification problem.
Figure 3: Zoomed-out versions of the results in Figure 2 for binary classification on a toy data set using (a) MAP estimates, (b) ensemble of neural networks, last-layer Gaussian approximations over the weights obtained via (c) Laplace approximation, (d) moment matching method, (e) deterministic second-order dynamical sampler. Background colour depicts the confidence in classification.
In Figure 3, we show a zoomed out version of the results in Figure 2 to capture confidence levels significantly away from training data. It can be seen that the MLE estimate and ensemble learning demonstrate high confidence in the entire domain. The Bayesian approximations, even when applied to just the last-layer of a ReLU network, give desirable uncertainty estimates.

### Uncertainty on out-of-distribution data

In this experiment, we generate test samples on a large grid, such that the test points deviate significantly from the training samples. For each point in the test data set, we evaluate its Euclidean distance \(\delta>0\) from the nearest training point. We then evaluate the model's confidence in classifying these out-of-distribution (OOD) samples as a function of \(\delta\). The results can be found as in Figure 4. It can be seen that as distance from the training set increases, the MLE technique is extremely overconfident in its predictions everywhere. MLE trained network always yields arbitrarily 'overconfident' predictions away from training data and thus, is not robust. Using an ensemble of ReLU networks also does not improve uncertainty estimates and has little affect on the confidence on OOD data. However, last layer-Bayesian approximations assign lower confidence for OOD test data. As the distance from the training set increases, the model is less confident in its prediction and the level of confidence converges to a constant. Furthermore, our approaches assign maximum confidence (higher than Laplace approximations) when \(\delta=0\), i.e., the in-distribution data.
Figure 4: Confidence of MAP, ensembles of neural networks, last-layer Laplace approximation, moment matching method, and deterministic second-order dynamical sampler as functions of \(\delta\) over the test set. Thick blue lines and shades correspond to means and \(\pm\) standard deviations, respectively. Dashed black lines signify the desirable confidence for \(\delta\) sufficiently high.
### Effect of varying ensemble sizes on predictive uncertainty

We also analyze the effect of varying ensemble sizes on inference of last-layer network parameters.

It can be observed that as ensemble size increases, the region of high confidence is identified a lot closer to the training data, while closely following the decision boundary obtained by the MLE estimator. For parameters \(D=50\), as ensemble size \(J\rightarrow\infty\) (here \(J=300\)), we observe that the interacting particle systems converge to their mean-field limit. It can also be concluded that a large ensemble size (\(J>D\)) provides better uncertainty estimates than small ensemble sizes (\(J\leq D\)). However, using ensemble size (\(J<D\)) still results in better uncertainty estimates near the decision boundary than those obtained by MLE estimates.

## 7 Conclusions

In this paper, we have presented two extensions of EnKF and related interacting particle systems to Bayesian logistic regression. We have proven quantitative convergence rates for these systems to their mean field limits as the number of particles tends to infinity. We have employed both methods for Bayesian inference in ReLU networks with cross-entropy loss function. The numerical results confirm the effectiveness of the proposed methods for quantifying uncertainty. They have also shown that these uncertainty estimates make ReLU networks more robust with respect to distributional shifts.

**Acknowledgements:** The research has been partially funded by the Deutsche Forschungsgemeinschaft (DFG)- Project-ID 318763901 - SFB1294. The authors would also like to thank the Isaac

Figure 5: Effect of varying ensemble sizes (\(J\)) on confidence in prediction for binary classification using proposed ensemble sampling methods for Bayesian inference over the network’s output (last) layer.
Newton Institute for Mathematical Sciences, Cambridge, for support and hospitality during the programme _The Mathematical and Statistical Foundation of Future Data-Driven Engineering_ where work on this paper was undertaken. This work was supported by EPSRC grant no EP/R014604/1.
