SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Cuong Pham
Cuong C. Nguyen
Trung  $\mathbf{L}\mathbf{e}^{1}$
Dinh Phung 1,4
Gustavo Carneiro<sup>3</sup>
$^{2}$ Australian Institute for Machine Learning, University of Adelaide, Australia
<sup>3</sup>Centre for Vision, Speech and Signal Processing, University of Surrey, United Kingdom
$^{4}$ VinAI, Vietnam
{cuong.pham1, trunglm, dinh.phung, toan.do}@monash.edu
cuong.nguyen@adelaide.edu.au, g.carneiro@surrey.ac.uk
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.
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  $L1$  or  $L2$  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.
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 = 1}^{N}$  with  $x_{n}\in \mathcal{X}$  being an input and  $y_{n}\in \mathcal{Y}$  being the corresponding output. The aim is to find a model  $f(:,w):\mathcal{X}\to \mathcal{Y}$ , parameterized by  $w\in \mathcal{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:
$$
p (w | \mathcal {S}) = \prod_ {n = 1} ^ {N} \frac {p \left(y _ {n} \mid x _ {n} , w\right) p (w)}{\int p \left(y _ {n} \mid x _ {n} , w\right) \mathrm {d} w}. \tag {1}
$$
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|\mathcal{S})$  by a variational posterior  $q(w;\theta)$  parameterized by  $\theta$ . The variational posterior  $q(w;\theta)$  can be obtained by minimizing the following KL divergence:
$$
\begin{array}{l} \theta^{*} = \operatorname *{arg  min}_{\theta}\mathrm{KL}[q(w;\theta)\| p(w|\mathcal{S})] \\ = \arg \min  _ {\theta} \mathrm {K L} [ q (w; \theta) \| p (w) ] - \int q (w; \theta) \sum_ {n = 1} ^ {N} \log p \left(y _ {n} \mid x _ {n}, w\right) \mathrm {d} w \tag {2} \\ = \arg \min _ {\theta} \mathrm {K L} [ q (w; \theta) \| p (w) ] + \mathbb {E} _ {q (w; \theta)} \left[ - \sum_ {i = n} ^ {N} \log p (y _ {n} | x _ {n}, w) \right], \\ \end{array}
$$
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; \theta) = \mathcal{N}(\mu, \mathrm{diag}(\sigma^2))$  and apply a parametrization trick to represent each weight in the neural network by a mean and a standard deviation:
$$
w = \mu + \sigma \odot \epsilon , \quad \epsilon \sim \mathcal {N} (0, I), \tag {3}
$$
where:  $\odot$  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:
$$
w = \mu + \sigma \odot \frac {\epsilon}{\| \epsilon \|} r, \quad r = | \rho | \quad \text {f o r} \quad \rho \sim \mathcal {N} (0, I), \tag {4}
$$
where  $\|\cdot\|$  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  $L1$  or  $L2$  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.1 Deep Mutual Learning for Bayesian Neural Networks
We formulate our proposed method by considering two peer Bayesian neural networks,  $B1$  and  $B2$ , 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:
$$
q (w; \theta_ {i}) = \mathcal {N} (w; \mu_ {\theta_ {i}}, \operatorname {d i a g} \left(\sigma_ {\theta_ {i}} ^ {2}\right)), \quad i = 1: 2. \tag {5}
$$
where  $\mu_{\theta_i},\sigma_{\theta_i}\in \mathbb{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:
$$
\mathcal {L} \left(\theta_ {i}\right) = \mathrm {K L} \left[ q (w; \theta_ {i}) \| p (w) \right] + \mathbb {E} _ {q \left(w; \theta_ {i}\right)} \left[ - \sum_ {n = 1} ^ {N} \log p \left(y _ {n} \mid x _ {n}, w\right) \right]. \tag {6}
$$
Let  $z^{B1}$  and  $z^{B2}$  be the logit outputs of the Bayesian neural networks  $B1$  and  $B2$ ,  $\mathcal{T}$  be the temperature parameter to control the smoothness of the output distribution. We denote soft logit  $p^{B1} = \mathrm{softmax}(z^{B1} / \tau)$ , and  $p^{B2} = \mathrm{softmax}(z^{B2} / \tau)$ . The logit loss terms are defined as follows:
$$
\mathcal {L} _ {\text {l o g i t s}} ^ {B 1} = \mathcal {L} \left(\theta_ {1}\right) + \mathcal {T} ^ {2} \times \mathrm {K L} \left[ p ^ {B 2} \| p ^ {B 1} \right], \tag {7}
$$
$$
\mathcal {L} _ {\text {l o g i t s}} ^ {B 2} = \mathcal {L} \left(\theta_ {2}\right) + \mathcal {T} ^ {2} \times \mathrm {K L} \left[ p ^ {B 1} \| p ^ {B 2} \right]. \tag {8}
$$
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  $\theta_{1} = (\mu_{\theta_{1}},\Sigma_{\theta_{1}})$  and  $\theta_{2} = (\mu_{\theta_{2}},\Sigma_{\theta_{2}})$  for BNN models  $B1$  and  $B2$ , respectively. We then measure the distance between two posterior distributions  $D(q(w;\theta_1),q(w;\theta_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]:
$$
D \left(q (w; \theta_ {1}), q (w; \theta_ {2})\right) = \mathcal {W} _ {2} \left(q (w; \theta_ {1}), q (w; \theta_ {2})\right) ^ {2} = \left\| \mu_ {\theta_ {1}} - \mu_ {\theta_ {2}} \right\| ^ {2} + \mathcal {B} \left(\Sigma_ {\theta_ {1}}, \Sigma_ {\theta_ {2}}\right) ^ {2}, \tag {9}
$$
where  $\mathcal{B}(.)$  is the Bures distance [2] represented by:
$$
\mathcal {B} \left(\Sigma_ {\theta_ {1}}, \Sigma_ {\theta_ {2}}\right) ^ {2} = \operatorname {t r} \left(\Sigma_ {\theta_ {1}} + \Sigma_ {\theta_ {2}} - 2 \left(\Sigma_ {\theta_ {1}} ^ {1 / 2} \Sigma_ {\theta_ {2}} \Sigma_ {\theta_ {1}} ^ {1 / 2}\right) ^ {1 / 2}\right). \tag {10}
$$
With diagonal covariances  $\Sigma_{\theta_1} = \mathrm{diag}(\sigma_{\theta_1}^2)$ , and  $\Sigma_{\theta_2} = \mathrm{diag}(\sigma_{\theta_2}^2)$ , we have a simple form of distance between two approximate posteriors:
$$
D \left(q \left(w; \theta_ {1}\right), q \left(w; \theta_ {2}\right)\right) = \left\| \mu_ {\theta_ {1}} - \mu_ {\theta_ {2}} \right\| ^ {2} + \left\| \sigma_ {\theta_ {1}} - \sigma_ {\theta_ {2}} \right\| ^ {2}. \tag {11}
$$
To promote diversity in the parameter space, a straightforward approach would be to add a negative  $D(q(w;\theta_1),q(w;\theta_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;\theta_1),q(w;\theta_2))$  which would impair the training process. Therefore, we propose minimizing the following loss function:
$$
\mathcal {L} _ {\text {d i v e r s e p a r a m}} = \log \left(1 + \exp \left(- D \left(q (w; \theta_ {1}), q (w; \theta_ {2})\right)\right)\right). \tag {12}
$$
The above loss function exhibits desirable properties:  $\mathcal{L}_{\mathrm{diverse~param}} \geq 0$ , and minimizing  $\mathcal{L}_{\mathrm{diverse~param}}$  implies maximizing  $D(q(w; \theta_1), q(w; \theta_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  $L2$  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_{k}^{B1}$  and  $F_{k}^{B2}$  be the extracted features at the  $k^{th}$  block of the BNN models,  $B1$  and  $B2$ , 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_{k}^{B1}$  and  $F_{k}^{B2}$  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  $B1$ , the combined feature using cross-attention can be expressed as:
$$
F ^ {\prime B 1} = u \left(F _ {k} ^ {B 1}, F _ {k + 1} ^ {B 1}\right) = \operatorname {s o f t m a x} \left[ \left(\mathbf {W} _ {Q} F _ {k + 1} ^ {B 1}\right) \left(\mathbf {W} _ {K} F _ {k} ^ {B 1}\right) ^ {\top} \right] \mathbf {W} _ {V} F _ {k} ^ {B 1}, \tag {13}
$$
where  $\mathbf{W}_Q$ ,  $\mathbf{W}_K$ , and  $\mathbf{W}_V$  represent learnable parameters for query, key, and value, respectively. Similarly, we have combined features  $F'^{B2} = u(F_k^{B2}, F_{k+1}^{B2})$ . Specifically, in a mini-batch with  $n$  samples, we denote  $G^{B1} = \{F_j'^{B1}\}_{j=1}^n$ , and  $G^{B2} = \{F_j'^{B2}\}_{j=1}^n$  represented fused features derived from BNN models  $B1$  and  $B2$ , 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