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

Authors: david  smerkous
qinxun  bai
li  fuxin
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 01:21:44

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.

Main paper: 
 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 38th Conference on Neural Information Processing Systems (NeurIPS 2024). 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. 
related work: 
 Uncertainty Estimation. A large 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 (Beck et al., , 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. 
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 asf (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.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) ,EQUATION)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) .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 m l } m=1,...,M l=1,...,L . The mean pairwise loss across all layers L is as follows,EQUATIONIn its current form, CKA pw 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 Km = 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 ′ ) = Km⊤ Km ′. 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,CKA pw (K) = 1 LM (M -1) L l=1 1 ⊺ zd(K l K ⊺ l )1 s.t K l =    Kl 1 . . . Kl M    ∈ R M ×N 2 (3)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.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 layer l we treat the M model Gram vectors Km l as particles on the hypersphere, its geodesic on the hypersphere is thend i,j = arccos(CKA(K i l , K j l )) = arccos( Ki⊺ l Kj 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.EQUATIONwhere 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. 
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 withL layers f (x, θ) = f L (f (...) (f 1 (x, θ 1 ), • • • ), θ L ).We denote the ensemble of target network layer at layer l asE 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: The MAP estimate of the posterior in Eq. ( 5) results in the following objective:EQUATIONEQUATIONThe 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.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,csize , 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 .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 li corresponds to a latent vector for a single convolution filter i. We feed each filter code i through a filter generator g l (c li ) separately to generate the filter for layer l. An example architecture can be seen in Fig. 4 .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.  
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.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.2x) (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.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. 
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. 
 
 Appendix: 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) , grad KL(q t |p) = -∇ log p + ∇ log q t . (7) 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.EQUATIONwhere K HE-CKA is the HE-CKA kernel defined by Eq. ( 4). In practice, Eq. 8 can be approximated by the empirical expectation,EQUATIONIn this paper, we apply Eq. ( 7) and Eq. ( 9) to the supervised tasks of classification and regression.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 ,EQUATION)where H is the RKHS with corresponding HE kernel K HE-CKA , and the following identities are used,q ϕ (x) = q(x) det ∂ϕ ∂x , d log | det A| = Tr(A -1 dA), ∂v i (x) ∂x j = ⟨∇ x j K HE-CKA (x, •), v i (x)⟩ H .By definition of gradient, dJ ϕ (v) ϕ=id = ⟨∇ ϕ J ϕ=id , v⟩ H , comparing with Eq. 10,∇ ϕ J (ϕ)(•) ϕ=id = E x∼q [∇ x K HE-CKA (x, •)] .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) .EQUATIONThe 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 Riesz-s based kernel with an exponential one, ie e -s arccos( Km⊺ l Km ′ l /(1.0+ϵ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. 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 5e-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. 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.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. 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.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.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.  


Decision: In this work, the authors propose a novel approach to increase feature diversity in deep ensembles by minimizing hyperspherical energy (MHE) on top of Centered Kernel Alignment (CKA) kernels. The reviewers praise the work for its originality, well-written presentation, and comprehensive experiments demonstrating significant improvements in uncertainty quantification. However, they are critical of the limited evaluation on small datasets and the lack of ablation studies on particle numbers and model architectures. In their rebuttal, the authors clarify their contributions, address concerns about dataset size and model performance, and provide additional experiments on TinyImageNet. This addressed the reviewers' concerns about dataset size and model performance but did not fully address their concerns about the technical contribution and scalability.

After the rebuttal and discussion, the paper is very borderline, where two reviewers lean towards acceptance, while one leans towards rejection. The main criticism on which the decision hinges seems to be the small scale of the experiments and the lack of ablation studies. Both of these have been provided by the authors in their rebuttal.  I trust that the authors add those additional results to the camera-ready version and might even add more ablation studies that have not yet been provided, which would further strengthen the paper. Based on this, I recommend acceptance. However, I would further recommend that the authors seriously take all the reviewer feedback into account for the camera-ready version.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: See weakness
limitations: limitations: The authors do not address potential negative social impact since the paper is predominantly theoretical.
paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: - 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.
 - 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.
 - (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.
 - There seems to be a mistake in table 2 row Ensemble + HE, column PE
 - 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.
 - 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.
questions: questions: - 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?
- Why was the normalization-free variant of LeNet used for some methods whereas the standard variant for the DDU baseline?
- 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?
limitations: limitations: The 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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: 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.

Comments: comment: ## TinyImageNet Correction
We discovered a bug in our TinyImageNet dataloader for the HE + OOD experiment during the rebuttal period, which skewed our results for that one experiment. After fixing it and retraining, we obtained the following updated results.

| Model | NLL $\downarrow$ | ID Accuracy $\uparrow$ | ID ECE $\downarrow$ | AUROC SVHN $\uparrow$ | AUROC CIFAR-10/CIFAR-100 $\uparrow$ | AUROC Textures (DTD) $\uparrow$ |
| -------- | -------- | -------- | -------- | -------- | -------- | -------- | 
| ResNet18 + $\text{HE}$ OOD (5)   | 0.786     | 61.88     | 8.02     | **99.31**     | **81.56**/**87.64**     | **90.94** |
We thank the reviewer for the recognition of our experimental results and the raised rating. We also appreciate the continuing discussion. 

We have discussed in our rebuttal the nontrivial changes and key properties of our proposed regularization method versus previous regularization techniques in ensemble training. As recognized by other reviewers, the combination of $\text{CKA}$ and $\text{HE}$ is a novel approach and a fresh perspective, and most importantly, it leads to improvements in practice for a long-standing important problem in ensemble learning, i.e., enabling better ensemble diversity (not attainable by merely switching channels within a network) which leads to significantly improved OOD performance. 

We disagree with the reviewer's attempt to dismiss any paper that works on regularization of networks to be not novel, without actually judging the contribution of the paper. Besides, our novel $\text{CKA}$+$\text{HE}$ kernel also allowed us to train for increased diversity on the synthetically generated outlier images which further significantly improved the OOD performance in all cases. To the best of our knowledge, this has not been explored in existing literature of ParVI methods.
Thank you for pointing out the missing response to your comments on the weaknesses. We somehow forgot to include them and will post the answers here in addition to the experiments on the larger dataset.

> Fair comparison against DDU.

Instead of adding OOD examples to DDU as you suggested, which is also viable, in Tables (Paper 1-4. Rebuttal Table 1), we showed that even without OOD samples, our method/kernel still outperforms ensembles and DDU (Tables 1/2) in OOD detection, which are the apples-to-apples comparisons. Although, we can mark that row more clearly to indicate synthetic OOD sample usage.

> In 2D with OOD examples the problem becomes quite trivial. 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.

We agree that adding OOD to synthetic experiments is easy, however as mentioned earlier our methods outperform baselines even without OOD examples. For images, our main point is to show that we do not have to cover the entire OOD volume but just have to generate some random images to make sure the diversity is high on random outlier images. Our OOD detection experiments show that although our generated synthetic images (Fig. 7/8) are nowhere similar to the tested outliers, our approach still performs well for outlier detection.

> Larger datasets

This one is answered in the general remarks.

> Table 2 and Figure 1:

Thank you for pointing these out. We will correct the extra zero in the table, and further clarify the cossim/he figure name.
Thanks for answering the questions. I would have appreciated some comments on the weaknesses of the work I pointed out. Will leave my score unchanged as solely the issue of dataset size was addressed and other points remained unmentioned.
Thanks for the rebuttal, which addresses my concern about the model perfoemance. I am willing to raise my score to 6.
For the question you asked about larger dataset and larger models, results can also be found in the general remarks.
I acknowledge that the evaluations do cover MNIST/CIFAR10/CIFAR100. I am willing to raise my score to reflect that, along with the new experiments added in the rebuttal.

However, I still do not think the technical contributions mentioned in the rebuttal are convincing. The rebuttal only discusses the meaning of this regulation term.