Input-gradient space particle inference for neural network ensembles

TL;DR

This paper addresses the limited functional diversity in deep ensembles by proposing First-order Repulsive Deep Ensembles (FoRDE), which performs ParVI in the space of input gradients $\nabla_x f(x;\theta)_y$ to promote diverse feature learning. It defines a data-dependent kernel $k(\theta_i,\theta_j)$ based on gradient similarities and derives a Wasserstein gradient-descent update that combines a driving force toward the posterior with a repulsion force that encourages gradient diversity. The authors show that FaRDE, especially with PCA-based lengthscales (FoRDE-PCA), achieves superior accuracy and calibration under covariate shift from input perturbations and demonstrates strong transfer learning performance, outperforming standard deep ensembles and other repulsive methods, while offering notable robustness advantages. The practical impact is a scalable, gradient-space ensemble method that yields functionally diverse models with improved uncertainty estimation under distribution shifts and corruptions.

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 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.

Figures (12)

• Figure 1: Input-gradient repulsion increases functional diversity. 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.
• Figure 2: Input gradient ensembles (FoRDE) capture higher uncertainty than baselines. Each panel shows predictive uncertainty in 1D regression for different (repulsive) deep ensemble methods.
• Figure 3: Uncertainty of FoRDE is high in all input regions outside the training data, and is higher than baselines. Each panel shows the entropy of the predictive posteriors in 2D classification.
• Figure 4: FoRDE outperforms competing methods in transfer learning.First three columns: We report NLL, ECE and accuracy on in-distribution test sets and under covariate shift. For cifar-10, we use cinic10darlow2018cinic to evaluate models under natural shift and cifar-10-c for corruption shift. For cifar-100, we evaluate on cifar-100-c. FoRDE performs better than the baselines in all cases. Last column: We evaluate functional diversity by calculating epistemic uncertainty of ensembles on out-of-distribution (OOD) datasets using the formula in pmlr-v80-depeweg18a. We use cifar-100 as the OOD test set for cifar-10 and we use cifar-10 and cinic10 as OOD test sets for cifar-100. FoRDE exhibits higher functional diversity than the baselines.
• Figure 5: FoRDE is competitive on in-distribution and outperforms DEs under domain shifts by corruption. Performance of wideresnet16x4 on cifar-100 over 5 seeds.