SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Algorithm 1 Diversity for Bayesian Neural Networks in Mutual Learning.
1: procedure TRAIN(S,  $n_{\mathrm{batch}}$  size,  $n_{\mathrm{epochs}}$    
2:  $\triangleright S = \{x,y\}$  training dataset 3: nbatch size: size of a mini-batch 4: nepochs: number of training epochs 5:  $\mu_{\theta_1},\sigma_{\theta_1}\gets$  RANDOM INITIALIZATION Initialize net B1   
6:  $\mu_{\theta_2}\gets$  PRE-TRAINED MODEL Initialize net B2   
7:  $\sigma_{\theta_2}\gets$  RANDOM INITIALIZATION   
8: for epoch  $= 1\rightarrow n_{\mathrm{epochs}}$  do   
9: for iter in iterations do   
10: sample a mini-batch:  $\mathbb{B} = \{(x_i,y_i):(x_i,y_i)\sim \mathcal{S}\}_{i = 1}^{n_{\mathrm{batch}}}$    
11: Compute losses shared between two peer networks 12: compute:  $D_{\mathrm{OT}}\leftarrow$  WASSERSTEIN DISTANCE(  $\mathcal{N}(\mu_{\theta_1},\sigma_{\theta_1}^2),\mathcal{N}(\mu_{\theta_2},\sigma_{\theta_2}^2))$  Eq. (11)   
13: compute diverse parameter distribution loss:  $\mathcal{L}_{\mathrm{diverse~param}}$  Eq. (12)   
14: Network B1 15: compute logit loss for B1:  $\mathcal{L}_{\mathrm{logits}}^{B1}$  Eq. (7)   
16: compute diverse feature loss for B1:  $\mathcal{L}_{\mathrm{diverse~feat}}^{B1}$  Eq. (16)   
17: compute total loss for B1:  $\mathcal{L}^{B1} = \mathcal{L}_{\mathrm{logits}}^{B1} + \alpha \mathcal{L}_{\mathrm{diverse~param}} + \beta \mathcal{L}_{\mathrm{diverse~feat}}^{B1}$  Eq. (18)   
18: update parameter:  $\mu_{\theta_1},\sigma_{\theta_1}\gets$  ADAM(LB1)   
19: Network B2 20: compute logit loss for B2:  $\mathcal{L}_{\mathrm{logits}}^{B2}$  Eq. (8)   
21: compute diverse feature loss for B2:  $\mathcal{L}_{\mathrm{diverse~feat}}^{B2}$  Eq. (17)   
22: compute total loss for B2:  $\mathcal{L}^{B2} = \mathcal{L}_{\mathrm{logits}}^{B2} + \alpha \mathcal{L}_{\mathrm{diverse~param}} + \beta \mathcal{L}_{\mathrm{diverse~feat}}^{B2}$  Eq. (19)   
23: update parameter:  $\mu_{\theta_2},\sigma_{\theta_2}\gets$  ADAM(LB2)   
24: return  $\mu_{\theta_1},\sigma_{\theta_1},\mu_{\theta_2},\sigma_{\theta_2}$
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:
$$
p _ {i \mid j} = \frac {K \left(F _ {i} ^ {\prime} , F _ {j} ^ {\prime}\right)}{\sum_ {k = 1 , k \neq i} ^ {n} K \left(F _ {k} ^ {\prime} , F _ {j} ^ {\prime}\right)}, \tag {14}
$$
where  $K(a,b)$  is a kernel function. In our work, we adopt the similarity metric  $K(a,b) = \frac{1}{2} \left( \frac{a^T b}{\|a\| \|b\|} + 1 \right)$  which has been used in [25] for offline knowledge distillation.
Given that  $\mathcal{Q}$  and  $\mathcal{P}$  are the probability distributions of fused features from BNN models  $B1$  and  $B2$ , respectively, the distance between fused feature distributions is represented as:
$$
D _ {\mathrm {K L}} \left(G ^ {B 2}, G ^ {B 1}\right) = \mathrm {K L} [ \mathcal {P} \| \mathcal {Q} ] = \int \mathcal {P} (x) \log \left(\frac {\mathcal {P} (x)}{\mathcal {Q} (x)}\right) \mathrm {d} x. \tag {15}
$$
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  $B1$ , and  $B2$ :
$$
\mathcal {L} _ {\text {d i v e r s e f e a t}} ^ {B 1} = \log \left(1 + \exp \left(- D _ {\mathrm {K L}} \left(G ^ {B 2}, G ^ {B 1}\right)\right)\right), \tag {16}
$$
$$
\mathcal {L} _ {\text {d i v e r s e f e a t}} ^ {B 2} = \log \left(1 + \exp \left(- D _ {\mathrm {K L}} \left(G ^ {B 1}, G ^ {B 2}\right)\right)\right). \tag {17}
$$
Finally, the overall loss of our proposed Bayesian mutual learning for Bayesian neural networks  $B1$  and  $B2$  are defined as:
$$
\mathcal {L} ^ {B 1} = \mathcal {L} _ {\text {l o g i t s}} ^ {B 1} + \alpha \mathcal {L} _ {\text {d i v e r s e p a r a m}} + \beta \mathcal {L} _ {\text {d i v e r s e f e a t}} ^ {B 1}, \tag {18}
$$
$$
\mathcal {L} ^ {B 2} = \mathcal {L} _ {\text {l o g i t s}} ^ {B 2} + \alpha \mathcal {L} _ {\text {d i v e r s e p a r a m}} + \beta \mathcal {L} _ {\text {d i v e r s e f e a t}} ^ {B 2}, \tag {19}
$$
where  $\alpha$ , and  $\beta$  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.
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  $\beta$  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  $\mathcal{T}$  is set to 3 (Eq. 7 and Eq. 8), and  $\alpha$  and  $\beta$  (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  $\beta$  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  $\mathcal{T}$ ,  $\alpha$ , and  $\beta$  (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  $\mu$  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.
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
Table 1: Top-1 classification accuracy, NLL, and ECE on CIFAR-100 dataset. *Bayesian neural networks are initialized with the mean value from the pre-trained deterministic model.
ACC↑	NLL↓	ECE↓
DNN	Vanilla	DML	Ours	DNN	Vanilla	DML	Ours	DNN	Vanilla	DML	Ours
ResNet20	69.13	63.18	67.27	68.32	1.106	1.447	1.174	1.101	0.065	0.119	0.057	0.041
ResNet20*	-	67.62	69.61	70.45	-	1.438	1.073	1.043	-	0.152	0.047	0.038
Average	-	65.40	68.44	69.39	-	1.443	1.124	1.072	-	0.136	0.052	0.040
ResNet32	71.36	65.20	68.59	70.53	1.074	1.530	1.169	1.029	0.080	0.150	0.087	0.043
ResNet32*	-	69.76	71.45	72.14	-	1.542	1.012	0.975	-	0.173	0.064	0.040
Average	-	67.48	70.53	71.34	-	1.536	1.091	1.002	-	0.162	0.076	0.042
ResNet56	72.88	66.60	70.44	71.90	1.146	1.853	1.117	1.024	0.153	0.204	0.092	0.066
ResNet56*	-	71.94	73.79	74.20	-	1.764	0.997	0.954	-	0.189	0.085	0.069
Average	-	69.27	72.12	73.05	-	1.809	1.057	0.989	-	0.197	0.089	0.068
Table 2: Top-1 classification accuracy, NLL, and ECE on CIFAR-10 dataset. *Bayesian neural networks are initialized with the mean value from the pre-trained deterministic model.
ACC↑	NLL↓	ECE↓
DNN	Vanilla	DML	Ours	DNN	Vanilla	DML	Ours	DNN	Vanilla	DML	Ours
ResNet20	91.83	89.93	91.09	92.03	0.302	0.445	0.295	0.256	0.067	0.066	0.040	0.029
ResNet20*	-	92.47	92.67	92.89	-	0.383	0.232	0.228	-	0.055	0.031	0.025
Average	-	91.20	91.88	92.46	-	0.414	0.264	0.242	-	0.061	0.036	0.027
ResNet56	92.67	91.64	92.47	92.58	0.272	0.488	0.311	0.295	0.063	0.064	0.057	0.053
ResNet56*	-	93.33	93.50	93.91	-	0.431	0.243	0.237	-	0.054	0.041	0.039
Average	-	92.49	92.99	93.25	-	0.460	0.277	0.266	-	0.059	0.049	0.046
Table 3: Top-1 classification accuracy, NLL, and ECE on ImageNet dataset. *Bayesian neural networks are initialized with the mean value from the pre-trained deterministic model.
ACC↑	NLL↓	ECE↓
DNN	Vanilla	DML	Ours	DNN	Vanilla	DML	Ours	DNN	Vanilla	DML	Ours
ResNet18	69.76	64.30	66.88	67.36	1.247	1.505	1.357	1.343	0.130	0.134	0.134	0.131
ResNet18*	-	66.95	67.65	68.16	-	1.389	1.335	1.311	-	0.132	0.130	0.129
Average	-	66.63	67.27	67.76	-	1.447	1.346	1.327	-	0.133	0.132	0.130
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.
Table 4: Impact of diversity for each term in our proposed method in terms of top-1 classification accuracy. *Bayesian neural networks are initialized with the mean value from the pre-trained deterministic model.
Setting	\( \mathcal{L}_{\text{logits}} \)	\( \mathcal{L}_{\text{diverse param}} \)	\( \mathcal{L}_{\text{diverse feat}} \)	CIFAR100
ResNet20	ResNet20*	Average
(a)	✓	67.27	69.61	68.44
(b)	✓	✓	67.78	70.22	69.00
(c)	✓	✓	67.57	70.04	68.81
(d)	✓	✓	✓	68.32	70.45	69.39
Table 5: Comparison of top-1 classification accuracy on CIFAR-100 dataset when only partial data are retained based on the uncertainties of the predictions. *Bayesian neural networks are initialized with the mean value from the pre-trained deterministic model.
Method	Model	20% data retained	40% data retained	60% data retained	80% data retained
Vanilla	ResNet32	98.90	92.75	83.22	74.34
ResNet32*	99.10	95.38	87.43	78.29
DML	ResNet32	99.20	94.83	86.33	77.43
ResNet32*	98.95	95.43	87.73	80.10
Ours	ResNet32	98.95	94.80	87.27	78.59
ResNet32*	99.25	95.88	88.25	80.60
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  $\mathcal{L}_{\mathrm{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.
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.
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.
[26] Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable laplace approximation for neural networks. In ICLR, 2018.
[27] Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio. FitNets: Hints for thin deep nets. In ICLR, 2015.
[28] Simone Rossi, Sebastien Marmin, and Maurizio Filippone. Walsh-Hadamard variational inference for Bayesian deep learning. In NeurIPS, 2020.
[29] 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.
[30] Gehui Shen, Xi Chen, and Zhihong Deng. Variational learning of Bayesian neural networks via Bayesian dark knowledge. In IJCAI, 2021.
[31] 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.
[32] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 2008.
[33] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. NIPS, 30, 2017.
[34] Cédric Villani. Topics in optimal transportation. American Mathematical Soc., 2021.
[35] Cédric Villani et al. Optimal transport: old and new. Springer, 2009.
[36] Kuan-Chieh Wang, Paul Vicol, James Lucas, Li Gu, Roger Grosse, and Richard Zemel. Adversarial distillation of Bayesian neural network posteriors. In ICML, 2018.
[37] Max Welling and Yee W Teh. Bayesian learning via stochastic gradient Langevin dynamics. In ICML, 2011.
[38] Xiaobing Zhang, Shijian Lu, Haigang Gong, Zhipeng Luo, and Ming Liu. Amln: adversarial-based mutual learning network for online knowledge distillation. In ECCV, 2020.
[39] Ying Zhang, Tao Xiang, Timothy M Hospedales, and Huchuan Lu. Deep mutual learning. In CVPR, 2018.
[40] Xiatian Zhu, Shaogang Gong, et al. Knowledge distillation by on-the-fly native ensemble. NeurIPS, 2018.