Title: Enhancing Diversity in Bayesian Deep Learning via Hyperspherical Energy Minimization of CKA
Venue: NeurIPS 2024
Accepted: yes
Keywords: bayesian inference

variational inference

uncertainty quantification

deep learning

hypernetworks
Abstract: Particle-based Bayesian deep learning often requires a similarity metric to compare two networks. However, naive similarity metrics lack permutation invariance and are inappropriate for comparing networks. Centered Kernel Alignment (CKA) on feature kernels has been proposed to compare deep networks but has not been used as an optimization objective in Bayesian deep learning. In this paper, we explore the use of CKA in Bayesian deep learning to generate diverse ensembles and hypernetworks that output a network posterior. Noting that CKA projects kernels onto a unit hypersphere and that directly optimizing the CKA objective leads to diminishing gradients when two networks are very similar. We propose adopting the approach of hyperspherical energy (HE) on top of CKA kernels to address this drawback and improve training stability. Additionally, by leveraging CKA-based feature kernels, we derive feature repulsive terms applied to synthetically generated outlier examples. Experiments on both diverse ensembles and hypernetworks show that our approach significantly outperforms baselines in terms of uncertainty quantification in both synthetic and realistic outlier detection tasks.

Enhancing Diversity in Bayesian Deep Learning via Hyperspherical Energy Minimization of CKA

Abstract

Particle-based Bayesian deep learning often requires a similarity metric to compare two networks. However, naive similarity metrics lack permutation invariance and are inappropriate for comparing networks. Centered Kernel Alignment (CKA) on feature kernels has been proposed to compare deep networks but has not been used as an optimization objective in Bayesian deep learning. In this paper, we explore the use of CKA in Bayesian deep learning to generate diverse ensembles and hypernetworks that output a network posterior. Noting that CKA projects kernels onto a unit hypersphere and that directly optimizing the CKA objective leads to diminishing gradients when two networks are very similar. We propose adopting the approach of hyperspherical energy (HE) on top of CKA kernels to address this drawback and improve training stability. Additionally, by leveraging CKA-based feature kernels, we derive feature repulsive terms applied to synthetically generated outlier examples. Experiments on both diverse ensembles and hypernetworks show that our approach significantly outperforms baselines in terms of uncertainty quantification in both synthetic and realistic outlier detection tasks.

1 Introduction

Bayesian deep learning has always garnered substantial interest in the machine learning community. Instead of a point estimate which most deep learning algorithms obtain, a posterior distribution of trained models could significantly improve our understanding about prediction uncertainty and avoid overconfident predictions. Bayesian deep learning has potential applications in transfer learning, fairness, active learning, and even reinforcement learning, where reducing uncertainty can be used as a powerful intrinsic reward function (Yang & Loog, 2016; Ratzlaff et al., 2020; Wang et al., 2023).

One line of approach to Bayesian deep learning is to add noise to a single trained model. Such noises can either be injected during the training process, e.g. as in the stochastic gradient Langevin dynamics (Welling & Teh, 2011), or after the training process (Maddox et al., 2019). However, many such approaches often underperform the simple ensemble method (Lakshminarayanan et al., 2017b) which merely trains several deep networks with different random seeds. Intuitively, an ensemble, because it starts from different random initializations, might be able to 'explore" a larger portion of the parameter space than those that are always nearby one specific model or training path. Because of this, ensembles may capture different modes and therefore better represent the posterior distribution of 'well-trained" network functions (Fort et al., 2019; Wilson & Izmailov, 2020).

However, a critical question is, how different are the networks in an ensemble from one another? And can we utilize the idea of diversification to further improve these networks by making them even more diverse? In order to answer these questions, we first need a metric to compare those networks, which is in itself a significant problem; regular L1/L2 distances, either in the space of the network parameters, or in the space of the network activations, are not likely to work well. First, they suffer from the curse

Input

Layer

Inlier

(optional)

Models

Centering and Projection

Hyperspherical Energy

Model 1

of dimensionality due to the excessive number of parameters in modern deep networks. Moreover, there is the peculiar permutation invariance , where one can randomly permute the different channels of each layer and result in a network that has vastly different parameters and activations, yet represents the same function. The popular RBF kernel lacks this permutation invariance inhibiting methods like Stein Variational Gradient Descent (SVGD) from working effectively on larger networks (D' Angelo &Fortuin, 2021). Therefore, a proper kernel for comparing network functions should address these critical issues by being effective in high-dimensional spaces and invariant to permutations of neural network channels.

Kornblith et al. (2019) proposed an interesting approach for performing this comparison based on Centered Kernel Alignment (CKA). The idea is, instead of directly comparing activations or parameters, comparison is made between the Gram matrices of the same dataset fed into two different networks. Each example will generate a feature vector at each layer of the network, and a kernel matrix can be constructed based on the similarity between all example pairs in the dataset. Then, a CKA metric measures the similarity of these two Gram matrices as the similarity of the two networks. This idea addresses the permutation invariance issue and generates meaningful comparisons between deep networks.

In this paper, we propose to explicitly promote diversity of network functions by adding CKA-based loss terms to deep ensemble learning. Given that CKA projects all kernels on a hypersphere, we further propose to use Hyperspherical Energy ( HE ) minimization as an approach to more evenly distribute the ensemble of neural networks on the hypersphere. Experiments on synthetic data, MNIST, CIFAR, and TinyImageNet show that our approach maintains the predictive accuracy of ensemble models while boosting their performance in uncertainty estimation across both synthetic and realistic datasets. Besides, we demonstrate that our method can also be applied to training hypernetworks, improving the diversity and uncertainty estimation of the networks generated by a hypernetwork. Additionally, we propose using synthetic out-of-distribution (OOD) examples, to reduce their likelihood, and introducing feature repulsive terms on synthetic outlier examples to enhance OOD detection performance. We hope that our approach provides a different perspective to variational inference methods and contributes to improving uncertainty estimation in deep networks. Code is publicly available at https://github.com/Deep-Machine-Vision/he-cka-ensembles .

2 Related Work

Uncertainty Estimation. Alarge body of literature has studied the problem of uncertainty estimation for neural networks. Bayesian neural networks (Gal & Ghahramani, 2016; Krueger et al., 2017;

Nazaret & Blei, 2022) approximate a posterior distribution over the model parameters and therefore estimate the epistemic uncertainty of the predictions. Non-Bayesian approaches, on the other hand, rely on bootstrap (Osband et al., 2016), ensemble (Lakshminarayanan et al., 2017b; Wen et al., 2020; Park & Kim, 2022), and conformal prediction (Bhatnagar et al., 2023) to generate multiple neural networks of the same structure. Our approach is most closely related to ensemble methods for estimating predictive uncertainty. We follow the common practice of evaluating uncertainty by distinguishing between inlier and outlier images for datasets like CIFAR/SVHN and MNIST/FMNIST. Most approaches typically evaluate the separation or distance within the feature space of a model between inliers and outliers (Mukhoti et al., 2023; Van Amersfoort et al., 2020; D' Angelo & Fortuin, 2021; Ovadia et al., 2019b; Lakshminarayanan et al., 2017a; Liu et al., 2020). Deep deterministic uncertainty (DDU) utilizes a single deterministic network with a Gaussian mixture model (GMM) fitted on the observed inlier features before the last layer and calculates separation of inliers and outliers using feature log density. We refer the reader to Mukhoti et al. (2023) for more information. For a more comprehensive survey and benchmarking of different uncertainty estimation approaches, refer to Ovadia et al. (2019a); Gawlikowski et al. (2022).

ParVI. Particle-based variational inference methods (ParVI), such as Stein Variational Gradient Descent (SVGD) (Liu & Wang, 2016; Chen et al., 2018; Liu & Zhu, 2018), use particles to approximate the Bayes posterior. Our work most closely resembles work done by D' Angelo & Fortuin (2021), which explores adapting kernelized repulsive terms in both the weight and function space of deep ensembles to increase model diversity and improve uncertainty estimation. Our work, however, focuses more on constructing a new kernel rather than exploring new repulsive terms that utilize an RBF kernel on weights or network activations.

Hypernetworks. Hypernetworks have been used for various specific tasks, some are conditioned on the input data to generate the target network weights, such as in image conditioning or restoration (Alaluf et al., 2022; Aharon & Ben-Artzi, 2023). It has seen popular use in meta-learning tasks related to reinforcement learning (Beck et al., 2023, 2024; Sarafian et al., 2021), and few-shot learning in Zhmoginov et al. (2022). Hypernetworks conditioned on a noise vector to approximate Bayesian inference have been proposed (Krueger et al., 2018; Ratzlaff & Fuxin, 2019), but either require an invertible hypernet or do not diversify target features explicitly. Our motivation is to provide Bayesian hypernetworks by explicitly promoting feature diversity in target networks.

3 Measurements of Network Diversity

In order to generate an ensemble of diverse networks, we first need a measurement of similarity between internal network features. Throughout this paper, we denote a deep network with L layers as f ( x, θ ) = f L ( f ( ... ) ( f 1 ( x, θ 1 ) , · · · ) , θ L ) , where f l ( x, θ l ) ∈ R p l are the features of layer l parameterized by θ l , where p l ∈ N is the output feature dimension.

3.1 Comparing two networks with CKA

We will compare two networks with the same architecture layer-by-layer. Given two networks at layer l with weights θ 1 l , θ 2 l and feature activations f l ( x, θ 1 l ) , f l ( x, θ 2 l ) , a naive approach would be to take some Euclidean L k norm between the weights ∥ θ 1 l -θ 2 l ∥ k or features ∥ f l ( x, θ 1 l ) -f l ( x, θ 2 l ) ∥ k , but those tend to be bad heuristics for similarity measures in high-dimensional vector spaces due to the curse of dimensionality (Reddi et al., 2014; Aggarwal et al., 2001; Weber et al., 1998). A better approach to measuring similarity would be to analyze the statistical independence or alignment of features between networks, through Canonical Correlation Analysis (CCA), Singular Vector CCA (SVCCA), Projection-Weighted CCA (PWCCA), Orthogonal Procrustes (OP), Hilbert-Schmidt Independence Criterion (HSIC), or Centered Kernel Alignment (CKA)(Raghu et al., 2017; Gretton et al., 2005; Kornblith et al., 2019). Ideally, the chosen metric should be computationally efficient, invariant to isotropic scaling, orthogonal transformations, permutations, and be easily differentiable. However, CCA methods and OP require the use of Singular Value Decomposition (SVD) or iterative approximation methods, which can be computationally intensive. Additionally, HSIC and OP are not invariant to isotropic scaling of f l .

As a comparison metric between networks, Kornblith et al. (2019) propose to utilize CKA on Gram matrices, obtained by evaluating the neural network on a finite sample. CKA is based on the nonparametric statistical independence criterion HSIC, which has been a popular method of measuring

statistical independence as a covariance operator in the kernel Hilbert spaces (Gretton et al., 2005). An empirical estimation of HSIC on a dataset of N examples is given by 1 / ( N -1) 2 tr ( K 1 HK 2 H ) , where the two Gram matrices K 1 i,j = k ( f l ( x i , θ 1 l ) , f l ( x j , θ 1 l )) and K 2 i,j = l ( f l ( x i , θ 2 l ) , f l ( x j , θ 1 l )) are constructed through the k ( · , · ) kernel function, and H = I -1 N 11 ⊺ a centering matrix to center the Gram matrices around the row means, where 1 denotes the all ones vector, and I as the identity matrix. This function, however, is not invariant to isotropic scaling. The isotropic scaling invariant version of HSIC is termed Centered Kernel Alignment (CKA) (Kornblith et al., 2019),

We stick with the linear kernel for k , unless otherwise specified, due to its computational simplicity. The RBF kernel works, but it is computationally expensive and requires the use of heuristic like the median heuristic to perform well and make CKA isotropic scaling invariant (Reddi et al., 2014; Kornblith et al., 2019).

3.2 Generalizing to multiple networks

Given an ensemble of M models, a simple approach to generalizing Eq. (1) to measure the similarity of an ensemble would be to construct a pairwise alignment metric. For each layer l of each member of the ensemble m , we construct the set of kernel matrices K = { K l m } m =1 ,...,M l =1 ,...,L . The mean pairwise loss across all layers L is as follows,

̸

In its current form, CKApw provides a good approximate metric to evaluate the similarity among members of an ensemble. We found that rewriting Eq. (2) gives us another perspective on optimizing CKA. First, to simplify notation, let ¯ K m = 1 ∥ K m H ∥ F vec( K m H ) be the centered and normalized Gram matrix, and rewriting the inner product in Eq. (1) results in the cosine similarity metric CKA ( K m , K m ′ ) = ¯ K m ⊤ ¯ K m ′ . The matrix of the vectorized kernels from the set K l can be represented in a compact form K l , and CKA pw can be rewritten using this compact form,

where zd ( X ) = X ⊙ ( 11 ⊺ -I ) is a function that zeros out the diagonal of a matrix.

Now each row m of K l is a vectorized Gram matrix with unit length from the model m . We can view these vectors as the Gram matrices projected on the unit hypersphere as shown in Fig. 1. For each pair of models i, j on the hypersphere, with an angle ϕ i,j between the feature gram vectors, CKA is equivalent to cos( ϕ i,j ) . Thus minimizing pairwise CKA would reduce the sum of cos( ϕ i,j ) , pushing Gram matrices between model pairs apart.

3.3 Comparing Networks with Hyperspherical Energy

Note that CKA suffices as a differentiable measure between deep networks and one can directly minimize CKA to push different models in the ensemble apart from each other. However, CKA may have a specific deficiency as an optimization objective in that the gradient of cos( ϕ ) is -sin( ϕ ) , which is close to 0 when ϕ is close to 0 . In other words, if two models are already very similar to each other (their CKA being close to 1 ), then optimizing with CKA may not provide enough gradient to move them apart. Hence, we explore further techniques to alleviate this drawback. Minimum Hyperspherical Energy (MHE) (Liu et al., 2021) aims to distribute particles uniformly on a hypersphere, which maximizes their geodesic distances from each other. In physics, this is analogous to distributing electrons with a repellent Coloumb's force.

Inspired by MHE, we propose to adopt the idea of hyperspherical energy ( HE ) on top of the CKA kernel to compare neural networks, termed HE-CKA , which is novel to our knowledge. For each

X

1.0

0.5

0.0

0.5

1.0

Y

1.0

0.5

0.0

0.5

1.0

Z

1.0

0.5

0.0

0.5

1.0

(a) Initial Points

X

1.0

0.5

0.0

0.5

1.0

Y

1.0

0.5

0.0

0.5

1.0

Z

1.0

0.5

0.0

0.5

1.0

(b) cossim Minimized

X

1.0

0.5

0.0

0.5

1.0

Y

1.0

0.5

0.0

0.5

1.0

Z

1.0

0.5

0.0

0.5

1.0

(c)

HE

Minimized

0

10

20

30

40

50

Iterations

0.0

0.1

0.2

0.3

0.4

0.5

0.6

Cosine Similarity

MHE

(d) cossim Objective

0

10

20

30

40

50

Iterations

2

4

6

8

10

Mean Hyperspherical Energy

MHE

(e)

HE

Objective

layer l we treat the M model Gram vectors ¯ K l m as particles on the hypersphere, its geodesic on the hypersphere is then d i,j = arccos( CKA ( K i l , K j l )) = arccos( ¯ K i ⊺ l ¯ K j l ) , we define the energy function by simulating a repellent force on the particles via F i,j = ( d i,j ) -s as shown in Fig 1. Incorporating this across all layers, weighted by w l , and model pairs results in the overall hyperspherical energy of CKA between all models is.

̸

where s > 0 is the Riesz s -kernel function parameter. For more information regarding the layer weighting w l and smoothing terms please see Appendix C.

HE has been shown as a proper kernel (Liu et al., 2021). The minimization of HE , as mentioned in Liu et al. (2021), asymptotically corresponds to the uniform distribution on the hypersphere, In order to demonstrate the difference between HE and the pairwise cosine similarity, we conducted a test on a synthetic dataset by generating random vectors from two Gaussian distributions in R 3 , and projecting on the unit hypersphere. We then minimized pairwise cosine similarity and HE respectively. Figure 2 illustrates that minimizing HE converges faster and achieves a more uniform distribution compared to minimizing cosine similarity. Specifically, as observed in Figure 2 (b), minimizing the cosine similarity loss caused particles to cluster towards two opposite sides of the sphere, as the gradient of this optimization - as mentioned in the beginning of the subsection, - becomes very small between particles that are clustered together. In Fig. 2(d), we show that minimizing HE actually leads to lower cosine similarity than directly minimizing cosine similarity, showing that minimizing cosine similarity could fall into local optima as described.

4 Particle-based Variational Inference by Minimizing Model Similarity

Armed with the comparison metrics between deep networks, we now proceed to incorporate the minimization of network similarity into deep ensemble training. In this section, we explore two different types of ensembles. The first is a regular ensemble where deep networks are trained to maximize the data likelihood, and we would add a term minimizing model similarity to it. Afterwards, we also explore the application of the idea on generative ensembles by hypernetworks, which aims to train a generator that generates network weights so that one can directly sample the posterior from it. Such a generator can easily exhibit mode collapse by always generating the same function, and we hope the idea of minimizing the similarity of generated networks would help alleviate this issue.

Suppose we are given a deep network with L layers f ( x, θ ) = f L ( f ( ... ) ( f 1 ( x, θ 1 ) , · · · ) , θ L ) . We denote the ensemble of target network layer at layer l as E l ( x, θ ) = [ f l ( x, θ 1 ) , ..., f l ( x, θ M )] , with the ensemble parameters θ = { θ m } M m =1 , and the training set of N examples as D = { x i , y i } N i =1 . From a Bayesian perspective, incorporating CKA pw / HE-CKA into the ensemble training can be interpreted as imposing a Boltzmann prior with HE-CKA over the ensemble network parameters θ that produce feature Gram matrices uniformly distributed on the unit hypersphere. Specifically, p ( θ ) ∝ exp( -γ HE-CKA( K ( E ( · ) ( · , θ )) , where K ( E ( · ) ( · , θ )) is the set of feature Gram matrices, constructed from the ensemble E , as described in Sec. 3.2. The posterior distribution now becomes:

(a) Ensemble

(b) SVGD + RBF

(c) fSVGD + RBF

RBF

(e) fKDE WGD +

RBF

(f) SVGD +

HE-CKA

(g) Ensemble +

HE-CKA

(h) Ensemble + OOD

HE-CKA

(i) Hypernetwork

(j) Hypernetwork +

OOD

HE-CKA

10

10

10

10

10

5

1.0

5

1.0

5

1.0

5

1.0

5

1.0

0

0

0

0

0

0.5

0.5

0.5

0.5

0.5

5

5

-5

-5

-5

0.0

-10

0.0

-10

0.0

-10

0.0

0

10

0

10

-10

0

10

-10

0

10

-10

0

10

10

10

10

10

10

5

1.0

5

1.0

5

1.0

5

1.0

5

1.0

0

0

0

0

0

0.5

0.5

0.5

0.5

0.5

-5

-5

-5

5

-5

-10

0.0-10

0.0-10

10.0

-10

0.0-10

0.0

-10

0

10

-10

0

10

-10

0

10

-10

0

10

-10

0

10

The MAP estimate of the posterior in Eq. (5) results in the following objective:

The left hand side is the negative log likelihood term where L ( x, y ) is the target loss, such as crossentropy or MSE. Minimizing θ , while adjusting the constant γ used in the Boltzmann prior, allows us to balance between gram matrix hyperspherical uniformity and fitting the training data. Further explanation of Eq. (6)'s relationship to ParVI is given in Appendix A. Note that the same approach can be used to derive the formula for the CKA kernel in Eq. (2) as well.

4.1 Diverse Generative Ensemble with Hypernetworks

Besides diversifying ensemble models, we also explore using CKA pw / HE-CKA in learning a nondeterministic generator (Krueger et al., 2018) which gives us the ability to sample from a continuous nonlinear posterior distribution of network weights. This is appealing since it can generate any amount of network with a single training run of the generator, without being restricted by the fixed amount of posterior samples one can access with a regular ensemble.

The approach we take uses the concept of hypernetworks (Ha et al., 2016; Krueger et al., 2018). However, current variational inference methods are not scalable to larger models and generally require a change of variables or invertible functions (Krueger et al., 2018). Naively using a hypernetwork to transform a prior distribution to generate θ of the target network may result in the collapse of the posterior θ distribution. Hence, it would be interesting to explore using CKA pw / HE-CKA to avoid such mode collapses. We use the surrogate diversity loss in Eq. (4) to impose non-parametric independence of feature distributions. With hypernetworks we aim to transform, using a network h ( z ) , some prior distribution z ∼ N ( 0 , I ) , z ∈ R P to h ( z ) = θ ∈ R ∑ l w l , where P is the dimensionality of the latent space, and w l the number of parameters for layer l . To learn the function h ( · ) we sample a batch M of θ 's, feed through the ensemble E ( x, θ ) and calculate loss, similar to a fixed ensemble as in Eq. (6). With the difference being that now we are backpropagating gradients to h ( · ) accumulated from the M ensemble members.

Using a plain MLP for the hypernetwork h would require the last layer's weight matrix to contain ∑ l w l × J entries, where J is the activation dimension right before the last layer. This could possibly result in a matrix of millions of trainable parameters. To overcome this challenge we follow the approach by Ratzlaff & Fuxin (2019) of decomposing h into several parts. First a layer code generator h ( z ) = c ∈ R L,c size , and the layer generators θ l = g l ( c l ) , where each layer generator g is a separate smaller network per layer l . See Fig. 4 for a visualization. Note c is a matrix with L layer codes of size c size .

Conv2D

Conv2D

Linear

Linear

To further reduce size of the hypernetwork, for convolutional networks, we use the assumption that filters in convolutional layers can be independently sampled. For each convolutional layer l we create layer code vectors via the layer code generator c l = h ( z l ) , where each code vector i in c l i corresponds to a latent vector for a single convolution filter i . We feed each filter code i through a filter generator g l ( c l i ) separately to generate the filter for layer l . An example architecture can be seen in Fig. 4.

4.2 Synthetic OOD Feature Diversity

Striking a balance between ensemble member diversity and inlier performance is a challenge. Enforcing strong feature dissimilarity on observed inlier examples could degrade inlier performance if not tuned correctly. ParVI methods that only observe inlier points, like SVGD, can achieve better diversity but often at the expense of inlier accuracy (D' Angelo & Fortuin, 2021). We have found that a more effective strategy is to reduce the feature similarity on obvious OOD examples, and reduce their likelihood, which could be synthetically generated. Intuitively, we want more diverse features on obvious outlier examples to indicate uncertainty because the networks trained on these examples should not be confident. We found this approach to generate OOD examples and increase their feature diversity to be very effective.

Importantly, the OOD points do not need to be close to the inlier data manifold at all. For images, we generate outlier points via random grids, lines, perlin noise, simplex noise, and vastly distorted and broken input samples. See Appendix E.2 for more details and example images. For vector datasets, such as the test 2D datasets presented in Fig. 3, we identify outlier points by locating the minimum and maximum values across training examples. Generally, the boundary does not need to be close to the in-distribution (ID) dataset to achieve good results. We split the Gram matrices into K ID and K OOD and apply HE-CKA to them separately, with respective hyperparameters γ ID and γ OOD. The parameter value γ ID can be adjusted to be smaller than γ OOD. Additionally, for classification tasks, we add an entropy-maximizing term, scaled by hyperparameter β , for synthetic OOD points to Eq. (6). Similar loss terms may be constructed for other tasks, such as variance for regression tasks, but we have not explored them yet.

(a) Ensemble

(b) Ensemble +

HE-CKA

(c) Hypernetwork

(d) Hypernetwork +

Function

Function

Function

Function

Posterior Mean

Posterior Mean

Posterior Mean

Posterior Mean

Training

Training

Training

Training

HE-CKA

(a) Ensemble

(b) SVGD + RBF

(c) SVGD +

HE-CKA

(d) Ensemble + OOD

HE-CKA

1.00

1.00

1.00

1.00

0.75

0.75

0.75

0.75

Fraction

0.3

Fraction

0.3

Fraction

Fraction

0.3

0.2

0.2

0.2

0.2

0.1

0.1

0.1

0.1

0.0

Entropy

0.0

Entropy

0.0

Entropy

0.0

Entropy

MNIST

AMNIST

FMNIST

MNIST

AMNIST

FMNIST

MNIST

AMNIST

FMNIST

MNIST

AMNIST

FMNIST

5 Experiments

In this section, we conduct experiments on several datasets, ranging from synthetic tasks to realistic out-of-distribution (OOD) detection problems, to validate our approach. We compare OOD results between ensembles, ParVI, and other baselines.

5.1 Synthetic Data

We start by testing our approach on two synthetic tasks to visually assess the uncertainty estimation capability on both classification and regression problems. The first task is a 2D four-class classification problem, where each class is distributed in one quadrant of the 2D space with points sampled from Gaussians with σ = ( . 4 , . 4) and µ = ( ± 2 , ± 2) . The objective is to evaluate whether the models can accurately predict uncertainty, ideally showing low uncertainty near training examples and high uncertainty elsewhere. We employed a three-layer MLP trained with cross-entropy on the four classes and measured the predictive entropy of points sampled uniformly from a 10 × 10 grid.

When using cross-entropy alone, the decision boundaries among the four classes tend to be very similar. Deep ensembles classify with high confidence in most areas where they have never observed data before (Fig. 3(a)). Introducing the HE-CKA diversity term to the ensemble significantly reduces the ensemble's confidence on points outside the in-distribution set (Fig.3(g)). Furthermore, incorporating the HE-CKA and entropy term for OOD points allows the model to better estimate uncertainty, with only inliers being confident (Fig. 3(h)). In the case of hypernetworks, we observe the importance of a diversity term. Without it, hypernetwork predictions tend to be overconfident on outliers (Fig. 3(i)). However, when introducing HE-CKA hypernetwork, we achieve results closely resembling that of the ensemble + HE-CKA term (Fig. 3(j)).

In our second test, we perform a 1D regression modeling task. We aim to learn the function y ( x ) = -sin (1 . 2 x ) (1 + x ) within x ∈ ( -6 , 6) with high certainty everywhere except in x ∈ ( -2 , 2) . The training dataset involves sampling the function with 40 points uniformly from both ( -6 , -2) and (2 , 6) , with 2 points from ( -2 , 2) . We then fit a four layer MLP to approximate y ( x ) .

The visual result of each method is shown in Fig. 5. The fixed ensemble (Fig. 5(a)) has little diversity between the areas with low density, in contrast to the ensemble plus the HE-CKA term (Fig. 5(b)). The hypernetwork, without any feature diversity term (Fig. 5(c)) collapses, producing very similar weights. However, adding the HE-CKA term to the hypernetwork (Fig. 5(d)) alleviates this issue.

5.2 OODDetection on Real Datasets

We evaluated our proposed approach on a variety of real-world datasets, including Dirty-MNIST, Fashion-MNIST, CIFAR-10/100, SVHN, and TinyImageNet. We employ different CNN architectures such as LeNet, ResNet32, and ResNet18 to demonstrate the versatility of our method across models of varying complexity. Our experiments compare the out-of-distribution (OOD) detection performance of our approach against several approaches, including Deep Deterministic Uncertainty (DDU), deep ensembles and Stein Variational Gradient Descent (SVGD) equipped with the RBF kernel.

We provide experimental settings and training details here and additionally in Appendix C. Limitations of this approach are discussed in Appendix D, while further insights into memory usage and computational efficiency are discussed in Appendix G. Details regarding synthetic OOD example generation is described in Appendix E.2.

Dirty-MNIST vs Fashion MNIST. The Dirty-MNIST vs Fashion MNIST OOD benchmark (Mukhoti et al., 2023) examines the capability of models to discern inliers, OOD data in a similar distribution and OOD data in a more dissimilar distribution. This dataset combines MNIST with more ambiguous and challenging examples known as ambiguous MNIST (AMNIST). We examine the ability to to distinguish MNIST and AMNIST from the out-of-distribution Fashion MNIST (FMNIST) (Xiao et al., 2017). We trained a LeNet5 on Dirty-MNIST and assessed OOD classification between Dirty-MNIST and FMNIST using predictive entropy (PE) and mutual information (MI). The area under the receiver operating characteristic (AUROC) is used to assess the separability between MNIST and FMNIST (Lecun et al., 1998; Bradley, 1997). For DDU a GMM is fit to the second to last layer's features over Dirty-MNIST, using the log density of features to distinguish inliers from outliers, rather than predictive entropy.

Results shown in Fig. 6 and Table 1. Ensembles (Fig. 6(a)) do not exhibit significant separation using predictive entropy (PE) alone, resulting in low AUROC with PE for FashionMNIST. While methods like SVGD, equipped with an RBF kernel, improve separation (Fig. 6(b)), our approach demonstrates that using HE-CKA on an ensemble alone surpasses both RBF kernels and DDU's approach, and when paired with OOD examples, and OOD likelihood minimization, results in almost perfect separation with 99.99% AUROC. CKA plots of each method are presented in Appendix F.

CIFAR-10/100 vs SVHN. We further evaluated our method on CIFAR-10 and CIFAR-100 datasets, testing outlier detection performance on SVHN (Table 2) (Netzer et al., 2011). For a fair comparison with D' Angelo & Fortuin (2021), we trained ResNet32 ensembles following the training procedure and parameters described by D' Angelo & Fortuin (2021). For more details regarding model architecture please refer to the aforementioned paper and published code. We used predictive entropy and mutual information for the OOD classification, with the exception of DDU using feature space density (Mukhoti et al., 2023).

Given that the network presented in D' Angelo & Fortuin (2021) has significantly fewer parameters than a typical ResNet, it is expected to see an inferior classification accuracy to that of standard ResNet. In order to show that our approach generalizes to larger networks, we trained on larger ResNet18 ensembles. Results in Table. 3 show that HE-CKA can maintain similar accuracy as regular deep ensembles while significantly improving on ECE and AUROC of outliers. For the CIFAR-100 results please see Appendix C.3. Our ensemble with a standard ResNet18 with batch normalization even slightly outperforms a WideResNet-28-10 (WRN) using the approach by Mukhoti et al. (2023). Additionally, the mean inference time for a WRN is 13ms compared to 9ms for the ResNet18 ensemble on a Quadro RTX 8000.

TinyImageNet vs SVHN/CIFAR-10/CIFAR-100/DTD. To further evaluate the effectiveness of our approach to larger models and more complex datasets, we conducted experiments using the TinyImageNet dataset (Le & Yang, 2015). We trained ensembles of ResNet18 models and tested their ability to detect OOD samples from SVHN (Netzer et al., 2011), CIFAR-10/100 (Krizhevsky, 2009), and the Describable Textures Dataset (DTD) (Cimpoi et al., 2014). Our objective was to assess whether the proposed methods could generalize to large-scale settings and improve OOD detection performance without compromising in-distribution accuracy. Training details, and data splits, are provided in Appendix C.4.

Our proposed methods, especially Ensemble+OOD HE-CKA , enhanced OOD detection performance. Notably, Ensemble+OOD HE-CKA achieved an AUROC of 99.31% on SVHN and substantial improvements on CIFAR-10/100 and DTD datasets (Table. 4), with AUROC scores of 81.56%/87.64% and 90.94%, respectively. This improvement in OOD detection did not come at a major expense of ID accuracy.

6 Conclusion

In this paper, we explored the novel usage of CKA and MHE on feature kernels to diversify deep networks. We demonstrated that HE-CKA is an effective way to minimize pairwise cosine similarity, thereby enhancing feature diversity in ensembles and hypernetworks when applied on top of CKA. Our approach significantly improves the uncertainty estimation capabilities of both deep ensembles and hypernetworks, as evidenced by experiments on synthetic classification/regression tasks and real image outlier detection tasks. We showed that diverse ensembles utilizing predictive entropy alone can outperform other feature space density approaches, while synthetically generated OOD examples, far from the inlier distribution, can further significantly improve the OOD detection performance. While our current method requires fine-tuning several hyperparameters, such as layer weighting, we believe that future work could explore strategies for automatically estimating these parameters. We hope that our method inspires further advancements in Bayesian deep learning, extending its application to a wider range of tasks that require robust uncertainty estimation.

References

Aggarwal, C. C., Hinneburg, A., and Keim, D. A. On the surprising behavior of distance metrics in high dimensional spaces. In Proceedings of the 8th International Conference on Database Theory , ICDT '01, pp. 420-434, Berlin, Heidelberg, 2001. Springer-Verlag. ISBN 3540414568.

Aharon, S. and Ben-Artzi, G. Hypernetwork-based adaptive image restoration. In ICASSP 2023 2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) , pp. 1-5, 2023. doi: 10.1109/ICASSP49357.2023.10095537.

Alaluf, Y., Tov, O., Mokady, R., Gal, R., and Bermano, A. Hyperstyle: Stylegan inversion with hypernetworks for real image editing. In 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pp. 18490-18500, 2022. doi: 10.1109/CVPR52688.2022.01796.

Ambrosio, L., Gigli, N., and Savare, G. Gradient Flows: In Metric Spaces and in the Space of Probability Measures . Springer Science & Business Media, 2008.

Beck, J., Jackson, M. T., Vuorio, R., and Whiteson, S. Hypernetworks in meta-reinforcement learning. In Liu, K., Kulic, D., and Ichnowski, J. (eds.), Proceedings of The 6th Conference on Robot Learning , volume 205 of Proceedings of Machine Learning Research , pp. 1478-1487. PMLR, 14-18 Dec 2023. URL https://proceedings.mlr.press/v205/beck23a.html .

Beck, J., Vuorio, R., Xiong, Z., and Whiteson, S. Recurrent hypernetworks are surprisingly strong in meta-rl. Advances in Neural Information Processing Systems , 36, 2024.

Bhatnagar, A., Wang, H., Xiong, C., and Bai, Y. Improved online conformal prediction via strongly adaptive online learning. In Proceedings of the 40th International Conference on Machine Learning , ICML'23. JMLR.org, 2023.

Bradley, A. P. The use of the area under the roc curve in the evaluation of machine learning algorithms. Pattern Recognition , 30(7):1145-1159, 1997. ISSN 0031-3203. doi: https://doi.org/10.1016/ S0031-3203(96)00142-2. URL https://www.sciencedirect.com/science/article/pii/ S0031320396001422 .

Chen, C., Zhang, R., Wang, W., Li, B., and Chen, L. A unified particle-optimization framework for scalable bayesian sampling. arXiv preprint arXiv:1805.11659 , 2018.

Cimpoi, M., Maji, S., Kokkinos, I., Mohamed, S., and Vedaldi, A. Describing textures in the wild. In Proceedings of the IEEE Conf. on Computer Vision and Pattern Recognition (CVPR) , 2014.

D' Angelo, F. and Fortuin, V. Repulsive deep ensembles are bayesian. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems , volume 34, pp. 3451-3465. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper_files/paper/2021/file/ 1c63926ebcabda26b5cdb31b5cc91efb-Paper.pdf .

Fort, S., Hu, H., and Lakshminarayanan, B. Deep ensembles: A loss landscape perspective. arxiv 2019. arXiv preprint arXiv:1912.02757 , 2019.

Gal, Y. and Ghahramani, Z. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning , pp. 1050-1059. PMLR, 2016.

Gawlikowski, J., Tassi, C. R. N., Ali, M., Lee, J., Humt, M., Feng, J., Kruspe, A., Triebel, R., Jung, P., Roscher, R., Shahzad, M., Yang, W., Bamler, R., and Zhu, X. X. A survey of uncertainty in deep neural networks, 2022. URL https://arxiv.org/abs/2107.03342 .

Gretton, A., Bousquet, O., Smola, A., and Schölkopf, B. Measuring statistical dependence with hilbert-schmidt norms. In Jain, S., Simon, H. U., and Tomita, E. (eds.), Algorithmic Learning Theory , pp. 63-77, Berlin, Heidelberg, 2005. Springer Berlin Heidelberg. ISBN 978-3-540-316961.

Ha, D., Dai, A., and Le, Q. V. Hypernetworks, 2016.

Kornblith, S., Norouzi, M., Lee, H., and Hinton, G. Similarity of neural network representations revisited. In Chaudhuri, K. and Salakhutdinov, R. (eds.), Proceedings of the 36th International Conference on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pp. 35193529. PMLR, 09-15 Jun 2019. URL https://proceedings.mlr.press/v97/kornblith19a. html .

Krizhevsky, A. Learning multiple layers of features from tiny images. Technical report, 2009.

Krueger, D., Huang, C.-W., Islam, R., Turner, R., Lacoste, A., and Courville, A. Bayesian hypernetworks. arXiv preprint arXiv:1710.04759 , 2017.

Krueger, D., Huang, C.-W., Islam, R., Turner, R., Lacoste, A., and Courville, A. Bayesian hypernetworks, 2018.

Lakshminarayanan, B., Pritzel, A., and Blundell, C. Simple and scalable predictive uncertainty estimation using deep ensembles. In Proceedings of the 31st International Conference on Neural Information Processing Systems , NIPS'17, pp. 6405-6416, Red Hook, NY, USA, 2017a. Curran Associates Inc. ISBN 9781510860964.

Lakshminarayanan, B., Pritzel, A., and Blundell, C. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems , 30, 2017b.

Le, Y. and Yang, X. Tiny imagenet visual recognition challenge. CS 231N , 7(7):3, 2015.

Lecun, Y., Bottou, L., Bengio, Y., and Haffner, P. Gradient-based learning applied to document recognition. Proceedings of the IEEE , 86(11):2278-2324, 1998. doi: 10.1109/5.726791.

Liu, C. and Zhu, J. Riemannian stein variational gradient descent for bayesian inference. In Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence and Thirtieth Innovative Applications of Artificial Intelligence Conference and Eighth AAAI Symposium on Educational Advances in Artificial Intelligence , AAAI'18/IAAI'18/EAAI'18. AAAI Press, 2018. ISBN 978-157735-800-8.

Liu, C. and Zhu, J. Geometry in sampling methods: A review on manifold mcmc and particle-based variational inference methods. Advancements in Bayesian Methods and Implementations , 47:239, 2022.

Liu, C., Zhuo, J., Cheng, P., Zhang, R., and Zhu, J. Understanding and accelerating particle-based variational inference. In International Conference on Machine Learning , pp. 4082-4092. PMLR, 2019.

Liu, J., Lin, Z., Padhy, S., Tran, D., Bedrax Weiss, T., and Lakshminarayanan, B. Simple and principled uncertainty estimation with deterministic deep learning via distance awareness. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems , volume 33, pp. 7498-7512. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper_files/paper/2020/file/ 543e83748234f7cbab21aa0ade66565f-Paper.pdf .

Liu, Q. and Wang, D. Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances in neural information processing systems , 29, 2016.

Liu, W., Lin, R., Liu, Z., Xiong, L., Schölkopf, B., and Weller, A. Learning with hyperspherical uniformity. In International Conference On Artificial Intelligence and Statistics , pp. 1180-1188. PMLR, 2021.

Maddox, W. J., Izmailov, P., Garipov, T., Vetrov, D. P., and Wilson, A. G. A simple baseline for bayesian uncertainty in deep learning. Advances in neural information processing systems , 32, 2019.

Mukhoti, J., Kirsch, A., van Amersfoort, J., Torr, P. H., and Gal, Y. Deep deterministic uncertainty: A new simple baseline. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pp. 24384-24394, June 2023.

Nazaret, A. and Blei, D. Variational inference for infinitely deep neural networks. In International Conference on Machine Learning , pp. 16447-16461. PMLR, 2022.

Netzer, Y., Wang, T., Coates, A., Bissacco, A., Wu, B., Ng, A. Y., et al. Reading digits in natural images with unsupervised feature learning. In NIPS workshop on deep learning and unsupervised feature learning , volume 2011, pp. 4. Granada, 2011.

Osband, I., Blundell, C., Pritzel, A., and Van Roy, B. Deep exploration via bootstrapped dqn. Advances in neural information processing systems , 29, 2016.

Ovadia, Y., Fertig, E., Ren, J., Nado, Z., Sculley, D., Nowozin, S., Dillon, J., Lakshminarayanan, B., and Snoek, J. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R. (eds.), Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019a. URL https://proceedings.neurips.cc/paper_files/paper/2019/file/ 8558cb408c1d76621371888657d2eb1d-Paper.pdf .

Ovadia, Y., Fertig, E., Ren, J., Nado, Z., Sculley, D., Nowozin, S., Dillon, J. V ., Lakshminarayanan, B., and Snoek, J. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift . Curran Associates Inc., Red Hook, NY, USA, 2019b.

Park, N. and Kim, S. Blurs behave like ensembles: Spatial smoothings to improve accuracy, uncertainty, and robustness. In International Conference on Machine Learning , pp. 17390-17419. PMLR, 2022.

Raghu, M., Gilmer, J., Yosinski, J., and Sohl-Dickstein, J. Svcca: Singular vector canonical correlation analysis for deep learning dynamics and interpretability. Advances in neural information processing systems , 30, 2017.

Ratzlaff, N. and Fuxin, L. HyperGAN: A generative model for diverse, performant neural networks. In Chaudhuri, K. and Salakhutdinov, R. (eds.), Proceedings of the 36th International Conference on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pp. 5361-5369. PMLR, 09-15 Jun 2019. URL https://proceedings.mlr.press/v97/ratzlaff19a.html .

Ratzlaff, N., Bai, Q., Fuxin, L., and Xu, W. Implicit generative modeling for efficient exploration. In III, H. D. and Singh, A. (eds.), Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pp. 7985-7995. PMLR, 13-18 Jul 2020. URL https://proceedings.mlr.press/v119/ratzlaff20a.html .

Reddi, S. J., Ramdas, A., Póczos, B., Singh, A., and Wasserman, L. On the decreasing power of kernel and distance based nonparametric hypothesis tests in high dimensions, 2014. URL https://arxiv.org/abs/1406.2083 .

Sarafian, E., Keynan, S., and Kraus, S. Recomposing the reinforcement learning building blocks with hypernetworks. In Meila, M. and Zhang, T. (eds.), Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pp. 9301-9312. PMLR, 18-24 Jul 2021. URL https://proceedings.mlr.press/v139/sarafian21a.html .

Van Amersfoort, J., Smith, L., Teh, Y. W., and Gal, Y. Uncertainty estimation using a single deep deterministic neural network. In Proceedings of the 37th International Conference on Machine Learning , ICML'20. JMLR.org, 2020.

Villani, C. et al. Optimal transport: old and new , volume 338. Springer, 2009.

Wang, X., Zheng, R., Sun, Y., Jia, R., Wongkamjan, W., Xu, H., and Huang, F. Coplanner: Plan to roll out conservatively but to explore optimistically for model-based rl, 2023. URL https://arxiv.org/abs/2310.07220 .

Weber, R., Schek, H.-J., and Blott, S. A quantitative analysis and performance study for similaritysearch methods in high-dimensional spaces. In Proceedings of the 24rd International Conference on Very Large Data Bases , VLDB '98, pp. 194-205, San Francisco, CA, USA, 1998. Morgan Kaufmann Publishers Inc. ISBN 1558605665.

Welling, M. and Teh, Y. W. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11) , pp. 681-688. Citeseer, 2011.

Wen, Y., Tran, D., and Ba, J. Batchensemble: An alternative approach to efficient ensemble and lifelong learning, 2020. URL https://arxiv.org/abs/2002.06715 .

Wilson, A. G. and Izmailov, P. Bayesian deep learning and a probabilistic perspective of generalization. arXiv preprint arXiv:2002.08791 , 2020.

Xiao, H., Rasul, K., and Vollgraf, R. Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms. arXiv e-prints , art. arXiv:1708.07747, August 2017. doi: 10.48550/ arXiv.1708.07747.

Yang, Y. and Loog, M. Active learning using uncertainty information. In 2016 23rd International Conference on Pattern Recognition (ICPR) , pp. 2646-2651. IEEE, 2016.

Zhmoginov, A., Sandler, M., and Vladymyrov, M. Hypertransformer: Model generation for supervised and semi-supervised few-shot learning. In International Conference on Machine Learning , pp. 27075-27098. PMLR, 2022.

A Minimizing Model Similarity as Particle-based Variational Inference

In Bayesian deep learning, each network in an ensemble can be seen as a particle sampled from a distribution. Hence, the training process can be seen as a variational inference problem in terms of minimizing the KL-divergence between the empirical distribution defined by the particles and the training data. In this section, we relate Eq. (4) with an RKHS and apply it under a particlebased variational inference framework for supervised learning. Particle-based variational inference methods (Liu & Wang, 2016; Chen et al., 2018; Liu et al., 2019) can be viewed from a geometric perspective as approximating the gradient flow line on the Wasserstein space P 2 ( X ) (Liu & Zhu, 2022). Let q t denote the gradient flow line of the KL-divergence w.r.t. some target (data) distribution p ∈ P 2 ( X ) , for an absolutely continuous curve q t , its tangent vector at each t is given by (Villani et al., 2009; Ambrosio et al., 2008),

The core idea of particle-based VI is to represent q t by a set of particles { x i } and adopt a first-order approximation of q t + ϵ through a perturbation of { x i } . In Eq. (7), while the first term corresponds to the maximum (data) likelihood term of supervised learning, the second term is intractable. Different variants of particle-based VI methods tackle this term via different approximation/smoothing methods. Inspired by SVGD (Liu & Wang, 2016), we approximate ∇ log q t in an RKHS corresponding to the HE-CKA kernel.

In particular, let T = { ϕ : X → X} denote the space of transformations on space X of particles, a direction of perturbation can be viewed as a vector field on X , which is a tangent vector in the tangent space T ϕ = id T , where id is the identity transformation. Under the particle representation { x i } ∼ q , let q ϕ ( x ) denote the distribution represented by transformed particles { ϕ ( x i ) } . To approximate ∇ log q , we want to find the perturbation direction of { x i } that corresponds to the steepest ascend direction of the loss J ( ϕ ) = E x ∼ q [log q ϕ ( x )] at ϕ = id, which is the gradient of J ( ϕ ) in the tangent space T ϕ = id T . This gradient is given by the following Lemma, with the proof given in Appendix B.

Lemma A.1. For J ( ϕ ) = E x ∼ q [log q ϕ ( x )]

where K HE-CKA is the HE-CKA kernel defined by Eq. (4). In practice, Eq. 8 can be approximated by the empirical expectation,

In this paper, we apply Eq. (7) and Eq. (9) to the supervised tasks of classification and regression.

B Proof of Lemma A.1

Proof. To compute the gradient of J ( ϕ ) = E x ∼ q [log q ϕ ( x )] at ϕ = id, by definition, we compute as follows the differential of J at ϕ , ∀ v ∈ T ϕ = id T ,

where H is the RKHS with corresponding HE kernel K HE-CKA , and the following identities are used,

C Training Details

C.1 Smoothing Terms and Layer Weighting

To effectively train with the HE-CKA kernel for the repulsive term we found that it is essential to smooth out the particle energy using an ϵ dist on the geodesics and ϵ arc on the cosine similarity values. With larger smoothing terms we can reduce the large gradients on very similar particles, with CKA values near 1 , and ensure other particles still receive some repulsive force. Additionally, Eq. (4) equally weighs every layer in the network. It has been empirically shown that the first few layers of deep neural networks have high similarity (Kornblith et al., 2019), which indicates that initial layers learn more aligned features. Enforcing strong hyperspherical uniformity, or low CKA, of feature Gram matrices may remove useful features. We have noticed that it is difficult to train models with a uniform HE-CKA layer weighting scheme of 1 /L . To fix this we applied a custom weighting scheme w that typically increases linearly with the number of layers, with latter layers weighted higher. We found that using a weighting scheme in Eq. 4 allowed for finer control of the repulsive term. Typically the first layer in a CNN is a simple feature extractor, and depending on the depth of the network could assign too high of a repulsive term on the first layer. Additionally, the last layer could have too high of a weight and ruin inlier performance. We utilize a custom weighting scheme using the vector w = { w 1 , · · · , w L } , where ∥ w ∥ 1 is typically 1. We define the smoothed HE-CKA version for training as HE smooth (Eq. 11).

̸

The smoothing terms gives us finer control over the interaction between particles and prevents exploding or vanishing from the energy term. Although both ϵ dist and ϵ arc have a similar effect it is more important to include the ϵ arc as the gradient of arccos approaches ±∞ near -1 and 1 without any smoothing term. As demonstrated with the cosine similarity feature kernel used in Fig. 7. It is advised to set ϵ dist as a small constant then vary ϵ arc and γ parameters when searching for the right kernel. Optionally, one may replace the Rieszs based kernel with an exponential one, -s arccos( ¯ K m ⊺ ¯ K m ′ / (1 . 0+ ϵ )) -ϵ

ie e l l arc dist , which provides a more numerically stable gradient, and more intuitive to understand growth term s . As discussed in Appendix. D the parameters for γ , β , and w need to be selected. For MNIST experiments we performed a bayes sweep across parameters to select the layer weighting schemes, smoothing terms, and repulsive terms. For larger models, such as the CIFAR and TinyImageNet experiments, we selected, by trying a few combinations, a weighting schemes by testing values γ ∈ [0 . 25 , 1 . 5] , β ∈ [0 . 01 , 10 . 0] , and using layer weighting scheme w l that increases proportionally with l , and with different first layer and last layer values.

By definition of gradient, comparing with Eq. 10,

Earc = 0.010

Earc

= 0.015

Earc = 0.020

Earc

= 0.025

Earc = 0.030

C.2 Dirty-MNIST

The Dirty-MNIST experiments utilized an ensemble of 5 LeNet5 models with a modified variance preserving gelu activation function. Models were trained using AdamW with lr = 0 . 0065 and weight decay of 0 . 001 for 50 epochs, except for Hypernetwork training which was trained for 85 epochs with AdamW with lr = 0 . 0025 and weight decay 0 . 0025 . Details such as lr warmup, gradient clipping, repulsive terms, layer weighting, HE-CKA smoothing terms, and more can be found in the official repository.

C.3 CIFAR-10/100

Additionally, we have some results showing much improvement on CIFAR-100 OOD detection with SVHN when trained with synthetic OOD examples in Table 5. With about a 10% improvement in AUROC between the inlier and outlier sets. We applied HE-CKA to an ensemble of ResNet18 models and evaluated the approach on CIFAR-10 (Table 3) and CIFAR-100 (Table 5). The models were trained for 200 epochs using SGD with a learning rate of 0 . 1 and weight decay 5 e -4 . The HE-CKA kernel used a linear kernel for feature calculation with the exponential kernel s = 2 , and γ = 1 . 0 . For experiments with out-of-distribution (OOD) data, the following values were adjusted: γ = 0 . 5 , γ OOD = 0 . 75 , and β = 0 . 75 . Details regarding layer weighting and smoothing are available in the repository. Forty-eight OOD samples were taken per batch for all CIFAR experiments, where applicable.

The feature repulsion term was not applied to every convolution of the ResNet18 architecture. To conserve computational resources, only a subset of layers was included. Specifically, the selected layers comprised the initial convolutional layer, the output of every other ResNet block within the first two of the four layers, the output of all blocks in the last two layers, and the final linear layer.

Training details regarding the ResNet32 experiments follow the training procedure, learning rate scheduling, and hyperparameters given by D' Angelo & Fortuin (2021). The hypernetwork variant, due to the difficulty of training, was trained for 180 instead of 143 epochs, and utilized group based normalization to stabilize feature variance.

C.4 TinyImageNet and Particle Number Ablation

All models utilized a pretrained deep ensemble without any repulsive term and then fine-tuned using different methods, including our proposed approach. Methods utilizing CKA pw and HE-CKA employed a linear feature kernel. For OOD detection, we used predictive entropy (PE) computed from the ensemble predictions. Additionally, we generated synthetic OOD data from noise and augmented TinyImageNet samples to enhance the OOD detection capability. We utilized a training split of

80:10:10 for training, validation, and testing respectively. Training utilized SGD with a learning rate of 0 . 005 and weight decay of 5 e -4 . We additionally performed a particle number ablation on the ResNet18 + HE-CKA ensemble, utilizing the same repulsive term, showing improvements in accuracy, and outlier detection, when going from 2 particles to 5 (Table. 6).

D Limitations

With our approach, we are able to resolve some of the issues related to tackling permutation of feature channels, which normally pose challenges for Euclidean-based kernels like RBF. However, constructing a model kernel based on layer features requires tuning the repulsive term ( γ ), the likelihood term ( β ), and the layer weighting terms ( w ). This introduces numerous hyperparameters that need to be adjusted depending on the dataset and the architecture in use. Future work could explore automating the estimation of these parameters or simplifying the HE-CKA kernel. Although the assumption that the first few layers should have small repulsive terms seems clear, the weighting and smoothing of later layers remain unclear. This work only explored repulsive terms that increased with layer depth; the dynamics of which layers should have more repulsion are not well understood and have not yet been explored. Additionally, feature-based kernels based on CKA pw are sensitive to the number of particles and the batch size sampled, as the dimensionality of the hypersphere changes, impacting the repulsive terms. One possible solution could be to construct a normalized HE-CKA variant, which precomputes the minimum and maximum energy available on the ( N 2 -1) -sphere with M models.

E Synthetic OOD examples

E.1 OODpoints for 2D datasets

The selection of an OOD set can help force ensemble members to learn unique internal features for outliers, resulting in more diverse output features. For simple datasets it can be easy to construct an OOD set by simply taking min/max + padding as shown in Fig. 8 for the 2D classification tasks.

10

5

1.0

0

0.5

-5

-10

0.0

-10

0

10

E.2 OODImages for MNIST and CIFAR

Images are harder to define boundary/ood points, but we found in practice that generating images by transforming inliers to outliers via typical augmentations, and generating synthetic random channel data worked well in practice. Our approach to transforming inlier points to outlier points consists of the following augmentations in random order: blurring, affine transform, perspective transform, elastic deformations, erasures, Gaussian noise, Gaussian blurring and inversions. Examples of such images are shown in Fig. 9 for MNIST, Fig. 10 for CIFAR, and Fig. 11 for TinyImageNet. The ID to OOD set of images accounts for roughly 30-40% of the dataset, whereas the other 60% are randomly generated. Synthetic OOD images are generated via combinations of perlin noise, simplex noise, gaussian noise, lines, alternating grids, inversions, and random area thresholding. For images with more than one channel, such as CIFAR and TinyImageNet, we either apply different noise to each channel, use the same method but different seed, or occasionally broadcast one channel along all channels.

F MNIST CKA Results

We present the CKA plots of each method: Ensemble (Fig. 12), Hypernetwork (Fig. 13), SVGD + RBF (Fig. 14), SVGD + CKA pw (Fig. 15), SVGD + HE-CKA (Fig. 16), Ensemble + HE-CKA (Fig. 17), Ensemble + OOD HE-CKA (Fig. 18), and Hypernetwork + OOD HE-CKA (Fig. 19). Each plot shows a grid comparing layerwise estimation of pairwise CKA, equipped with a linear feature kernel on the inlier Dirty-MNIST dataset, while using an unbiased estimator for CKA. To our surprise we found that CKA pw and HE-CKA results in fairly similar unbiased CKA estimates across the ensemble, but overall performance of the models in uncertainty estimation and accuracy presented in Table. 1 by HE-CKA were better. Nevertheless, methods utilizing CKA pw and HE-CKA kernels significantly reduce similarity of features compared to methods with no repulsive terms or RBF based kernels. This was especially true for our hypernetwork tests.

G Memory Footprint and Time Complexity

We compare the training runtime of our HE-CKA term to ParVI based methods presented in D' Angelo & Fortuin (2021). We evaluated mini-batch training time averaged over 50 batches on a Quadro RTX 8000. Each method used a ResNet18 fed with batches of 128 images from CIFAR-10. The results are presented in Table 8. Reported CUDA memory includes all ensemble members, loss, batch statistics, feature kernels (if applicable), and gradients. We see that all ParVI methods increase training time by 2.2x for 5 ensemble members and 1.2x for 10. Given that HE-CKA is applied layerwise, our method does require a slight increase in memory compared to the other ParVI methods, but is comparable to other methods in terms of training batch time increase. The time complexity for each minibatch of HE is O ( LN 2 n 2 ) compared to typical function space O ( N 2 n 2 ) or weight space O ( n 2 ) kernels, where n is the number of particles, N is the mini-batch size and L is the number of layers we use to compute HE-CKA . However, n and N are typically small, 5 and up to 128 respectively in our experiments, and one can use a subset of the layers L , as we did with our experiments. While using the HE-CKA kernel does have an increased memory and computation cost than using the RBF kernel, the benefit of having a kernel invariant to feature permutations and scaling are worth the minor additional cost.

0

10

20

0

10

20

0

10

20

0

10

20

10

20

0

10

20

0

10

20

0

20

0

20

0

200

20

0

200

20

0

20

0

20

0

20

0

20

0

20

0

20

0

20

0

20

0

20

0

20

0

20

0

20

0

20

0

25

50

0

25

50

0

25

50

0

25

50

0

25

50

0

25

50

0

25

50

0

50

0

500

50

0

500

50

0

50

1.0

4

0.59

0.66

0.77

0.83

0.89

0.8

3

0.77

0.82

0.89

0.90

0.83

0.6

ayer

2

0.86

0.90

0.93

0.89

0.77

0.4

1

0.93

0.96

0.90

0.82

0.66

0.2

0

0.99

0.93

0.86

0.77

0.59

0.0

0

2

4

Layer

1.0

4

0.58

0.67

0.80

0.90

1.00

0.8

3

0.64

0.71

0.88

1.00

0.90

0.6

Layer

2

0.80

0.86

1.00

0.88

0.80

0.4

1

0.91

1.00

0.86

0.71

0.67

0.2

0

1.00

0.91

0.80

0.64

0.58

0.0

0

2

4

Layer

1.0

4

0.57

0.61

0.66

0.62

0.73

0.8

3

0.72

0.74

0.76

0.68

0.62

0.6

Layer

2

0.84

0.87

0.87

0.76

0.66

0.4

1

0.91

0.94

0.87

0.74

0.61

0.2

0

0.98

0.91

0.84

0.72

0.57

0.0

0

2

4

Layer

1.0

4

0.39

0.27

0.14

0.12

0.12

0.8

3

0.40

0.28

0.15

0.13

0.12

0.6

Layer

2

0.45

0.31

0.17

0.15

0.14

0.4

1

0.65

0.45

0.31

0.28

0.27

0.2

0

0.96

0.65

0.45

0.40

0.39

0.0

0

2

4

Layer

1.0

4

0.41

0.26

0.13

0.10

0.10

0.8

3

0.42

0.27

0.13

0.11

0.10

0.6

Layer

2

0.46

0.30

0.15

0.13

0.13

0.4

1

0.66

0.47

0.30

0.27

0.26

0.2

0

0.97

0.66

0.46

0.42

0.41

0.0

0

2

4

Layer

1.0

4

0.35

0.22

0.13

0.11

0.10

0.8

3

0.36

0.22

0.14

0.11

0.11

0.6

Layer

2

0.39

0.24

0.15

0.14

0.13

0.4

1

0.58

0.38

0.24

0.22

0.22

0.2

0

0.81

0.58

0.39

0.36

0.35

0.0

0

2

4

Layer

1.0

4

0.45

0.46

0.23

0.17

0.15

0.8

3

0.48

0.48

0.24

0.17

0.17

0.6

Layer

2

0.56

0.56

0.29

0.24

0.23

0.4

1

0.78

0.73

0.56

0.48

0.46

0.2

0

0.98

0.78

0.56

0.48

0.45

0.0

0

2

4

Layer

1.0

4

0.18

0.15

0.11

0.07

0.07

0.8

3

0.18

0.16

0.12

0.07

0.07

0.6

Layer

2

0.24

0.22

0.17

0.12

0.11

0.4

1

0.55

0.43

0.22

0.16

0.15

0.2

0

0.88

0.55

0.24

0.18

0.18

0.0

0

2

4

Layer

H NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: See abstract and introduction.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Limitations on training time and memory consumption was shown in Table 8. This paper has no ethical limitations. Additional limitations of our approach are discussed in Appendix. D.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: See Appendix A and B.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The proposed algorithm is completely defined and the network structures are either defined/referenced or are standard.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The code is publicly available, and contains the relevant instructions at the following public repository https://github.com/Deep-Machine-Vision/ he-cka-ensembles .

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The experimental setups are described clearly. Training details and data splits are described in Appendix. C.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Standard deviations are provided whenever necessary by multiple runs.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: See Appendix. G for details regarding memory footprint, and a comparison to other ParVI methods.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We conducted the research ethically.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This paper does not have significant social impact to be discussed.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This paper does not have a high risk for misuse.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have cited all the appropriate sources.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: We didn't release new assets.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: We did not use crowdsourcing or human subjects.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: We did not conduct human subject research.

Peer reviews:
comments: SUMMARY:
This paper introduces the hyperspherical energy as an objective to prompt the diversity of particles in BNNs. It claims that the hyperspherical energy approach can avoid permutation invariance in traditional diversity metrics.

STRENGTHS:
I agree that the diversity of particles in BNNs is an important problem. The proposed method introduces hyperspherical energy, which can avoid permutation invariance. This invariance is mainly due to the complex structures of DNNs, which previous Bayesian Inference methods have long ignored.

WEAKNESSES:
The technical contribution of this paper is not strong enough. Only a new regularization term is added to the standard ensemble framework. The idea of hyperspherical energy comes from previous works, and I do not recognize significant changes from previous works.

QUESTIONS:
1. The diversity evaluations are limited to synthetic data. I am not convinced that the proposed method can be scaled up.

2. Important ablation studies are missing. For example, different model architectures and the number of particles.

LIMITATIONS:
1. This paper is restricted to relatively small Bayesian neural networks, and does not discuss the scalability.

2. The proposed method can improve uncertainty estimation, but fails to improve accuracy.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-10

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper “Minimizing Hyperspherical Energy for Diverse Deep Ensembling” explores the use of Centered Kernel Alignment (CKA) and Minimization of Hyperspherical Energy (MHE) in Bayesian deep learning to enhance the diversity of ensemble models and parameters generated by hypernetworks. By incorporating these techniques, the authors aim to improve uncertainty quantification in both synthetic and real-world tasks, which is quantified by measuring OOD detection performance and calibration. The key contributions include proposing CKA as an optimization objective and utilizing MHE to address the diminishing gradient issue, leading to more stable training and better performance in uncertainty estimation.\n\nSTRENGTHS:\n- **Original:** The paper introduces a novel approach by combining CKA and MHE to enhance the diversity of deep learning ensembles and hypernetworks, which is a fresh perspective in Bayesian deep learning.\n- **Detailed experiments:** The experimental results are comprehensive and demonstrate significant improvements in uncertainty quantification across various tasks, showing the practical effectiveness of the proposed methods.\n- **Well-written:** The paper is well-structured, with clear explanations of the methods and thorough discussions of the results. The figures and tables effectively illustrate the performance improvements.\n\nWEAKNESSES:\n- The comparison of the approach in terms of OOD detection performance is slightly biased due to leveraging of generated OOD samples. Here in order to differentiate the contribution of the method from the incorporation of additional information a further experiments comparing against DDU would be warranted. In particular one could imagine adding a term to the loss which encourages low likelihood of the fitted GMM on OOD samples.\n - The comparisons in Figure 2 paint a slightly overly optimistic picture for the OOD HE approaches. In 2 dimensions the problem becomes quite trivial given negative / OOD samples. The difficulty of generating enough samples to cover the OOD volume becomes exponentially more difficult with increasing dimensionality of the input space and is especially easy when the diversity of the data is limited.\n - (Related to prior point) the datasets used to show the efficacy of the OOD approach are rather simple and small such that sample diversity is limited. Running one experiment on a larger dataset like imagenet would be very beneficial to show that the method can also work in the context of highly diverse input distributions and high dimensionality of the input space.\n - There seems to be a mistake in table 2 row Ensemble + HE, column PE\n - I recommend renaming Fig1 d and e to Cossim evaluation and HE evaluation as using objective gives the impression that the data shown is from models trained using these objectives.\n - As the work points out that applying the method to larger datasets would in principle be possible, this might be a good additional comparison to include. Many comparison approaches will have issues in this case, but comparison to other ensembling based approaches would be feasible.\n\nQUESTIONS:\n- When running the LeNet comparisons agains DDU, was the model trained with spectral normalization (a requirement for DDU to function well) and how was the regularization strength determined?\n- Why was the normalization-free variant of LeNet used for some methods whereas the standard variant for the DDU baseline?\n- Page 9, line 289: Was the GMM fitted to the features of the pre-classification layer. I.e. to the features used by the last linear layer?\n\nLIMITATIONS:\nThe work correctly points out that the small size of the datasets is one of the main limitations. Generally, seeing the efficacy of the approach in a larger scale training setup would significantly improve the paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

comments: SUMMARY:
The authors proposed to improve the quantification of particle diversity in deep ensemble with hyperspherical energy (HE) on top of the CKA kernel. They further integrate the HE kernel in particle-based variational inference (ParVI) and generative ensemble with hypernetwork frameworks. The methods are evaluated on both synthetic experiments and small-scale classification datasets.

STRENGTHS:
1. The motivation of the paper is clear: addressing mode collapse in deep ensemble by using a kernel that is more suitable to measure particle diversity of neural networks (HE kernel).
2. The advantage of HE kernel has been demonstrated in two different ensemble frameworks (parVI & generative ensemble with hypernetwork), and the empiricall performance of the method looks Ok.

WEAKNESSES:
The datasets considered are a bit outdated and the networks considred seem to be quite small, such that overall the performance is on the lower end (e.g. 85% acc. for CIFAR10, while many BDL methods can easily achieve accuracy higher than 90%). Fuerthermore, recent BDL papers typically consider larger datasets (such as ImageNet) and deeper networks. It is necessary to consider larger datasets and larger models in order to assess the practical effectiveness of the method.

QUESTIONS:
See weakness

LIMITATIONS:
The authors do not address potential negative social impact since the paper is predominantly theoretical.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-14