Input-gradient space particle inference for neural network ensembles
Deep Ensembles (DEs) demonstrate improved accuracy, calibration and
robustness to perturbations over single neural networks partly due to their
functional diversity. Particle-based variational inference (ParVI) methods
enhance diversity by formalizing a repulsion term based on a network similarity
kernel. However, weight-space repulsion is inefficient due to
over-parameterization, while direct function-space repulsion has been found to
produce little improvement over DEs. To sidestep these difficulties, we propose
First-order Repulsive Deep Ensemble (FoRDE), an ensemble learning method based
on ParVI, which performs repulsion in the space of first-order input gradients.
As input gradients uniquely characterize a function up to translation and are
much smaller in dimension than the weights, this method guarantees that
ensemble members are functionally different. Intuitively, diversifying the
input gradients encourages each network to learn different features, which is
expected to improve the robustness of an ensemble. Experiments on image
classification datasets and transfer learning tasks show that FoRDE
significantly outperforms the gold-standard DEs and other ensemble methods in
accuracy and calibration under covariate shift due to input perturbations.
Trung Trinh, Markus Heinonen, Luigi Acerbi, Samuel Kaski
# Input gradient diversity for neural network ensembles

###### Abstract

Deep Ensembles (DEs) demonstrate improved accuracy, calibration and robustness to perturbations over single neural networks partly due to their functional diversity. Particle-based variational inference (ParVI) methods enhance diversity by formalizing a repulsion term based on a network similarity kernel. However, weight-space repulsion is inefficient due to over-parameterization, while direct function-space repulsion has been found to produce little improvement over DEs. To sidestep these difficulties, we propose First-order Repulsive Deep Ensemble (FoRDE), an ensemble learning method based on ParVI, which performs repulsion in the space of first-order input gradients. As input gradients uniquely characterize a function up to translation and are much smaller in dimension than the weights, this method guarantees that ensemble members are functionally different. Intuitively, diversifying the input gradients encourages each network to learn different features, which is expected to improve the robustness of an ensemble. Experiments on image classification datasets show that FoRDE significantly outperforms the gold-standard DEs and other ensemble methods in accuracy and calibration under covariate shift due to input perturbations.

## 1 Introduction

Ensemble methods, which combine predictions from multiple models, are a well-known strategy in machine learning to boost performance (Dietterich, 2000). Ensembles of neural networks have been shown to achieve great results in terms of predictive performance (Lakshminarayanan et al., 2017), uncertainty estimation (Ovadia et al., 2019), robustness to adversarial attacks (Pang et al., 2019) and corruptions (Hendrycks and Dietterich, 2019). A common strategy to create a neural network ensemble is called Deep Ensembles (DEs) (Lakshminarayanan et al., 2017), which generate different networks by repeating training from different weight initializations. While DEs are simple to implement and have promising performance, they are susceptible to redundancy in model averaging due to their reliance on training randomness to induce network diversity (Rame and Cord, 2021; D'Angelo and Fortuin, 2021; Yashima et al., 2022).

Particle-based variational inference methods (ParVIs) (Liu and Wang, 2016; Chen et al., 2018; Liu et al., 2019; Shen et al., 2021) have recently emerged as a direction to improve diversity in neural ensembles from the Bayesian perspective (Wang et al., 2019; D'Angelo and Fortuin, 2021). Notably, the ParVI update rule adds a kernelized repulsion term between the ensemble networks for explicit control of the network diversity. Typically repulsion is done in the weight space to capture different regions in the weight posterior. However, due to the over-parameterization of neural networks, weight-space repulsion is inefficient and does not resolve the redundancy problem of DEs. An alternative approach is to define the repulsion in function space (Wang et al., 2019; D'Angelo and Fortuin, 2021), which requires a kernel between functions. Previous works avoided this issue by
comparing functions only on training inputs, which leads to underfitting (D'Angelo et al., 2021). Neither weight- nor function-space repulsion has led to significant improvements over vanilla DEs.

The gravitation towards weight and function repulsion stems from the fact that weight and function views have been the two main ways to characterize a parametric model. From a functional perspective, a model can also be uniquely represented, up to translation, using its first-order derivatives, i.e., input _gradients_\(\nabla_{\mathbf{x}}f\). Promoting diversity in this third view of input gradients has notable advantages:

1. each ensemble member is guaranteed to correspond to a different function;
2. input gradients often have dramatically smaller dimensions than network weights and thus are more amenable to kernel comparison;
3. the training procedure can directly diversify the input gradients of the networks with respect to the training samples without resulting in underfitting ensembles (see Fig. 1), unlike function-space repulsion on the training samples;
4. as input gradients quantify the salience of input features, each ensemble is encouraged to learn different features, which can improve robustness.

In this work, we propose a ParVI neural network ensemble that promotes diversity in their input gradients, denoted First-order Repulsive deep ensemble (FoRDE). In addition to the above advantages, we devise a data-dependent kernel that allows FoRDE to significantly outperform other ensemble methods under input corruptions on image classification tasks.

## 2 Background

Bayesian neural networksIn a Bayesian neural network (BNN), we treat the model's weights \(\theta\) as random variables with a prior \(p(\theta)\). Given a dataset \(\mathcal{D}=\{(\mathbf{x}_{n},y_{n})\}_{n=1}^{N}\) and a likelihood function \(p(y|\mathbf{x},\theta)\) per data point, we infer the posterior over weights \(p(\theta|\mathcal{D})\) using Bayes' rule

\[p(\theta|\mathcal{D})=\frac{p(\mathcal{D}|\theta)p(\theta)}{\int_{\theta}p( \mathcal{D}|\theta)p(\theta)\mathrm{d}\theta}=\frac{p(\theta)\prod_{n=1}^{N} p(y_{n}|\mathbf{x}_{n},\theta)}{\int_{\theta}p(\theta)\prod_{n=1}^{N}p(y_{n}| \mathbf{x}_{n},\theta)\mathrm{d}\theta}, \tag{1}\]

where the likelihood is assumed to factorize over data. To make a prediction on a test sample \(\mathbf{x}^{*}\), we integrate over the inferred posterior in Eq. (1), a practice called _Bayesian model averaging_ (BMA):

\[p(y|\mathbf{x}^{*},\mathcal{D})=\int_{\theta}p(y|\mathbf{x}^{*},\theta)p( \theta|\mathcal{D})\mathrm{d}\theta=\mathbb{E}_{p(\theta|\mathcal{D})}\big{[} p(y|\mathbf{x}^{*},\theta)\big{]}. \tag{2}\]

However, computing the integral in Eq. (2) is intractable for BNNs. Various approximate inference methods have been developed for BNNs, including variational inference (VI) (Graves, 2011; Blundell et al., 2015), Markov chain Monte Carlo (MCMC) (Neal, 2012; Welling and Teh, 2011; Zhang et al., 2020) and more recently ParVI (Liu and Wang, 2016; Wang et al., 2019; D'Angelo and Fortuin, 2021).

Figure 1: **An illustration of input gradient repulsion in 1D regression with 3 neural networks.****Left:** At some point during training, the models fit well to the training samples yet exhibit low functional diversity. **Middle**: As training proceeds, at each data point, the repulsion term gradually pushes the input gradients (represented by the arrows) away from each other on a unit sphere. **Right:** As a result, at the end of training, the ensemble has gained functional diversity.
Deep ensemblesAs opposed to BNNs, which attempt to learn the posterior distribution, DEs (Lakshminarayanan et al., 2017) consist of multiple maximum-a-posteriori (MAP) estimates trained independently from random initializations. They can capture diverse functions that explain the data well, as independent training runs under different random conditions will likely converge to different modes in the posterior landscape. DEs have been shown to be better than BNNs in both accuracy and uncertainty estimation (Ovadia et al., 2019; Ashukha et al., 2020; Gustafsson et al., 2020).

Particle-based variational inference for neural network ensemblesParVI methods (Liu and Wang, 2016; Chen et al., 2018; Liu et al., 2019; Shen et al., 2021) have been studied recently to formalize neural network ensembles. They approximate the target posterior using a set of samples, or particles, by deterministically transporting these particles from their initial positions to the target distribution (Liu and Wang, 2016). ParVI methods are expected to be more efficient than MCMC as they take into account the interactions between particles in their update rules (Liu et al., 2019). These repulsive interactions are driven by a kernel which measures the pairwise similarities between particles, ie. networks (Liu et al., 2019).

The main approaches for the kernel are comparing the networks in weight space \(\theta\) or in function space \(f(\cdot;\theta)\). Weight-space repulsion is ineffective due to difficulties in comparing extremely high-dimensional weight vectors and the existence of weight symmetries in the weight-space posterior (Fort et al., 2019; Enetezari et al., 2022). Unsurprisingly, previous studies showed that ParVI with weight-space repulsion did not improve performance over plain DEs (D'Angelo and Fortuin, 2021; Yashima et al., 2022). Comparing neural networks via a function-space kernel is also challenging since functions are infinite-dimensional objects. Previous works resort to comparing functions only on a subset of the input space called the measurement set (Wang et al., 2019; D'Angelo and Fortuin, 2021). Using the training data as the measurement set leads to underfitting (D'Angelo and Fortuin, 2021; Yashima et al., 2022), likely because these inputs have known labels thus leaving no room for diverse predictions without impairing performance.

## 3 FoRDE: First-order Repulsive Deep Ensembles

In this section, we present a framework to perform ParVI in the _input-gradient_ space. We start by introducing the update rule for particle optimization called Wasserstein gradient descent (WGD) in Section 3.1. We then describe how to apply WGD for input-gradient-space repulsion in Section 3.2. In Section 3.3, we discuss hyperparameter selection for the input-gradient kernel. Finally, in Section 3.4, we outline two practical considerations when training FoRDE.

Throughout this section, we assume a set of \(M\) weight particles \(\{\theta_{i}\}_{i=1}^{M}\) corresponding to a set of \(M\) neural networks \(\{f_{i}:\mathbf{x}\mapsto f(\mathbf{x};\theta_{i})\}_{i=1}^{M}\). We focus on the supervised classification setting: given a labelled dataset \(\mathcal{D}=\{(\mathbf{x}_{n},y_{n})\}_{n=1}^{N}\) with \(\mathcal{C}\) classes and inputs \(\mathbf{x}_{n}\in\mathbb{R}^{D}\), we approximate the posterior \(\log p(\theta|\mathcal{D})\) using the \(M\) particles. The output \(f(\mathbf{x};\theta)\) for input \(\mathbf{x}\) is a _log-softmax vector_ of size \(\mathcal{C}\) whose \(y\)-th entry \(f(\mathbf{x};\theta)_{y}\) equals the log-probability of the \(y\)-th class.

### Wasserstein gradient descent

Following D'Angelo and Fortuin (2021), we use a ParVI method called Wasserstein gradient descent (WGD) (Liu et al., 2019; Wang et al., 2022). Given a target posterior distribution \(\pi=p(\cdot|\mathcal{D})\) and a set of particles \(\{\theta_{i}\}_{i=1}^{M}\) from distribution \(\rho\), the goal of WGD is to find the particle distribution \(\rho^{*}\),

\[\rho^{*}=\operatorname*{argmin}_{\rho\in\mathcal{P}_{2}(\Theta)}\mathrm{KL}_{ \pi}(\rho), \tag{3}\]

where \(\mathrm{KL}_{\pi}(\rho)\) is a shorthand for the standard KL-divergence

\[\mathrm{KL}_{\pi}(\rho)=\mathbb{E}_{\rho(\theta)}\big{[}\log\rho(\theta)- \log\pi(\theta)\big{]}, \tag{4}\]

and \(\mathcal{P}_{2}(\Theta)\) is the Wasserstein space equipped with the Wasserstein distance \(W_{2}\)(Ambrosio et al., 2005; Villani, 2009). WGD solves the problem in Eq. (3) using a Wasserstein gradient flow \((\rho_{t})_{t}\), which is roughly the family of steepest descending curves of \(\mathrm{KL}_{\pi}(\cdot)\) in \(\mathcal{P}_{2}(\Theta)\). The tangent vector of this gradient flow at time \(t\) is

\[v_{t}(\theta)=\nabla\log\pi(\theta)-\nabla\log\rho_{t}(\theta) \tag{5}\]
whenever \(\rho_{t}\) is absolutely continuous (Villani, 2009; Liu et al., 2019). Intuitively, \(v_{t}(\theta)\) points to the direction where the probability mass at \(\theta\) of \(\rho_{t}\) should be transported in order to bring \(\rho_{t}\) closer to \(\pi\). Since Eq. (5) requires the analytical form of \(\rho_{t}\) which we do not have access to, we use kernel density estimation (KDE) to obtain a tractable approximation \(\hat{\rho}_{t}\) induced by particles \(\{\theta_{i}^{(t)}\}_{i=1}^{M}\),

\[\hat{\rho}_{t}(\theta)\propto\sum_{i=1}^{M}k\big{(}\theta,\theta_{i}^{(t)} \big{)}, \tag{6}\]

where \(k\) is a positive semi-definite (PSD) kernel. Then, the approximation gradient is

\[\nabla\log\hat{\rho}_{t}(\theta)=\frac{\sum_{i=1}^{M}\nabla_{\theta}k\big{(} \theta,\theta_{i}^{(t)}\big{)}}{\sum_{i=1}^{M}k\big{(}\theta,\theta_{i}^{(t)} \big{)}}. \tag{7}\]

Using Eq. (7) in Eq. (5), we obtain a practical update rule for each particle \(\theta^{(t)}\) of \(\hat{\rho}_{t}\):

\[\theta^{(t+1)}=\theta^{(t)}+\eta_{t}\Bigg{(}\underbrace{\nabla_{\theta^{(t)} }\log\pi\big{(}\theta^{(t)}\big{)}}_{\text{driving force}}-\underbrace{ \frac{\sum_{i=1}^{M}\nabla_{\theta^{(t)}}k\big{(}\theta^{(t)},\theta_{i}^{(t) }\big{)}}{\sum_{i=1}^{M}k\big{(}\theta^{(t)},\theta_{i}^{(t)}\big{)}}}_{\text {repulsion force}}\Bigg{)}, \tag{8}\]

where \(\eta_{t}>0\) is the step size at optimization time \(t\). Intuitively, we can interpret the first term in the particle gradient as the driving force directing the particles towards high density regions of the posterior, while the second term is the repulsion force pushing the particles away from each other.

### Defining the kernel for WGD in input gradient space

We propose to use a kernel comparing the _input gradients_ of the particles

\[k(\theta_{i},\theta_{j})\stackrel{{\text{def}}}{{=}}\mathbb{E}_ {(\mathbf{x},y)\sim p(\mathbf{x},y)}\Big{[}\kappa\big{(}\nabla_{\mathbf{x}}f( \mathbf{x};\theta_{i})_{y},\nabla_{\mathbf{x}}f(\mathbf{x};\theta_{j})_{y} \big{)}\Big{]}, \tag{9}\]

where \(\kappa\) is a _base kernel_ between gradients \(\nabla_{\mathbf{x}}f(\mathbf{x};\theta)_{y}\) that are of same size as the inputs \(\mathbf{x}\). In essence, we define \(k\) as the expected similarity between the input gradients of two networks with respect to the data distribution \(p(\mathbf{x},y)\). Interestingly, by using the kernel \(k\), the KDE approximation \(\hat{\rho}\) of the particle distribution not only depends on the particles themselves but also _depends on the data distribution_. We approximate the kernel \(k\) using the training samples, with linear complexity:

\[k(\theta_{i},\theta_{j})\approx k_{\mathcal{D}}(\theta_{i},\theta_{j})=\frac{ 1}{N}\sum_{n=1}^{N}\kappa\big{(}\nabla_{\mathbf{x}}f(\mathbf{x}_{n};\theta_{i })_{y_{n}},\nabla_{\mathbf{x}}f(\mathbf{x}_{n};\theta_{j})_{y_{n}}\big{)}. \tag{10}\]

The kernel compares the gradients of only the true label \(\nabla_{\mathbf{x}}f(\mathbf{x}_{n};\theta)_{y_{n}}\), as opposed to the entire Jacobian matrix \(\nabla_{\mathbf{x}}f(\mathbf{x}_{n};\theta)\), as our motivation is to encourage each particle to learn different features that could explain the training sample \((\mathbf{x}_{n},y_{n})\) well. This approach also reduces computational complexity, since automatic differentiation libraries such as JAX (Bradbury et al., 2018) or Pytorch (Paszke et al., 2019) would require \(\mathcal{C}\) passes, one per class, to calculate the full Jacobian.

Choosing the base kernelWe choose the RBF kernel on the unit sphere as our base kernel \(\kappa\):

\[\kappa(\mathbf{s},\mathbf{s}^{\prime};\boldsymbol{\Sigma})=\exp\left(-\frac{ 1}{2}(\mathbf{s}-\mathbf{s}^{\prime})^{\top}\boldsymbol{\Sigma}^{-1}(\mathbf{ s}-\mathbf{s}^{\prime})\right),\qquad\mathbf{s}=\frac{\nabla_{\mathbf{x}}f( \mathbf{x};\theta)_{y}}{||\nabla_{\mathbf{x}}f(\mathbf{x};\theta)_{y}||_{2}} \in\mathbb{R}^{D} \tag{11}\]

where \(\mathbf{s},\mathbf{s}^{\prime}\) denote the two normalized gradients of two particles with respect to one input, and \(\boldsymbol{\Sigma}\in\mathbb{R}^{D\times D}\) is a diagonal matrix containing squared lengthscales. We design \(\kappa\) to be norm-agnostic since the norm of the true label gradient \(||\nabla_{\mathbf{x}}f(\mathbf{x}_{n};\theta)_{y_{n}}||_{2}\) fluctuates during training and as training converges, the log-probability \(f(\mathbf{x}_{n};\theta)_{y_{n}}=\log p(y_{n}|\mathbf{x}_{n},\theta)\) will approach \(\log 1\), leading to the norm \(||\nabla_{\mathbf{x}}f(\mathbf{x}_{n};\theta)_{y_{n}}||_{2}\) approaching \(0\) due to the saturation of the log-softmax activation. Furthermore, comparing the normed input gradients between ensemble members teaches them to learn complementary explanatory patterns from the training samples, which could improve robustness of the ensemble. The RBF kernel is an apt kernel to compare unit vectors (Jayasumana et al., 2014), and we can control the variances of the gradients along input dimensions via the square lengthscales \(\boldsymbol{\Sigma}\).
### Selecting the lengthscales for the base kernel

In this section, we present a method to select the lengthscales for the base kernel. These lengthscales are important for the performance of FoRDE, since they control how much repulsion force is applied in each dimension of the input-gradient space. The dimension-wise repulsion is (Liu and Wang, 2016)

\[\frac{\partial}{\partial s_{d}}\kappa(\mathbf{s},\mathbf{s}^{\prime};\mathbf{ \Sigma})=-\frac{s_{d}-s^{\prime}_{d}}{\mathbf{\Sigma}_{dd}}\kappa(\mathbf{s}, \mathbf{s}^{\prime};\mathbf{\Sigma}), \tag{12}\]

where we can see that along the \(d\)-th dimension the inverse square lengthscale \(\mathbf{\Sigma}_{dd}\) controls the strength of the repulsion \(\nabla_{s_{d}}\kappa(\mathbf{s},\mathbf{s}^{\prime};\mathbf{\Sigma})\): a smaller lengthscale corresponds to a stronger force 1. Additionally, since the repulsion is restricted to the unit sphere in the input-gradient space, maximizing the distance of two particles in one dimension will minimize the distance in other dimensions. As a result, the repulsion motivates the ensemble members to depend more on dimensions with stronger repulsion in the input gradient space for their predictions, while focusing less on dimensions with weaker repulsion. From this observation, one should apply strong repulsion (small lengthscales), in the dimensions of the data manifold with high variances and vice versa.

Footnote 1: Here we assume that the lengthscales are set appropriately so that the kernel \(\kappa\) does not vanish, which is true since we use the median heuristic during training (Section 3.4).

To realize the intuition above in FoRDE, we first apply Principal Component Analysis (PCA) to discover important features in the training data. In PCA, we calculate the eigendecomposition

\[\mathbf{C}=\mathbf{U}\mathbf{\Lambda}\mathbf{U}^{T}\quad\in\mathbb{R}^{D\times D} \tag{13}\]

of the covariance matrix \(\mathbf{C}=\mathbf{X}^{T}\mathbf{X}/(N-1)\) of the centered training samples \(\mathbf{X}\in\mathbb{R}^{N\times D}\) to get the eigenvectors and eigenvalues \(\{\mathbf{u}_{d},\lambda_{d}\}_{k=1}^{D}\). The \(d\)-th eigenvalue \(\lambda_{d}\) is the variance of the data along eigenvector \(\mathbf{u}_{d}\), offering a natural choice of inverse eigenvalues \(\lambda_{d}^{-1}\) as the squared lengthscales \(\mathbf{\Sigma}_{dd}\) of the principal components. Let \(\tilde{\mathbf{x}}=\mathbf{U}^{T}\mathbf{x}\) denote the representation of the input \(\mathbf{x}\) in eigenbasis \(\mathbf{U}=[\mathbf{u}_{1}\,\mathbf{u}_{2}\,\dots\,\mathbf{u}_{D}]\). We compute the gradient kernel in PCA basis \(\mathbf{U}\) and set \(\mathbf{\Sigma}=\mathbf{\Lambda}^{-1}\):

\[\kappa(\mathbf{s},\mathbf{s}^{\prime})\stackrel{{\mathrm{def}}}{ {=}}\exp\left(-\frac{1}{2}(\mathbf{U}^{\top}\mathbf{s}-\mathbf{U}^{\top} \mathbf{s}^{\prime})^{\top}\,\mathbf{\Lambda}\left(\mathbf{U}^{\top}\mathbf{ s}-\mathbf{U}^{\top}\mathbf{s}^{\prime}\right)\right), \tag{14}\]

where \(\mathbf{\Lambda}\) is a diagonal matrix containing the eigenvalues. While setting the square inverse lengthscales equal to the eigenvalues seems problematic at first glance, since large eigenvalues will push the kernel \(\kappa\) towards 0, this problem is avoided in practice since we also employ the median heuristic, which introduces in the kernel a global bandwidth scaling term that adapts to the current pairwise distance between particles, as discussed below in Section 3.4.

Connection to the EmpCov priorRecently, Izmailov et al. (2021) propose the EmpCov prior for each weight column \(\mathbf{w}\) of the first layer of a BNN:

\[\mathbf{w}\sim\mathcal{N}(0,\alpha\mathbf{C}+\epsilon\mathbf{I}),\qquad \mathbf{w}\in\mathbb{R}^{D} \tag{15}\]

where \(\alpha>0\) determines the prior scale and \(\epsilon>0\) is a small constant ensuring that the covariance matrix is positive definite. This prior encourages these weights to vary more along high variance dimensions in the data manifold and less along the low variance dimensions. Samples from this prior will have large input gradients along the dimensions in the data manifold with high variances, and small input gradients along the low variance dimensions. In this sense, the EmpCov prior is similar to our approach. The difference is that while Izmailov et al. (2021) incorporates knowledge of the data manifold into the prior, we embed this knowledge into our approximate posterior via the kernel.

### Practical considerations

In this section, we detail two important considerations to make FoRDEs work in practice. We include the full training algorithm in the Appendix.
Mini-batchingTo make FoRDE amenable to mini-batch gradient optimization, we adapt the kernel in Eq. (10) to a mini-batch of samples \(\mathcal{B}=\{(x_{b},y_{b})\}_{b=1}^{B}\):

\[k_{\mathcal{B}}(\theta_{i},\theta_{j})=\frac{1}{B}\sum_{b=1}^{B}\kappa\Big{(} \nabla_{\mathbf{x}}f(\mathbf{x}_{b};\theta_{i})_{y_{b}},\nabla_{\mathbf{x}}f( \mathbf{x}_{b};\theta_{j})_{y_{b}}\Big{)} \tag{16}\]

However, using this kernel in the update rule in Eq. (8) leads to biased stochastic gradients of the repulsion term because the average over batch samples in Eq. (16) is inside the logarithm. In practice, we found no performance degradation or convergence issues in our experiments.

Median heuristicsSince we perform particle optimization with an RBF kernel \(\kappa\), following earlier works (Liu and Wang, 2016; Liu et al., 2019), we adopt the median heuristic (Scholkopf et al., 2002). Besides the lengthscales, we introduce a global bandwidth \(h\) in our base kernel in Eq. (11):

\[\kappa(\mathbf{s}_{i},\mathbf{s}_{j};\boldsymbol{\Sigma})=\exp\left(-\frac{1} {2h}(\mathbf{s}_{i}-\mathbf{s}_{j})^{\top}\boldsymbol{\Sigma}^{-1}(\mathbf{s} _{i}-\mathbf{s}_{j})\right),\qquad\mathbf{s}_{i}=\frac{\nabla_{\mathbf{x}}f( \mathbf{x};\theta_{i})_{y}}{||\nabla_{\mathbf{x}}f(\mathbf{x};\theta_{i})_{y} ||_{2}}\in\mathbb{R}^{D} \tag{17}\]

During training, the bandwidth \(h\) is adaptively set to \(\mathrm{med}^{2}/(2\log M)\), where \(\mathrm{med}^{2}\) is the median of the pairwise distance \((\mathbf{s}_{i}-\mathbf{s}_{j})^{\top}\boldsymbol{\Sigma}^{-1}(\mathbf{s}_{i} -\mathbf{s}_{j})\) between the weight samples \(\{\theta_{i}\}_{i=1}^{M}\).

### Computational complexity

Compared to DEs, FoRDEs take roughly three times longer to train. In addition to a forward-backward pass to calculate the log-likelihood, we need an additional forward-backward pass to calculate the input gradients, and another backward pass to calculate the gradients of the input gradient repulsion with respect to the weights. This analysis is confirmed in practice: in resnet18/cifar-100 experiments of Section 4.1 with an ensemble size of 10, a DE took around \(\sim\)31 seconds per epoch on an Nvidia A100 GPU, while FoRDE took around \(\sim\)101 seconds per epoch.

## 4 Experiments

In this section, we present experimental results of FoRDE on image classification tasks. We use cifar(Krizhevsky, 2009) and tinyimagenet(Le and Yang, 2015) datasets, which both have corrupted versions of the test set (Hendrycks and Dietterich, 2019). We use resnet18 (He et al., 2016) and preactresnet18 (He et al., 2016) architectures. We report mean and standard deviations over multiple runs with varying seeds.

### Comparisons to other repulsive DE methods and BNNs

We report performance of FoRDE against other methods on cifar-10/100 in Tables 1-2 and tinyimagenet in Table 3. Besides the PCA kernel introduced in Section 3.3, we also experiment with FoRDE where we set the lengthscales \(\boldsymbol{\Sigma}=\mathbf{I}\). For the repulsive DE (RDE) baselines, we choose weight-space RDE (D'Angelo and Fortuin, 2021), function-space RDE (D'Angelo and Fortuin, 2021) and feature-space RDE (Yashima et al., 2022). We also include the results of node-BNNs (Trinh et al.,

\begin{table}
\begin{tabular}{l c c c c} \hline \hline Method & NLL \(\downarrow\) & Accuracy (\% ) & ECE \(\downarrow\) & CA / CNLL / cECE \\ \hline node-BNNs & \(0.74\pm 0.01\) & \(79.7\pm 0.3\) & \(0.054\pm 0.002\) & 54.8 / 1.96 / 0.05 \\ SWAG & \(0.73\pm 0.01\) & \(79.4\pm 0.1\) & \(\mathbf{0.038\pm 0.001}\) & 53.0 / 2.03 / 0.05 \\ \hline Deep ensembles & \(\mathbf{0.70\pm 0.00}\) & \(81.8\pm 0.2\) & \(0.041\pm 0.003\) & 54.3 / 1.99 / 0.05 \\ weight-RDE & \(\mathbf{0.70\pm 0.01}\) & \(81.7\pm 0.3\) & \(0.043\pm 0.004\) & 54.2 / 2.01 / 0.06 \\ function-RDE & \(0.76\pm 0.02\) & \(80.1\pm 0.4\) & \(0.042\pm 0.005\) & 51.9 / 2.08 / 0.07 \\ feature-RDE & \(0.75\pm 0.04\) & \(\mathbf{82.1\pm 0.3}\) & \(0.072\pm 0.023\) & 54.8 / 2.02 / 0.06 \\ \hline FoRDE-PCA (ours) & \(0.71\pm 0.00\) & \(81.4\pm 0.2\) & \(0.039\pm 0.002\) & **56.1 / 1.90 / 0.05** \\ FoRDE-Identity (ours) & \(\mathbf{0.70\pm 0.00}\) & \(\mathbf{82.1\pm 0.2}\) & \(0.043\pm 0.001\) & 54.1 / 2.02 / 0.05 \\ \hline \hline \end{tabular}
\end{table}
Table 1: Results of resnet18 on cifar-100 evaluated over 5 seeds. Each ensemble has 10 members. ca, CNLL and cECE are accuracy, NLL, and ECE on cifar-100-c. FoRDE-PCA achieves the best performance under corruptions while FoRDE-Identity outperforms the baselines on clean data.
2022) and SWAG (Maddox et al., 2019) as reported in Trinh et al. (2022). We use a default ensemble size of 10. We use standard performance metrics of expected calibration error (ECE) (Naeini et al., 2015), negative log-likelihood (NLL) and predictive accuracy. For evaluations on input perturbations, we use the corrupted test images provided by Hendrycks and Gimpel (2017). We use resnet18 for cifar-10/100 and preactresnet18 for tinyimagenet. We report the hyperparameters and training procedure in the Appendix.

Tables 1-3 show that FoRDE-PCA outperforms other methods under input corruptions across all metrics, while maintaining competitive performance on clean data. For instance, FoRDE-PCA exhibits a \(+1.3\%\) gain on cifar-100-c and \(+2.4\%\) gain on cifar-10-c in accuracy compared to the second-best results. As stated in Section 3.3, the PCA kernel encourages FoRDE to rely more on features with high variances in the data manifold to make predictions, while being less dependent on features with low variances. This effect has been shown in Izmailov et al. (2021) to boost model robustness against perturbations, which explains why FoRDE with the PCA kernel performs better than the baselines on input corruptions.

On the other hand, Tables 1 and 2 show that FoRDE with identity lengthscales outperforms the baselines in terms of NLL on cifar-10 and has the best accuracy on cifar-100. However, FoRDE-Identity is slightly less robust than DE against corruptions. We suspect that with the identity lengthscales, FoRDE also learns to rely on low-variance features to make predictions, which is harmful to performance under corruptions (Izmailov et al., 2021).

The results of FoRDE-Identity suggest that an ensemble still needs to learn low-variance features in the data to generalize well on clean images. Hence, we believe that the optimal lengthscales for good performance on both clean and corrupted data lie somewhere between identity lengthscales and using the inverse square root eigenvalues as lengthscales.
et al. (2021) for all ensemble methods except FoRDE. By comparing Table 4 to Table 2, we see that the EmCov prior slightly improves robustness of the baseline ensemble methods against corruptions, while also leads to a small reduction in performance on clean data. These small improvements in robustness are not surprising, since Izmailov et al. (2021) stated that for ensemble methods consisting of approximate MAP solutions, the isotropic Gaussian prior already minimizes the influences of low variance data features on the ensemble's predictions. We argue that besides minimizing the influences of low variance features on predictions via the PCA kernel, FoRDE also encourages its members to learn complementary patterns that can explain the data well, and these two effects act in synergy to improve the robustness of the resulting ensemble. Thus, Table 4 shows that FoRDE is still more robust against corruptions than the baseline methods with the EmpCov prior.

### Performance under different ensemble sizes

We report the NLL of FoRDE and DE under different ensemble sizes on both cifar-100 and cifar-100-c in Fig. 2. We use the widerensnet16x4 (Zagoruyko and Komodakis, 2016) architecture for this experiment. These figures show that both methods enjoy significant improvements in performance as the ensemble size increases. While Fig. 1(a) shows that FoRDE underperforms DE on clean images, Fig. 1(b) shows that FoRDE significantly outperforms DE on corrupted images, such that a FoRDE with 10 members has the same corruption robustness of a DE with 30 members.

## 5 Related works

ParVI for neural networksD'Angelo and Fortuin (2021) studied repulsive DE in the weight and function space. They showed that by formulating the repulsion term according to ParVI, the ensemble will converge to the target posterior in the infinite particle limit. However, due to the shortcomings of the weight-space and function-space repulsion, their method showed little performance improvement over plain DEs. Subsequent work by Yashima et al. (2022) introduced feature-space repulsion, which is the output space of the penultimate layer. They argue that feature-space repulsion motivates ensemble members to learn diverse patterns in the training data, and showed performance improvements over DEs. Here, we propose the input-gradient space as a good

\begin{table}
\begin{tabular}{l c c c c} \hline \hline Method & NLL \(\downarrow\) & Accuracy (\%) \(\uparrow\) & ECE \(\downarrow\) & CA / CNLL / cECE \\ \hline Deep ensembles & \(0.119\pm 0.001\) & \(\mathbf{96.2\pm 0.1}\) & \(0.006\pm 0.001\) & 78.7 / 0.76 / 0.08 \\ weight-RDE & \(0.120\pm 0.001\) & \(96.0\pm 0.1\) & \(0.006\pm 0.001\) & 78.8 / 0.76 / 0.08 \\ function-RDE & \(0.132\pm 0.001\) & \(95.6\pm 0.3\) & \(0.007\pm 0.001\) & 77.8 / 0.79 / 0.08 \\ feature-RDE & \(\mathbf{0.118\pm 0.001}\) & \(\mathbf{96.2\pm 0.1}\) & \(\mathbf{0.005\pm 0.001}\) & 78.9 / 0.74 / 0.07 \\ \hline FoRDE-PCA (ours) & \(0.125\pm 0.001\) & \(96.1\pm 0.1\) & \(0.006\pm 0.001\) & \(\mathbf{80.5\text{ }/\mathbf{0.71}\text{ }/\mathbf{0.07}}\) \\ \hline \hline \end{tabular}
\end{table}
Table 4: Results of resnet18 on cifar-10 evaluated over 5 seeds. Each ensemble has 10 members. ca, CNLL and cECE are accuracy, NLL, and ECE on cifar-10-c. Here we use the EmpCov prior for all methods except FoRDE. FoRDE achieves the best performance under corruptions while maintaining competitive performance on clean data.

Figure 2: **FoRDE is competitive on in-distribution and outperforms DEs under domain shifts by corruption**. Performance of widerensnet16x4 on cifar-100 over 5 seeds.
alternative space in which to perform repulsion. We also note that an earlier work by Wang et al. (2019) also studied ParVI methods for function-space inference of BNNs.

Gradient-based attribution methods for deep modelsA well-known application of input gradients is to build attribution (or saliency) maps, which assign importance to visually-interpretable input features for a specified output (Simonyan et al., 2014). In addition to the pioneering work of Simonyan et al. (2014) defining attribution as absolute input gradients, subsequent works include Layer-wise Relevance Propagation (Bach et al., 2015), Gradient*Input (Shrikumar et al., 2016), Integrated Gradients (Sundararajan et al., 2017) and deepLIFT (Shrikumar et al., 2017). Our method intuitively utilizes the attribution perspective of input gradients to encourage ensemble members to learn different patterns from training data.

Improving corruption robustness of BNNsPrevious works have evaluated the predictive uncertainty of BNNs under covariate shift (Ovadia et al., 2019; Izmailov et al., 2021), with Izmailov et al. (2021) showing that standard BNNs with high-fidelity posteriors perform worse than MAP solutions on under corruptions. Izmailov et al. (2021) attributed this phenomenon to the lack of posterior contraction in the null space of the data manifold and proposed the EmpCov prior as a remedy. As discussed in Section 3.3, the PCA kernel works in the same manner as the EmpCov prior and thus significantly improves robustness of FoRDE against corruptions. Trinh et al. (2022) studied the robustness of node-BNNs, an efficient alternative to weight-based BNNs, and showed that by increasing the entropy of the posterior, node-BNNs become more robust against corruptions. Here, we incorporate the knowledge of the data manifold into the KDE approximation of the particle distribution via the PCA kernel, thus resulting in a manifold-structured approximate posterior that performs well under corruptions. Wang and Aitchison (2023) allow BNNs to adapt to the distribution shift at test time by using test data statistics, and this method can be applied to FoRDE to further boost its robustness against corruptions.

## 6 Discussion

In this section, we outline two directions to further improve FoRDE.

Choosing more optimal lengthscalesAs we have shown in our experiments in Section 4.1, choosing different lengthscale settings controls the predictive performance of FoRDE, as the lengthscales determine the strength of the repulsion force in each dimension in the input-gradient space. Here, we only experimented with two settings: the lengthscale setting based on eigenvalues presented in Section 3.3, and the identity lengthscales. We believe that there is an optimal lengthscale setting which allows FoRDE to perform well on both clean and corrupted data. Intuitively, the optimal lengthscales are still inversely proportional to the eigenvalues for corruption robustness, yet allowing some repulsion in the low-variance dimension in the data manifold for good in-distribution performance.

Choosing different orthonormal basisIn Section 3.3, we propose to perform repulsion in the eigenbasis \(\mathbf{U}\). Here, the role of \(\mathbf{U}\) is to rotate the repulsion forces so that they are aligned with the data manifold. With this observation, we can use different orthonormal basis to control the directions of the repulsion forces depending on the applications. For instance, we can align the repulsion forces with the orthonormal basis of the discrete cosine transformation (Ahmed et al., 1974) for image data.

## 7 Conclusion

In this work, we propose FoRDE, an ensemble learning method that promotes diversity in the input-gradient space among ensemble members. We detail the update rule and devise a data-dependent kernel suitable for input-gradient repulsion. Experiments on image classification tasks show that FoRDE outperforms other ensemble methods on corrupted images. Future directions include more theoretical understanding of our method and its behavior under different settings as discussed in Section 6.