Title: Model and Feature Diversity for Bayesian Neural Networks in Mutual Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: Bayesian Neural Networks

Deep Mutual Learning
Abstract: Bayesian Neural Networks (BNNs) offer probability distributions for model parameters, enabling uncertainty quantification in predictions. However, they often underperform compared to deterministic neural networks. Utilizing mutual learning can effectively enhance the performance of peer BNNs. In this paper, we propose a novel approach to improve BNNs performance through deep mutual learning. The proposed approaches aim to increase diversity in both network parameter distributions and feature distributions, promoting peer networks to acquire distinct features that capture different characteristics of the input, which enhances the effectiveness of mutual learning. Experimental results demonstrate significant improvements in the classification accuracy, negative log-likelihood, and expected calibration error when compared to traditional mutual learning for BNNs.

Model and Feature Diversity for Bayesian Neural Networks in Mutual Learning

Abstract

Bayesian Neural Networks (BNNs) offer probability distributions for model parameters, enabling uncertainty quantification in predictions. However, they often underperform compared to deterministic neural networks. Utilizing mutual learning can effectively enhance the performance of peer BNNs. In this paper, we propose a novel approach to improve BNNs performance through deep mutual learning. The proposed approaches aim to increase diversity in both network parameter distributions and feature distributions, promoting peer networks to acquire distinct features that capture different characteristics of the input, which enhances the effectiveness of mutual learning. Experimental results demonstrate significant improvements in the classification accuracy, negative log-likelihood, and expected calibration error when compared to traditional mutual learning for BNNs.

1 Introduction

Bayesian neural networks [22] (BNNs) provide a probability distribution over the model parameters, which allows for uncertainty quantification in the model's predictions. This is especially useful in situations where making accurate predictions is crucial, and having a measure of uncertainty can inform decision-making. This measurement is essential for many applications, including continual learning, active learning, and robust operation. To measure uncertainty, we can learn the posterior distribution after observing data through methods like Variational Inference (VI) [1, 16] and Markov Chain Monte Carlo (MCMC) [37, 21, 3]. The variational approach for BNNs involves selecting prior distributions and approximate posterior distributions over neural network weights, offering a more memory-efficient and flexible approach compared to MCMC.

However, despite their advantages, Bayesian neural networks using variational inference often find it difficult to match the performance of BNNs using the MCMC approach, as well as deterministic neural networks. To address this performance gap, various methods have been explored [26, 28, 31, 8, 24, 14]. One promising method is knowledge distillation [13, 27, 17, 25, 12], which involves training a smaller network (student) to mimic the behaviour of a larger, more complex network (teacher) to achieve comparable performance with fewer computational resources. This idea has been investigated for compressing BNNs [18, 36, 30]. Although it is efficient to train the student in this manner, selecting the right teacher model which can guide the student during training may be difficult to a certain extent. In [18, 36, 30], the authors use teacher models, which are BNNs trained with Stochastic Gradient Langevin Dynamics (SGLD), resulting in prohibitive storage costs. Additionally, this knowledge

distillation requires a multi-stage offline process involving pre-training the teacher, which could potentially limit the mutual learning between peer BNNs.

Mutual learning or online knowledge distillation [39, 40, 15, 4, 38] has received increasing interest in recent years. Mutual learning simplifies the training process of knowledge distillation to a single stage, which allows two or more networks to learn collaboratively and teach each other simultaneously, even with identical model architectures. In [39], the authors simply encourage multiple networks to learn from each other by reducing Kullback-Leibler (KL) divergence between pairs of predictions. In the same line of research, [40] proposes to train a multi-branch architecture by sharing low-level layers and concurrently creating an effective teacher based on soft logits to enhance the learning of the target network. Meanwhile, in [38, 4], the authors match the distribution of feature maps through adversarial training. In [4], the authors also indicate that a direct alignment method, such as the L 1 or L 2 distance between features of peer models could make networks become identical. Consequently, this does not improve the performance of peer networks in mutual learning. Even though these approaches show effective results, they mainly concentrate on the soft logits or features of peer networks, with limited research exploring the impact of diversity in the model parameter space for mutual learning. In the context of deep mutual learning, increasing diversity in parameter distributions and feature representations could encourage peer networks to learn distinct perspectives of data from one another. This process would enable models to acquire richer information, leading to more accurate decision-making.

In this work, we propose a novel approach to enhance the performance of BNNs by leveraging deep mutual learning. Specifically, we aim to increase diversity for deep mutual learning applied to BNNs. First, we propose increasing diversity in model parameter space by forcing posterior distributions to diversify. Next, we improve the BNNs' performance in the context of mutual learning by diversifying feature distributions. By promoting diversity in terms of model parameters and feature distributions we encourage peer networks to learn distinct features and characteristics from one another, ultimately boosting their overall performance. Our experimental results demonstrate that increasing diversity in both model parameter distributions and feature distributions is crucial for enhancing the performance of peer BNNs. We conduct experiments on the CIFAR-10, CIFAR-100, and ImageNet datasets and observe substantial improvements in classification accuracy, negative log-likelihood, and expected calibration error compared to the traditional mutual learning for BNNs. Our contributions can be summarized as follows:

We propose a novel approach for improving the accuracy of BNNs by leveraging the effectiveness of mutual learning. Specifically, we propose promoting diversity in both model parameter space and feature space during the mutual training of peer BNNs. To our best knowledge, this is the first work that leverages mutual learning in the context of BNNs, and also investigates the usefulness of model parameter diversity in mutual learning. Additionally, we are the first to propose maximizing the distance between feature distributions to promote diversity in mutual learning for BNNs.

We extensively evaluate our approach on CIFAR-10, CIFAR-100, and ImageNet datasets and demonstrate that it outperforms traditional mutual learning for BNNs in terms of accuracy, negative log-likelihood, and expected calibration error.

2 Related work

Bayesian Neural Networks have emerged as an active area of research due to their ability to capture uncertainty in predictions. While traditional deep learning models represent model parameters by point estimation and tend to be overconfident in their predictions [9], BNNs focus on estimating the posterior distribution over the weight parameters of neural networks. This approach can result in more calibrated predictions and better uncertainty estimations [1]. However, BNNs present several challenges, including intractable posterior distributions, unclear prior distributions, and a large number of parameters. To address these issues, researchers have proposed various methods for BNNs such as variational inference [1, 16, 5, 23], and Markov Chain Monte Carlo (MCMC) [22]. MCMC often results in a good performance, but this approach requires prohibitive computational cost and memory storage, making it impractical for real-world applications. In contrast, variational inference or mean-field variational inference (MFVI) presents a more convenient alternative by approximating the complicated posterior distribution with a simpler variational approximation. For instance, this

assumes a Gaussian posterior with diagonal covariance matrix [1]. Variational methods require storing only one model for evaluation, which is more scalable and useful in downstream tasks.

Variational inference for Bayesian neural networks. Given training dataset S = { ( x n , y n ) } N n =1 with x n ∈ X being an input and y n ∈ Y being the corresponding output. The aim is to find a model f ( . ; w ) : X → Y , parameterized by w ∈ W , that can accurately predict the output y of each instance x . In Bayesian machine learning, that aim is equivalent to finding the posterior p ( w |S ) of w given the training dataset S with a prior p ( w ) through the Bayes' theorem:

The denominator in Eq. (1) is, however, intractable to be calculated exactly except for some simple models. To tackle this issue, variational inference is introduced to approximate the true posterior p ( w |S ) by a variational posterior q ( w ; θ ) parameterized by θ . The variational posterior q ( w ; θ ) can be obtained by minimizing the following KL divergence:

where p ( w ) is the prior of the parameter w .

Bayes by Backprop (BBB) [1] is a simple but efficient variational approach. They suppose the approximate posterior following a Gaussian distribution with a diagonal covariance matrix, i.e., q ( w ; θ ) = N ( µ, diag( σ 2 )) and apply a parametrization trick to represent each weight in the neural network by a mean and a standard deviation:

where: ⊙ is the Hadamard product and I is the identity matrix.

However, the standard BBB method in practice is sensitive to hyperparameters. To address this issue, in [5], the authors introduce a variant of BBB called Radial Bayesian Neural Network, where weight parameters are sampled as follows:

where ∥ . ∥ denotes the vector norm.

Deep Mutual Learning (DML) , also known as online knowledge distillation, has gained considerable interest recently due to its adaptability in various applications. Unlike offline knowledge which requires selecting and training a teacher network in advance, mutual learning involves distilling and sharing knowledge among peer networks. In [39], the authors argue that multiple models, even with identical architecture, can support each other during the learning process by minimizing the KL divergence between their soft logits. Regarding mutual learning in logits, ONE [40] generates an importance score for each student through the gate module, resulting in ensemble logits used as alignment targets for each network. In [10], optimal ensemble logits are computed and distilled among networks. However, these methods only utilize output logits without considering the potential benefits of incorporating knowledge from intermediate feature maps in peer networks. In [4], the authors show that directly minimizing the distance between features using L 1 or L 2 severely degrades the accuracy of peer networks, as it may cause models to learn in identical ways. Therefore for each network, in addition to learning useful features for the task on its own, [4] also encourages each network to learn the feature distribution of its peer network through an adversarial training strategy. By doing this, they encourage both networks to learn features that generalize better. In the other words, they encourage the two networks to learn features that work well for both, ultimately resulting in better performance. Despite these advancements, existing methods in mutual learning only use output logits or intermediate feature maps in peer networks and do not take advantage of knowledge from parameter space in peer networks, which could further promote diversity and play a crucial role in mutual learning.

Bayesian Deep Mutual Learning. Several works [36, 18, 30] adopt knowledge distillation methods to enhance the efficiency of BNNs. In [18], the authors train a teacher using SGLD [37] and then distil knowledge into a student network, which is a deterministic model. In [36], the authors attempt to reduce the storage issues of MCMC methods by distilling the posterior distribution of the BNN using the Generative adversarial networks (GANs). However, this approach could be difficult for larger models with millions of parameters. With this regard, in [30], the authors treat the knowledge in the teacher as prior to constraining the variational objective of the student. However, the aforementioned works still require getting MCMC samples in advance for distilling knowledge into variational BNNs, leading to expensive computational resources and large storage. This could be unsuitable for training large BNNs. In our work, we propose to apply mutual learning to improve the performance of variational Bayesian neural networks to make them applicable in real-work tasks. Specifically, to diversify features of peer BNNs, we propose to enlarge the distance between feature distributions of peer networks which is achieved by a careful design of feature combination, an appropriate choice of feature distribution approximation and a distance function. More importantly, in addition to diversity in feature space, we also investigate and leverage the optimal transport for encouraging diversity in model parameter space. To our best knowledge, this is the first work that investigates the usefulness of model parameter distribution diversity and also the first to enlarge the distance between feature distributions in mutual learning.

3 Proposed method

3.1 Deep Mutual Learning for Bayesian Neural Networks

We formulate our proposed method by considering two peer Bayesian neural networks, B 1 and B 2 , with the goal of constructing an efficient framework for BNNs in mutual learning. The core idea of mutual learning [39] allows peer models to provide feedback on their learning to each other. This encourages peer models to acquire additional information from one another, ultimately leading to improved generalization performance. We adopt the two loss terms for logit-based learning; one is the conventional loss for variational Bayesian neural networks, and the other is mutual distillation loss between peer networks using KL divergence. In our setting, the approximate posterior distributions for each peer BNN is assumed to be a Gaussian distribution with a diagonal covariance matrix:

where µ θ i , σ θ i ∈ R d , with d denoting the number of model parameters.

The conventional loss for variational Bayesian neural network B i , i = 1 : 2 is expressed as:

Let z B 1 and z B 2 be the logit outputs of the Bayesian neural networks B 1 and B 2 , T be the temperature parameter to control the smoothness of the output distribution. We denote soft logit p B 1 = softmax( z B 1 / T ) , and p B 2 = softmax( z B 2 / T ) . The logit loss terms are defined as follows:

3.2 Diversity in model parameter space

We aim to enhance diversity by leveraging the benefits of Bayesian neural networks over deterministic neural networks, in which weights in a BNN are represented by a distribution. Specifically, we increase the distance of posterior distributions and subsequently enhance the diversity in mutual learning. To the best of our knowledge, this is the first work that exploits the model parameter space to enhance diversity in mutual learning.

During the training process, for each iteration, we obtain Gaussian distributions with parameters θ 1 = ( µ θ 1 , Σ θ 1 ) and θ 2 = ( µ θ 2 , Σ θ 2 ) for BNN models B 1 and B 2 , respectively. We then measure the distance between two posterior distributions D ( q ( w ; θ 1 ) , q ( w ; θ 2 )) . To this end, we exploit two well-known methods: Kullback-Liebler (KL) divergence and optimal transport. Both of them have

closed-form for Gaussian densities. While the KL divergence is well-known, it has limitations, such as being asymmetric and potentially infinite. Optimal transport theory [35] offers an alternative with the Wasserstein distance, which is symmetric and satisfies the triangle inequality. Furthermore, since the approximate posterior of the Bayesian neural network for the MFVI approach is Gaussian with diagonal covariance, this results in a simple form of optimal transport. Thus, we adopt the optimal transport, or Wasserstein distance, to measure the distance between two Gaussian distributions [34]:

where B ( ., . ) is the Bures distance [2] represented by:

With diagonal covariances Σ θ 1 = diag( σ 2 θ 1 ) , and Σ θ 2 = diag( σ 2 θ 2 ) , we have a simple form of distance between two approximate posteriors:

To promote diversity in the parameter space, a straightforward approach would be to add a negative D ( q ( w ; θ 1 ) , q ( w ; θ 2 )) term to the objective functions in Eqs. (7) and (8) for minimization. However, doing so could lead to rapid increases of D ( q ( w ; θ 1 ) , q ( w ; θ 2 )) which would impair the training process. Therefore, we propose minimizing the following loss function:

The above loss function exhibits desirable properties: L diverse param ≥ 0 , and minimizing L diverse param implies maximizing D ( q ( w ; θ 1 ) , q ( w ; θ 2 )) , thus enhancing diversity in the model parameter space. These properties make optimization more stable during training.

3.3 Diversity in intermediate feature space

We aim to leverage intermediate features to increase the distance between the distributions of features among peer networks, rather than directly matching features using L 2 distance as in offline distillation methods such as FitNets [27] and Attention transfer [17]. This is because directly matching features can lead to peer networks becoming identical, thereby reducing diversity in mutual learning. In our framework, the backbone networks are divided into the same blocks according to their depth. Let F B 1 k and F B 2 k be the extracted features at the k th block of the BNN models, B 1 and B 2 , respectively. To increase the diversity of features among peer BNN models, we aim to maximize the distance between their feature distributions. It is worth noting that, diversifying peer networks' features is not a straightforward task. Our experimental results (detailed in the Supplementary Material) show that directly maximizing the L 1 or L 2 distance between peer networks' features does not necessarily enhance the effectiveness of mutual learning. Similarly, we observed the same outcome when directly maximizing the distance between peer networks' feature distributions. The reason could be that directly maximizing the distance between the distributions of original features F B 1 k and F B 2 k could potentially keep pushing the features apart indefinitely. This leads to BNN models becoming too dissimilar to effectively learn from each other, and then in turn negatively affect the predictive abilities of the BNN models. To mitigate this, we propose to diversify fused feature distributions. By doing so, we indirectly diversify feature distributions through transformed features, which allows for more flexible feature adaptation. Additionally, the use of fused features enables to increase diversity across multiple feature levels. We leverage cross-attention [33] to combine features from different blocks of each BNN. Specifically, given extracted features from the k th block, and ( k +1) th block of the network B 1 , the combined feature using cross-attention can be expressed as:

where W Q , W K , and W V represent learnable parameters for query, key, and value, respectively. Similarly, we have combined features F ′ B 2 = u ( F B 2 k , F B 2 k +1 ) . Specifically, in a mini-batch with n samples, we denote G B 1 = { F ′ j B 1 } n j =1 , and G B 2 = { F ′ j B 2 } n j =1 represented fused features derived from BNN models B 1 and B 2 , respectively.

Similar to the measurement of distance in the parameter space, we can utilize optimal transport. However the feature distributions may not be clearly represented by Gaussian distributions and

other forms of optimal transport often require the computation of cost matrices, which are often computationally expensive. Therefore, we adopt KL divergence to diversify feature distributions. To this end, we generate the distribution of fused intermediate features by leveraging the conditional probability density of any two data points within the feature space [32], which is expressed as:

̸

where K ( a, b ) is a kernel function. In our work, we adopt the similarity metric K ( a, b ) = 1 2 ( a T b ∥ a ∥∥ b ∥ + 1) which has been used in [25] for offline knowledge distillation.

Given that Q and P are the probability distributions of fused features from BNN models B 1 and B 2 , respectively, the distance between fused feature distributions is represented as:

It is worth noting that feature distribution has been investigated for offline knowledge distillation, i.e., in [25], the authors minimize KL divergence to transfer the probability distribution from teacher to student, rather than transferring actual representations. Conversely, in the context of mutual learning, our approach seeks to maximize the KL divergence between fused feature distributions to promote feature diversity. Similar to promoting diversity in the parameter space, we propose minimizing the following loss functions for BNN models B 1 , and B 2 :

Finally, the overall loss of our proposed Bayesian mutual learning for Bayesian neural networks B 1 and B 2 are defined as:

where α , and β are the hyper-parameters of total loss for optimization.

In our proposed BNNs in mutual learning strategy, the peer networks are jointly optimized and collaboratively learned during training. The optimization details are presented in Algorithm 1.

4 Experiments

4.1 Experimental setup

Datasets. We evaluate our approach on widely used datasets, including CIFAR-10 [20], CIFAR-100 [19], and ImageNet [29]. The CIFAR-100 dataset consists of 60,000 images across 100 classes, with 50,000 and 10,000 images used for training and validation, respectively. Meanwhile, the CIFAR-10 dataset has the same number of images for both training and validation sets but only consists of 10 classes. Additionally, we conduct experiments on ImageNet - a large and challenging dataset with 1,000 classes. This dataset contains about 1.2 million images for training and 50,000 images for validation, which we use as a test set in our experiments.

Implementation details. We evaluate the proposed BNNs-mutual learning with different ResNet models [11]. Specifically, on CIFAR-10 and CIFAR100 datasets, we evaluate our method with ResNet20, ResNet32 and ResNet56 architectures, while on the large scale ImageNet dataset, we evaluate our method with ResNet18 architecture. On CIFAR-10 and CIFAR-100, we train all models for 300 epochs, divided into two stages. In the first stage, we train peer networks in a mutual learning manner until convergence without the loss term that enhances diversity in features by setting the β in Eq. 18 and Eq 19 to 0. Specifically, we train BNNs for up to 200 epochs. We use the Adam optimizer with an initial learning rate of 0.001, which is divided by 10 after 80, 120, 160, and 180 epochs. In the second stage, which comprises the last 100 epochs, we integrate the diversity in feature space into our BNNs-mutual learning and use the convergence from the first stage as an initial point for the training procedure. We continue training for the last 100 epochs in which the learning rate is reset to 0.001 and then is divided by 10 after 230, 260, and 290 epochs. The networks are divided into 3 blocks, and we combine features extracted from block 2 and block 3 of each BNN for feature diversity. In all experiments, the temperature T is set to 3 (Eq. 7 and Eq. 8), and α and β (Eq. 18 and Eq. 19) are set to 1 and 2, respectively. On the ImageNet dataset, we train BNN models for up to 45 epochs. For the first 30 epochs, we set the β to 0 and use the Adam optimizer with an initial learning rate of 0.001, which is divided by 10 after 15, 20, and 25 epochs. For the last 15 epochs, we reinitialize the learning rate to 0.001 and divide it by 10 after 35 and 40 epochs. We divide the networks into 4 blocks and combine features extracted from block 2, block 3, and block 4 of each BNN for feature diversity. The temperature T , α , and β (Eq. 18 and Eq. 19) are set to 1. For all settings, we use a pair of Bayesian neural network models with the variational parameters to approximate posteriors (Eq. 6); one is initialized from scratch, and the other is initialized with the mean µ from the deterministic model trained on the same dataset. We empirically find that this setting achieves better results compared to peer models that are initialized under the same conditions. We sample a single network during the training phase, while for the inference phase, we conduct an ensemble of 50 sample models across all settings. All the reported results are average of 3 trials.

4.2 Experimental results

In this section, we compare our proposed method against deterministic models (DNN), Bayesian neural networks trained with mutual learning (DML) and trained without mutual learning (vanilla). In the vanilla approach, each BNN model is trained individually. We evaluate BNNs using three metrics including top-1 classification accuracy (ACC), negative log-likelihood (NLL), and expected calibration error (ECE) [9]. We first sample from the posterior distribution of the weights. We then average the predicted probabilities to calculate ACC and evaluate the negative log-likelihood of these probabilities for NLL. Meanwhile, to compute ECE, we categorize data samples into 20 bins based

on their confidence scores and measure the difference between the model's predicted probability and its actual accuracy.

Comparative results on CIFAR-100 and CIFAR-10. We present top-1 accuracy (%) (ACC), NLL, and ECE on both the CIFAR-100 and CIFAR-10 datasets for various peer BNNs, as shown in Tables 1 and 2. In terms of the average accuracy of peer BNNs in mutual learning, which represents the mean accuracy from the peer models, our proposed method consistently outperforms the BNNs trained with traditional mutual learning [39], with the highest improvement of 0.95% for ResNet20 on CIFAR-100 and 0.58% for ResNet20 on CIFAR-10, respectively. In terms of BNNs initialized with the mean value from the pre-trained deterministic model, our approach consistently outperforms the others on both datasets. The highest improvement is with ResNet20 architecture, i.e., our method outperforms DML 0.84% in terms of top-1 accuracy on the CIFAR-100 dataset. Regarding calibration results and negative log-likelihood, we also achieve better performance compared to other approaches across all pairs. Additionally, our BNN models that are initialized with the mean value from the pre-trained deterministic models outperform the deterministic models (DNN) in all metrics ACC, NLL, and ECE.

Comparative results on ImageNet. To further demonstrate the effectiveness of our proposed method, we conduct an experiment on the large-scale and complex ImageNet dataset, employing peer BNN models of ResNet18. As shown in Table 3, BNN networks trained with our method outperform those trained with DML [39], showing a gain of 0.49% for the average accuracy of two peer BNN networks and a gain of 0.51% for the BNN initialized with the mean value from the pre-trained deterministic model. This indicates that our method is applicable to large-scale datasets.

4.3 Ablation studies

Impact of diversity in our proposed method. In this section, we investigate how different diversity approaches in our proposed method contribute to the performance of variational BNNs in a mutual learning setting. All experiments are conducted on the CIFAR-100 dataset with peer BNNs using ResNet20. In Table 4, we present the top-1 accuracy of BNNs-mutual learning in various settings. The results show that when promoting diversity in parameter space, there are 0 . 51% and 0 . 61% improvements over the mutual learning that only uses L logits for the BNNs initialized from scratch and initialized with the mean value from the deterministic neural network, respectively. When promoting diversity in feature space, the improvements are 0 . 3% and 0 . 43% over the same settings, respectively. This indicates that promoting diversity in parameter space has a more significant impact on model performance. By including both diversity in parameter and feature spaces, the highest improvement is achieved. This demonstrates that both loss terms are complementary to each other in achieving better performance in BNNs-mutual learning.

The ability of uncertainty estimation. To evaluate uncertainty estimation, we adopt Bayesian active learning by disagreement (BALD) [7], which quantifies the mutual information between the posterior distribution of parameters and the predictive distribution. We follow the evaluation methodology presented in Bayesian deep learning (BDL) benchmarks [6] to compare the accuracy in which only partial data are retained depending on uncertainty estimates obtained from the proposed method. Specifically, we compute uncertainty [7] for every testing image. Then we keep top images with the lowest uncertainty and report the top-1 accuracy of those kept images. This provides a more comprehensive evaluation metric of model performance and the ability to model uncertainty. Table 5 shows that accuracy increases when fewer data are retained on all networks. Furthermore, BNNs trained with our proposed method perform better than those trained without promoting diversity in mutual learning and those trained without mutual learning. This indicates that the proposed method produces more reliable uncertainty than the original mutual learning.

5 Conclusion

In this paper, we propose a novel approach to enhance BNNs by leveraging the effectiveness of mutual learning. Our approach encourages diversity in both model parameter space and feature space during the mutual learning of peer BNNs. By promoting such diversity, we encourage peer networks

to learn distinct features and perspectives from one another. This enhances the mutual learning of the peer networks. We extensively evaluate our method on CIFAR-10, CIFAR-100, and ImageNet datasets, demonstrating that the proposed method outperforms BNNs trained with and without mutual learning while producing reliable uncertainty estimates. We expect that our work will broaden the topic of mutual learning in which the model parameter space is taken into account, in addition to the traditional approaches that only consider the feature space.

Acknowledgments

This work was partly supported by ARC DP23 grant DP230101176 and by the Air Force Office of Scientific Research under award number FA2386-23-1-4044.

References

Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In ICML , 2015.

Donald Bures. An extension of Kakutani's theorem on infinite product measures to the tensor product of semifinite w ∗ -algebras. Transactions of the American Mathematical Society , 135:199-212, 1969.

Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In ICML , 2014.

Inseop Chung, SeongUk Park, Jangho Kim, and Nojun Kwak. Feature-map-level online adversarial knowledge distillation. In ICML , 2020.

Sebastian Farquhar, Michael Osborne, and Yarin Gal. Radial Bayesian Neural Networks: Beyond discrete support in large-scale Bayesian deep learning. AISTATS , 2020.

Angelos Filos, Sebastian Farquhar, Aidan N Gomez, Tim GJ Rudner, Zachary Kenton, Lewis Smith, Milad Alizadeh, Arnoud De Kroon, and Yarin Gal. A systematic comparison of bayesian deep learning robustness in diabetic retinopathy tasks. arXiv , 2019.

Yarin Gal and Zoubin Ghahramani. Dropout as a Bayesian approximation: Representing model uncertainty in deep learning. In ICML , 2016.

Soumya Ghosh, Jiayu Yao, and Finale Doshi-Velez. Structured variational learning of Bayesian neural networks with horseshoe priors. In ICML , 2018.

Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In ICML , 2017.

Qiushan Guo, Xinjiang Wang, Yichao Wu, Zhipeng Yu, Ding Liang, Xiaolin Hu, and Ping Luo. Online knowledge distillation via collaborative learning. In CVPR , 2020.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In CVPR , 2016.

Byeongho Heo, Jeesoo Kim, Sangdoo Yun, Hyojin Park, Nojun Kwak, and Jin Young Choi. A comprehensive overhaul of feature distillation. In ICCV , 2019.

Geoffrey Hinton, Oriol Vinyals, and Jeffrey Dean. Distilling the Knowledge in a Neural Network. In NIPS Deep Learning and Representation Learning Workshop , 2014.

Mohammad Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivastava. Fast and scalable Bayesian deep learning by weight-perturbation in adam. In ICML , 2018.

Jangho Kim, Minsung Hyun, Inseop Chung, and Nojun Kwak. Feature fusion for online mutual knowledge distillation. In ICPR , 2021.

Durk P Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparameterization trick. NeurIPS , 2015.

Nikos Komodakis and Sergey Zagoruyko. Paying more attention to attention: improving the performance of convolutional neural networks via attention transfer. In ICLR , 2017.

Anoop Korattikara Balan, Vivek Rathod, Kevin P Murphy, and Max Welling. Bayesian dark knowledge. NeurIPS , 2015.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Master's thesis, University of Toronto, 2009.

Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-10 (canadian institute for advanced research), 2010.

Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose Bayesian inference algorithm. NeurIPS , 2016.

Radford M Neal. Bayesian learning for neural networks , volume 118. Springer Science & Business Media, 2012.

Van-Anh Nguyen, Tung-Long Vuong, Hoang Phan, Thanh-Toan Do, Dinh Phung, and Trung Le. Flat seeking bayesian neural network. In NeurIPS , 2023.

Victor M-H Ong, David J Nott, and Michael S Smith. Gaussian variational approximation with a factor covariance structure. Journal of Computational and Graphical Statistics , 27(3):465-478, 2018.

Nikolaos Passalis and Anastasios Tefas. Learning deep representations with probabilistic knowledge transfer. In ECCV , 2018.

Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable laplace approximation for neural networks. In ICLR , 2018.

Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio. FitNets: Hints for thin deep nets. In ICLR , 2015.

Simone Rossi, Sebastien Marmin, and Maurizio Filippone. Walsh-Hadamard variational inference for Bayesian deep learning. In NeurIPS , 2020.

Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. IJCV , 2015.

Gehui Shen, Xi Chen, and Zhihong Deng. Variational learning of Bayesian neural networks via Bayesian dark knowledge. In IJCAI , 2021.

Jakub Swiatkowski, Kevin Roth, Bastiaan Veeling, Linh Tran, Joshua Dillon, Jasper Snoek, Stephan Mandt, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. The k-tied normal distribution: A compact parameterization of gaussian mean field posteriors in Bayesian neural networks. In ICML , 2020.

Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research , 2008.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. NIPS , 30, 2017.

Cédric Villani. Topics in optimal transportation . American Mathematical Soc., 2021.

Cédric Villani et al. Optimal transport: old and new . Springer, 2009.

Kuan-Chieh Wang, Paul Vicol, James Lucas, Li Gu, Roger Grosse, and Richard Zemel. Adversarial distillation of Bayesian neural network posteriors. In ICML , 2018.

Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In ICML , 2011.

Xiaobing Zhang, Shijian Lu, Haigang Gong, Zhipeng Luo, and Ming Liu. Amln: adversarial-based mutual learning network for online knowledge distillation. In ECCV , 2020.

Ying Zhang, Tao Xiang, Timothy M Hospedales, and Huchuan Lu. Deep mutual learning. In CVPR , 2018.

Xiatian Zhu, Shaogang Gong, et al. Knowledge distillation by on-the-fly native ensemble. NeurIPS , 2018.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposed a mutual learning approach to learn a pair of Bayesian Neural Network(BNN). The posterior of BNN is approximated by Variational Inference using a Gaussian distribution with a diagonal covariance matrix. To make the BNN learn different perspective of the data, the author proposed to increase the diversity in parameter space and intermediate feature space by adding the an estimate of distance between parameter distribution and fused feature distribution of two BNN models into the objective function. Empirically, the proposed method outperform existing mutual learning method and vanilla BNN model in terms of accuracy, negative log likelihood loss and expected calibration error. An ablation study is also provided to investigate the usefulness of each component.\n\nSTRENGTHS:\nThe paper is well written and easy to follow. Increasing the diversity of parameter distribution and intermediate feature distribution of peer BNN models to boost performance is an interesting idea. Experiments and detailed ablation study demonstrate the effectiveness of proposed method.\n\nWEAKNESSES:\n1. It is mentioned in the abstract and introduction that the BNN model with variational inference may underperform deterministic model or BNN obtained by MCMC, the baseline only involves BNN model trained with(DML) or without(vanilla) mutual learning. Would the proposed method close the gap to some extent? Data augmentation, optimizer may all affect performance, so it is still helpful to include deterministic model results follow with same training setup. I would expect the BNN model to outperform deterministic model at least in NLL and ECE, and with the 50 ensemble, it can outperform the accuracy. \n\n2. Continue with last point, for MCMC method (e.g. in line 81 of the paper), I agree that traditional MCMC method(e.g. Metropolis Hasting) may not be feasible for large model, and memory storage can be an issue for MCMC method. But I don't think the stochastic gradient MCMC cited in line 81 would require prohibitive computational cost, it behaves like adding a noise to at each step of standard SGD training. \n\n3. The code is not provided so it may hurt the reproducibility of the paper.\n\nQUESTIONS:\n1. To my knowledge, it is not very clear if variational distribution(e.g. Gaussian with diagonal covariance matrix) can approximate the true posterior very well, can the author comment a bit on this, e.g. how would different choice of variational family affect the model?\n\n2. In line 264 and line 6 of algorithm 1, it is mentioned that one BNN model is initialized with a trained model and this lead to better results empirically. Can the author discuss more on why this happened? It is a bit wired for me as it seems in the implementation detail, the pre-trained model and the model from scratch are trained with same optimizer and learning rate schedule.\n\n3. Seems like $\\alpha$ $\\beta$ are set to 1,2 for CIFAR and 1,1 for Imagenet, these two parameters controls the strength of proposed penalty to the model, can the author comments a bit more on how sensitive are the model to those parameters, It can help to illustrate how diversity helps model performance.\n\n4. As mentioned in line 268, results are average of 3 trials, I think it would be better to include the standard deviation as well to boost the significance of the results.\n\n5. In figure A.3 in supplementary material, looks like a sharp increase of KL divergence between the fused feature distributions at around 30 epochs, but the penalty for feature is only added for last 100 epochs, can the author explain more on this?\n\nLIMITATIONS:\nThe authors addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-01"}

comments: SUMMARY:
The paper proposes a method to combine deep mutual learning with BNN to diversify the weight distributions of each BNN networks in a pair or ensemble, to improve performance.

STRENGTHS:
1. AFAIK this is the first work combining mutual learning with BNN, so the authors can claim this point.
2. The paper is in general written clearly and easy to follow.
3. Experiments are adequate with ablation studies on individual features impact on diversity.

WEAKNESSES:
1. Some design choices are found to be "empirically" working well without too much discussion or hypothesis.
2. Would be interesting to see how the model performs for o.o.d test data, especially uncertainty performance.

QUESTIONS:
Line 178-179: The authors said adding the D(...) term will rapidly increase of this term and impact training. Wouldn't putting a smaller scaling factor for this term fix this issue?

LIMITATIONS:
None.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper titled addresses the challenge of improving the performance of Bayesian Neural Networks (BNNs) by leveraging the concept of mutual learning. BNNs provide a means for quantifying uncertainty in predictions through probability distributions of model parameters. However, BNNs often fall short in performance compared to their deterministic counterparts. The authors propose a novel approach that employs deep mutual learning to enhance the capabilities of BNNs.\n\nSTRENGTHS:\n1. Innovative Approach: The paper introduces a novel method that combines deep mutual learning with Bayesian Neural Networks. By promoting diversity in both network parameter distributions and feature distributions, the proposed approach enables peer networks to acquire distinct features, capturing different characteristics of the input data. This innovative technique enhances the effectiveness of mutual learning in BNNs.\n2. Detailed algorithm description: The paper provides a thorough and detailed description of the proposed algorithm for improving the performance of Bayesian Neural Networks (BNNs) through deep mutual learning.\n3. Comprehensive Experiments: The authors conduct extensive experiments to evaluate the proposed approach thoroughly. The experimental results are statistically sound and demonstrate significant improvements in classification accuracy, negative log-likelihood, and expected calibration error compared to traditional mutual learning methods for BNNs.\n\nWEAKNESSES:\n1. Limited variety in experimental validation: One weakness of the paper is that the proposed approach and its effectiveness are only verified through experiments conducted on Residual Neural Networks (ResNets). It would have been beneficial to include experiments on a diverse set of network architectures to demonstrate the approach's effectiveness across different model types and complexities. \n2. Lack of detailed explanation for temperature, α, and β: One weakness of the paper is the limited explanation provided for the temperature parameter (T), α, and β, which are crucial components of the proposed approach. These parameters play a significant role in controlling the diversity of network parameter distributions and feature distributions, but their specific effects and optimal values are not thoroughly discussed.\n3. Weakness in the conclusion: The current conclusion merely restates the experimental results and does not highlight the broader implications of the proposed approach or its potential impact on the field.\n\nQUESTIONS:\nThe author should supplement more experiments to prove its effectiveness.\n\nLIMITATIONS:\nThe author should supplement more experiments to prove its effectiveness and strengthen the conclusion.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper focuses on improving the accuracy of BNNs by promoting diversity in both parameter space and feature space while training two peer BNNs with mutual learning between them. More specifically, they train two variational BNNs with a mean-field Gaussian variational loss for each along with a KL divergence term between the (temperature-scaled) predictive distributions of the two models, a Wasserstein distance term between the corresponding approximate posterior distributions across the two models (added as a softplus(-distance) term), and a KL divergence term between corresponding feature distributions. On the latter term, instead of directly maximizing the distance between corresponding feature distributions, they instead do so on \"fused feature distributions\". To do so, they use learned cross-attention to fuse the features from multiple feature levels in a model (two at a time). Then, they use the KL divergence between the distributions of the fused feature distributions of the two peer networks. To derive the distributions, they use the conditional probability density defined as $p_{i|j} = \\frac{K(F'_i, F'_j)}{sum_{k=1, k \\noteq i}^n K(F'_k, F'_j)}$, where $K(F'_a, F'_b)$ is a kernel function between two fused feature representations. Given those conditional probs, they compute a KL divergence term. Similar to the parameter space diversity term, they add this term to the loss as softplus(-divergence). The paper claims to be the first to propose maximizing the distance between feature distributions to promote diversity. In terms of experiments, the paper includes results for ResNet models on CIFAR-10/100 and ImageNet, measuring accuracy, NLL, and ECE as metrics, and comparing different approaches.\n\nSTRENGTHS:\nThe paper does a great job of precisely articulating the modeling approach, and discussing the relevant background info. More specifically, the proposed approach of adding terms to promote diversity in parameter and feature space is clear and would be easy to reimplement.\n\nWEAKNESSES:\nMy main concern is with the experiment section. More specifically, a few key details are unclear in the text, and importantly a deterministic baseline is missing that I believe should be present given the framing of the paper and relevant literature. Please see the Questions below. Given updates, I believe the paper would be great and I would gladly update my rating.\n\nQUESTIONS:\nMain:\n- In the experiments, a few details are currently unclear. The following points are on Table 1, but generalize to all three tables. Please clarify these details here and in the paper.\n  - Consider the ResNet20 section of Table 1. Is my understanding correct that the \"ResNet20\" results are for a pair of BNNs trained from scratch, while the \"ResNet20*\" results are for a pair of BNNs trained with the approximate posterior means set to the values from a deterministic model?\n  - Is it correct that all results (all three metrics across all three approaches) are computed after averaging the predicted probs from the pair of models?\n- For the experiments, a deterministic baseline is missing. Given the intro that discusses how BNNs can lag behind deterministic models in acc (though not always), the experiments lack a comparison. It would be helpful to understand how the proposed approach compares to a deterministic baselines, specifically a single deterministic model and a size-2 deep ensemble. Could you add this as a baseline? I would consider this to be a blocker for the paper given the framing and relevant literature.\n\nOther:\n- The KL divergence term is scaled by the square of the temperature -- why?\n- How did you choose the values for temp, alpha, and beta? They differ between CIFAR-10/100 and ImageNet. Did you ablate values?\n\nMinor comments:\n- updating lines 17 & 22 of Alg 1 could be helpful for readability\n\nLIMITATIONS:\nNo limitations are included.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a novel method for enhancing the performance of Bayesian Neural Networks (BNNs) by employing deep mutual learning. The proposed approach aims to enhance the diversity of both network parameter distributions and feature distributions, encouraging individual networks to capture unique characteristics of the input data. The effectiveness of the proposed method is demonstrated on datasets, including CIFAR10, CIFAR100, and ImageNet.\n\nSTRENGTHS:\nThe proposed method improves performance and uncertainty estimation while reducing the expected calibration error (ECE).  The technical approach is novel as the method introduces mutual learning in the context of BNNs and first to propose maximizing the distance between feature distributions and parameter distributions. The paper includes large scale data experiments (ImageNet) and ablation studies to demonstrate the effectiveness of each technical contribution introduced in this paper.\n\nWEAKNESSES:\nThe previous studies mentioned in the paper utilize alignments on feature maps [4] or predictions [38], rather than diversifying them. In contrast, the proposed method diversifies both feature distributions and parameter distributions which is an opposite approach to the previous works. Interestingly, both alignment-based and diversification methods improves performance over vanilla BNNs, as indicated in Table 1, 2, 3, and 5. However, the paper does not explicitly explain the reasons behind the performance improvements resulting from these contrasting approaches.\n\nGiven the observed contradicting results in the experiments, where the alignment-based method (DML [38]) also enhances the performance of BNNs, an important question arises: could combining alignment-based methods with parameter diversification further improve BNN performance? Alternatively, is it necessary to diversify both feature and parameter distributions to achieve significant improvements?\n\nIn the experiment section, the proposed method is only compared with [38] and not with [4]. \n\nHyperparameters used for CIFAR experiments and ImageNet experiments are different. However, the paper does not describe details regarding the hyperparameter tuning or determination.\n\nQUESTIONS:\n3 block resnet is used for CIFAR experiments while 4 block resnet is used for ImageNet experiments. Why different form of resents are used for different datasets?\n\nLIMITATIONS:\nLimitations are shortly addressed in the supplementary.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-09"}