May 27, 2026
Knowledge distillation is a powerful method for model compression, enabling the efficient deployment of complex deep learning models (teachers), including large language models. However, its underlying statistical mechanisms remain unclear, and uncertainty evaluation is often overlooked, especially in real-world scenarios requiring diverse teacher expertise. To address these challenges, we introduce Multi-Teacher Bayesian Knowledge Distillation (MT-BKD), where a distilled student model learns from multiple teachers within the Bayesian framework. Our approach leverages Bayesian inference to capture inherent uncertainty in the distillation process. We introduce a teacher-informed prior, integrating external knowledge from teacher models and task-specific training data, offering better generalization, robustness, and scalability. Additionally, an entropy-based weighting mechanism adaptively adjusts each teacher’s influence, allowing the student to combine multiple sources of expertise effectively. MT-BKD enhances the interpretability of the student model’s learning process, improves predictive accuracy, and provides uncertainty quantification. We validate MT-BKD on both synthetic and real-world tasks, including protein subcellular location prediction and image classification. Our experiments show improved performance and robust uncertainty quantification, highlighting the strengths of our MT-BKD framework.
1
1
Multi-Teacher Bayesian Knowledge Distillation
Uncertainty Quantification; Bayesian Priors; Large Language Models; Protein Subcellular Prediction; Image Classification
The rapid advancements in deep learning have led to the development of increasingly complex models with billions of parameters, achieving state-of-the-art performance across numerous tasks [1], [2]. However, deploying these large models in real-world applications is often infeasible due to significant computational and storage demands. Knowledge distillation (KD) offers a powerful solution by transferring knowledge from a large, pre-trained teacher model to a smaller, more efficient student model [3]–[5]. By training the student to replicate the teacher’s predictions, KD produces lightweight models with competitive performance, suitable for resource-constrained environments [6]. This approach is particularly prevalent in the development of large language models, where vendors routinely release both full-scale and lightweight versions, with the latter often trained using KD [7].
Many real-world scenarios involve learning from multiple teacher models with diverse expertise. For instance, a model trained on structured data may provide insights different from one specialized in unstructured data. Multi-teacher KD extends KD by leveraging this diversity to improve student learning. Recent research addresses this need with data-dependent weighting schemes that balance each teacher’s influence, enabling the student to capture multiple teachers’ strengths [8], [9].
Despite the empirical success of Multi-teacher KD, significant challenges remain. First, the statistical insights of KD and their effects on student-model improvement remain unclear. Although some studies [10], [11] have explored its impact, they primarily evaluate improvements in predictive performance. A deeper investigation into the statistical and theoretical aspects is crucial for understanding its full potential. Second, most existing methods prioritize improving prediction accuracy but lack quantification of uncertainty in its prediction. This limitation reduces their ability to assess the reliability of the student model’s predictions, restricting their applicability in high-stakes domains where interpretability and robustness are crucial. Moreover, current approaches often treat teachers independently, combining their outputs through weighted averages without considering potential interactions among the teachers. This simplification can lead to suboptimal knowledge transfer, preventing the student from fully leveraging the combined expertise of multiple teachers.
To address these challenges, we propose a Bayesian framework for multi-teacher knowledge distillation, introducing a teacher-informed mixture prior that integrates guidance from multiple teachers into one probabilistically coherent model. By integrating both teacher knowledge and task-specific data, this prior captures diverse external expertise while remaining closely aligned with the task. This integration of knowledge offers better generalization and robustness than purely data-driven priors used in classical Bayesian analysis, especially in the case of limited training data. The mixture prior also reflects teacher interactions, allowing the data to favor the teachers that best explain the observations naturally. Additionally, we employ an entropy-based weighting scheme that dynamically adjusts each teacher’s influence on a per-sample basis, ensuring that the most knowledgeable teacher has the greatest impact. Beyond improving the distillation process, our Bayesian formulation naturally equips the student model with uncertainty quantification capabilities. By sampling from the student’s posterior via stochastic gradient Langevin dynamics [12], we derive robust uncertainty measures and estimate credible intervals. These uncertainty estimates improve interpretability by quantifying the model’s confidence in its predictions, which is crucial for reliable decision-making in critical applications.
Our contributions are summarized as follows:
Uncertainty Quantification: We develop Bayesian inference tools for the student model to provide uncertainty estimates that improve interpretability and enable robust decision-making.
Teacher-informed Prior: We introduce a Bayesian KD framework that integrates the external knowledge of teacher models and training data to provide the prior, offering better generalization, robustness, and scalability.
Unified Multi-teacher Framework: We leverage a mixture distribution to combine the knowledge of multiple teachers, which captures their interactions, and ensures effective knowledge transfer. Our entropy-driven, sample-specific weighting mechanism prioritizes important teachers for each data point, optimizing the student’s learning process.
Validation: We demonstrate the effectiveness of our approach using synthetic and real-world datasets, highlighting its ability to enhance both predictive performance and uncertainty quantification.
Knowledge distillation (KD) has been widely studied for its effectiveness and limitations. The original work by [4] transfers knowledge by minimizing the KL divergence between teacher and student outputs. Later studies explored strategies like mimicking intermediate-layer representations [10], [13], [14]. To leverage diverse teacher expertise, researchers have proposed multi-teacher KD strategies. [8], [9], [15]–[18]. For instance, [8] align intermediate-layer representations across teachers using relative dissimilarity constraints and aggregate their guidance through a voting mechanism. [9] combine multi-teacher KD with adversarial training by employing specialized teachers for clean and adversarial data, enhancing both accuracy and robustness. However, existing multi-teacher KD approaches predominantly focus on improving predictive accuracy and often overlook interactions between teachers or the uncertainty in the student model’s predictions.
Efforts to incorporate Bayesian principles into KD have led to methods such as Bayesian Dark Knowledge [19], which distills the teacher model’s posterior predictive distribution into a compact student model, enabling uncertainty quantification. Subsequent extensions [20]–[22] refine this framework, but these methods require access to the teacher model’s posterior predictive distribution, which is often impractical for large-scale pre-trained models.
Formulating data-based priors has a long history in Bayesian statistics [23]. Some studies incorporate information from previous work to construct priors [24]. In particular, the power prior is derived by raising the likelihood of historical data to a specified power. However, its properness requires the covariate matrix of either historical or current data to have full column rank [25]. Another approach is to use synthetic data to define the prior distribution by training a simpler model and generating samples from its predictive distribution [26], [27]. In cases with limited expert knowledge, this method is effective. However, with powerful teacher models, MT-BKD can integrate both expert insights and current data to create the prior.
We begin by introducing the preliminaries of deep neural network models, followed by an overview of knowledge distillation and its extension to multi-teacher knowledge distillation.
Deep Neural Network Model. For the convenience of presentation, we introduce our method in the context of classification problems. In the classification problems, we are given the training sample \(\mathcal{D}=\{(\mathbf{x}_i, \mathbf{y}_i)\}_{i=1}^N\), where \(\mathbf{x}_i \in \Omega_\mathbf{x}\subset \mathbb{R}^m\) is the predictor and \(\mathbf{y}_i=(y_{i1},\ldots,y_{iK})^T\) is the corresponding label represented in the form of indicator vector. Specifically, if the \(i^{\text{th }}\) data point belongs to class \(k\), we have \(y_{i k}=1\) and \(y_{i j}=0\) for all \(j \neq k\), where \(i=1, \ldots, N, j=1, \ldots, K\). Each data point is independently collected across \(i \in \{1,\ldots,N\}\). Our goal is to find a function \(\boldsymbol{h}: \mathbb{R}^m \to [0,1]^K\) from a function class that approximates the conditional probability \(\mathbb{P}(y_{ik} = 1 \mid \mathbf{x}_i)\) for each \(k \in \{1, \ldots, K\}\). Specifically, we seek a function \(\boldsymbol{h}(\mathbf{x}_i; \boldsymbol{\theta}) = (h_1(\mathbf{x}_i; \boldsymbol{\theta}), \ldots, h_K(\mathbf{x}_i; \boldsymbol{\theta}))^T\), where \(h_k(\mathbf{x}_i; \boldsymbol{\theta})\) corresponds to the approximation of \(\mathbb{P}(y_{ik} = 1 \mid \mathbf{x}_i)\). The function class can be very general, ranging from logistic regressions [28], nonlinear regressions [29], to deep neural networks [30]. In this article, we focus on finding \(\boldsymbol{h}\) in the function class of deep neural networks, including feed-forward networks, convolutional neural nets, transformers, etc, and please see [31] for more details. For example, a feed-forward network utilizes functions from the following compositional function class: \[\label{dnn95def} \boldsymbol{h}(\mathbf{x}; \boldsymbol{\theta})=\boldsymbol{\sigma}_{L}\left(\mathbf{W}_L \boldsymbol{\sigma}_{L-1}\left(\mathbf{W}_{L-1} \cdots \boldsymbol{\sigma}_2\left(\mathbf{W}_2 \boldsymbol{\sigma}_1\left(\mathbf{W}_1 \mathbf{x}\right)\right)\right)\right),\tag{1}\] where \(\boldsymbol{\theta}=\left\{\mathbf{W}_1, \ldots, \mathbf{W}_L\right\}\) are the parameters of the model, and \(\boldsymbol{\sigma}_l\), \(l=1, \ldots, L\) are some fixed nonlinear activation functions. A popular choice of \(\boldsymbol{\sigma}_l(\cdot)\) for \(l=1, \ldots, L-1\) is the ReLU (Rectified Linear Unit) function, which sets negative values to zero while keeping positive values unchanged, applied element-wise in multivariate settings. In classification problems, for the final layer \(L\), researchers often use a soft-max function [32].
As \(\boldsymbol{h}(\cdot; \cdot)\) is determined by the neural network structure, the unknown parameter is \(\boldsymbol{\theta}\). To estimate \(\boldsymbol{\theta}\), we minimize the empirical risk: \[\label{eq:loss} \mathcal{L}(\boldsymbol{\theta}\mid \mathcal{D}) = \frac{1}{N}\sum_{i=1}^N \text{CE}(\mathbf{y}_i, \boldsymbol{h}(\mathbf{x}_i;\boldsymbol{\theta})),\tag{2}\] where \(\text{CE}(\mathbf{y}_i, \boldsymbol{h}(\mathbf{x}_i;\boldsymbol{\theta})) = - \sum_{k=1}^K y_{ik} \log\left( h_k(\mathbf{x};\boldsymbol{\theta}) \right)\) is the cross-entropy loss between the probability vector \(\boldsymbol{h}(\mathbf{x}_i;\boldsymbol{\theta})\) and \(\mathbf{y}_i\).
Knowledge Distillation and its Multi-Teacher Variant. Knowledge distillation (KD), introduced by [4], is a procedure where a simpler model (student) learns from a larger model (teacher). In the classical setup of KD, a pre-trained complex deep neural network model serves as the teacher. For each data point \(i\) in the training set \(\mathcal{D}\), the teacher model \(M_t\) outputs a probability vector \(\boldsymbol{p}_{i}=(p_{i1},\ldots p_{iK})^T\), where \(p_{ij}\) represents the predicted class probability that \(\mathbf{x}_i\) belongs to class \(j\). Mathematically, \(p_{ij}\) approximates \(\mathbb{P}(y_{ik} = 1 \mid \mathbf{x}_i)\) according to model \(M_t\). In this work, teacher models are treated as closed-source models, i.e., their structure and parameters are unknown. This is common in practice as many popular pre-trained large language models, e.g., ChatGPT and Claudia, are closed-source.
In contrast, the student model \(M_s\) is transparent, defined by a deep neural network with a known structure \(\boldsymbol{h}(\cdot ;\cdot)\) and unknown parameter \(\boldsymbol{\theta}\). The student model \(M_s\) is trained using both the training sample \({\cal D}\) and the corresponding teacher model’s predicted class probabilities \(\boldsymbol{p}= \{\boldsymbol{p}_{i}\}_{i=1,\cdots,N}\). The student model \(M_s\) outputs a probability vector \(\boldsymbol{q}_{i}=(q_{i1},\ldots q_{iK})^T\) for data point \(i\), where \(\boldsymbol{q}_{i} = \boldsymbol{h}(\mathbf{x}_i;\boldsymbol{\theta})\) and \(q_{ij}\) represents the predicted class probability that \(\mathbf{x}_i\) belongs to class \(j\). KD measures the discrepancy between the student and teacher models with \[\label{eq:student95teach95discrepancy} \begin{align} \tilde{\mathcal{L}}( \boldsymbol{\theta}\mid \mathcal{D}, \boldsymbol{p}) &= \frac{1}{N}\sum_{i=1}^N \text{CE}(\boldsymbol{p}_i, \boldsymbol{h}(\mathbf{x}_i;\boldsymbol{\theta})) = - \frac{1}{N}\sum_{i=1}^N\sum_{k=1}^K p_{ik} \log\left( h_k(\mathbf{x}_i;\boldsymbol{\theta}) \right), \end{align}\tag{3}\] which is the sample mean of the cross-entropy \(\text{CE}(\boldsymbol{p}_i, \boldsymbol{h}(\mathbf{x}_i;\boldsymbol{\theta}))\) across the training data. To leverage the information from both the training data and the teacher model’s predictions, KD estimates the parameter \(\boldsymbol{\theta}\) by minimizing the objective function, \[\label{eq:opt95KD95loss} \begin{align} \mathcal{L}^{\mathrm{KD}}( \boldsymbol{\theta}\mid \mathcal{D}, \boldsymbol{p}, \lambda) = \mathcal{L}( \boldsymbol{\theta}\mid \mathcal{D} ) + \lambda\tilde{\mathcal{L}}( \boldsymbol{\theta}\mid \mathcal{D}, \boldsymbol{p}) , \end{align}\tag{4}\] where \(\mathcal{L}^{\mathrm{KD}}( \boldsymbol{\theta}\mid \mathcal{D}, \boldsymbol{p}, \lambda)\) combines the losses in Equation (2 ) and Equation (3 ) with a tuning parameter \(\lambda \in (0,\infty)\) to balance their contributions.
In practice, a student can learn from multiple teachers, with the collective guidance of these teachers providing a comprehensive framework beneficial for training the student network. Consider \(G\) pre-trained teacher models, denoted by \(M_t^{(g)}\) for \(g=1,\cdots,G\), we have the predicted class probabilities, denoted by \(\boldsymbol{p}^{(g)} = \{\boldsymbol{p}_{i}^{(g)}\}_{i=1,\cdots,N}\), where \(\boldsymbol{p}_i^{(g)}=(p_{i1}^{(g)},\ldots p_{iK}^{(g)})^T = M_t^{(g)}(\mathbf{x}_i)\). A natural way to formulate the multi-teacher KD loss is: \[\label{eq:opt95multi95KD95loss} \begin{align} \mathcal{L}^{\operatorname{Multi-KD}}( \boldsymbol{\theta}\mid \mathcal{D}, \boldsymbol{p}, \lambda) = \mathcal{L}( \boldsymbol{\theta}\mid \mathcal{D} ) + \lambda \sum_{g=1}^G w^{(g)}\tilde{\mathcal{L}}^{(g)}( \boldsymbol{\theta}\mid \mathcal{D}, \boldsymbol{p}^{(g)} ), \end{align}\tag{5}\] where \(w^{(g)}\in (0,\infty)\) is the weight assigned to \(M_t^{(g)}\), \(\sum_{g=1}^G w^{(g)}=1\), and \(\tilde{\mathcal{L}}^{(g)}( \boldsymbol{\theta}\mid \mathcal{D}, \boldsymbol{p}^{(g)} )\) is analogously defined as in Equation (3 ). Extensions, including assigning sample-wise weights and distilling intermediate-layer representations, have also been explored [8], [33], [34].
While existing research has advanced our understanding of KD and its extensions to multi-teacher scenarios, significant gaps remain in capturing interactions among teachers and providing uncertainty quantification for the model predictions. In this work, we bridge these gaps by introducing a Bayesian framework for multi-teacher KD that explicitly models teacher-student interactions through a mixture prior, enabling both robust knowledge integration and uncertainty-aware learning. We provide details in the following section.
Unlike the previous formulations of multi-teacher KD, in this work, we develop multi-teacher KD in a Bayesian paradigm, which naturally enables uncertainty quantifications.
In our proposed Bayesian model, since each observation belongs to exactly one class, we assume that the class label \(\mathbf{y}_i\) follows a multinomial distribution given the predictor \(\mathbf{x}_i\), \[\label{eq:32model} \mathbf{y}_i | \mathbf{x}_i \sim \text{Multinomial}(1, \boldsymbol{q}_i),\tag{6}\] where \(i=1, \ldots, N\), and the multinomial probability is provided by a deep neural network, i.e., \(\boldsymbol{q}_i=\boldsymbol{h}(\mathbf{x}_i;\boldsymbol{\theta})\).
In the Bayesian paradigm, parameter \(\boldsymbol{\theta}\) is assumed to be random. Our goal is to develop a proper prior distribution. A straightforward approach might be to directly place a prior on \(\boldsymbol{\theta}\) as commonly did in Bayesian generalized linear models, see Chapter 14 in [35]. However, in our study, the teacher models are close-source, meaning the teacher models’ architectures and parameters are inaccessible, with only access to their outputs. Even in the case that the teacher models are open-source, since the teacher and student models differ in structure, the teacher models’ knowledge cannot be directly encoded into a prior for the parameters of the student model. Without any further information, the only option for us to impose objective (aka noninformative) prior for \(\boldsymbol{\theta}\), which is not satisfiable. As an alternative, we propose using the teacher models’ predicted class probabilities to build the prior, referred to as the teacher-informed prior (TIP), for \(\boldsymbol{\theta}\). We set the following mixture distribution [36] as our TIP for \(\boldsymbol{\theta}\), \[\label{eq:prior95KD95theta} \pi(\boldsymbol{\theta}\mid\{\boldsymbol{p}_{i}^{(g)}\}_{i=1,\cdots,N}^{g=1,\cdots,G}) \propto \prod_{i=1}^N \sum_{g=1}^G w_i^{(g)} f( \boldsymbol{h}(\mathbf{x}_i; \boldsymbol{\theta})\mid\boldsymbol{p}_{i}^{(g)}),\tag{7}\] where \(\propto\) denotes proportionality, \(G\) is the number of mixture components, \(w_i^{(g)}\) are the weights assigned to the component \(g\) on data point \(\mathbf{x}_i\), \(f(\cdot\mid\boldsymbol{p}_{i}^{(g)})\) are probability density functions. In this work, we consider \(f(\cdot\mid\boldsymbol{p}_{i}^{(g)})\) is the density function of a Dirichlet distribution \(Dir(\mathbf{1}_K+\lambda \boldsymbol{p}_{i}^{(g)})\), where \(\lambda\) is a tuning parameter weighting the contribution of teacher model predictions. Hence, we have \[\label{eq:prior95KD95q} f(\boldsymbol{q}_i \mid\boldsymbol{p}_{i}^{(g)}) = \frac{1}{B(\mathbf{1}_K+\lambda \boldsymbol{p}_{i}^{(g)})}\prod_{k=1}^K (q_{ik})^{\lambda p_{ik}^{(g)}},\tag{8}\] where \(B(\cdot)\) is the multivariate Beta function. Note that the mode of \(f(\cdot \mid\boldsymbol{p}_{i}^{(g)})\) is \(\boldsymbol{p}_{i}^{(g)}\) for \(\lambda > 0\). Furthermore, we show that \(\pi(\boldsymbol{\theta}\mid\{\boldsymbol{p}_{i}^{(g)}\}_{i=1,\cdots,N}^{g=1,\cdots,G})\) is a proper prior with the following proposition. We relegate the proof to Appendix 12.
Proposition 1. Consider the probability density function \(f(\boldsymbol{q}\mid\boldsymbol{p}_{i}^{(g)})\) as defined in Equation (8 ) with a constant \(\lambda > 0\). Assuming that the parameters of the student model lie in a compact space, then \(\pi(\boldsymbol{\theta}\mid\{\boldsymbol{p}_{i}^{(g)}\}_{i=1,\cdots,N}^{g=1,\cdots,G})\) is a proper prior.
According to our best knowledge, the proposed TIP is novel in the statistical literature. The proposed TIP has four distinguishing features: (1) Broad knowledge integration. It is a prior integration of the information from both teacher models and task-specific training data. The teacher models are typically pre-trained using extraneous data, and our training data are fed into the teacher models to yield predicted probabilities. Therefore, it captures extensive external knowledge while maintaining direct relevance to the task. (2) Statistical robustness. It reduces overfitting risk by incorporating external information, making it more robust than purely data-driven priors used in classical Bayesian analysis. This is particularly advantageous with limited training data. (3) Adaptability. It offers adaptability by incorporating multiple teacher models with automatic weight adjustment based on model quality, and the prior can be updated when better teacher models become available. (4) Practically user-friendly. It relies less on expert elicitation and is scalable to different domains, including texts, images, and numerical data. The key innovation is the principled integration of external knowledge (LLMs) with empirical data, offering better generalization than either approach alone.
In the general case with multiple teachers, the posterior density of \(\boldsymbol{\theta}\) is,
\[\begin{align} \label{eq:post} \pi(\boldsymbol{\theta}\mid\mathcal{D}, \boldsymbol{p}, \lambda, \boldsymbol{h}(\cdot;\cdot) ) & \propto \exp\left\{ \sum_{i=1}^N \sum_{k=1}^K y_{i k} \log \left(h_k\left(\mathbf{x}_i; \boldsymbol{\theta}\right)\right) \right. \nonumber\\ & \left. + \sum_{i=1}^N \log \left[\sum_{g=1}^G w_i^{(g)} \frac{1}{B\left(\mathbf{1}_k+\lambda \boldsymbol{p}_i^{(g)}\right)} \prod_{k=1}^K\left(h_k\left(\mathbf{x}_i; \boldsymbol{\theta}\right)\right)^{\lambda p_{ik}^{(g)}}\right] \right\}. \end{align}\tag{9}\]
Please note that when \(G=1\), Equation (9 ) reduces to a single-teacher case [37]. We present the discussion of this special case in Appendix 11.1.
Applying the mixture prior defined in Equation (7 ) leads to improved results compared to the traditional multi-teacher knowledge distillation (KD) strategy described in Equation (5 ). Specifically, the traditional multi-teacher KD strategy is equivalent to independently deriving posterior distributions from each teacher model and combining them using logarithmic opinion pooling [38], [39]. We establish the following theorem, with the proof presented in Appendix 12.
Theorem 2. Given the prior distribution defined in Equation (7 ) and Equation (8 ), the traditional multi-teacher KD strategy in Equation (5 ) is equivalent to deriving each teacher model’s posterior independently and combining them using logarithmic opinion pooling.
Traditional multi-teacher KD algorithms typically provide only point estimates, overlooking the uncertainty in model parameters. In contrast, our proposed MT-\(BKD\) method estimates the full posterior distribution, equipping the student model with uncertainty quantification capabilities. This is particularly important when the posterior distribution is multimodal. Point estimates in such cases may fall between modes, failing to capture the true distribution and potentially leading to misleading conclusions. In addition, we employ a mixture prior approach in MT-\(BKD\). We begin with a mixture prior with fixed weights representing our initial beliefs about each component. Once data is observed, we compute the posterior distribution based on this mixture prior. Although the initial prior weights remain fixed, the posterior naturally shifts in favor of components that better explain the data, balancing our assumptions with empirical evidence. If multiple priors fit the data well, the posterior may be multimodal, reflecting multiple plausible parameter configurations. By maintaining a mixture of priors, we preserve these distinct possibilities while allowing the data to emphasize those that are most consistent. Conversely, combining individual posteriors via logarithmic opinion pooling first updates each prior with data to obtain individual posteriors and then combines them via a weighted geometric mean of their log-transformed values. As a result, each prior retains its predetermined weight in the final combined posterior, regardless of its performance on the data.
We illustrate the advantage of MT-BKD using a toy example with two Gaussian priors centered at \(\mu_1=2.0\) and \(\mu_2=6.0\) (\(\sigma^2=1.0\)). The observed data has a mean of \(3.5\), which lies between these priors, indicating no strong preference for either. Figure 2 shows that MTBKD (red line) yields a bimodal posterior, accurately reflecting the balanced support for both priors while slightly favoring the one better supported by the data. In contrast, traditional multi-teacher KD (green line) collapses to a point estimate, overlooking the bimodal nature implied by the priors and the data.
Equation (7 ) highlights the critical role of the weights \(w_i^{(g)}\) in distilling knowledge from multiple teacher models into a unified student model. Since teacher networks vary in their effectiveness across different data, identifying and leveraging these differences is essential for successfully integrating multiple teachers. We define \(w_i^{(g)}\) for the teacher model \(M_t^{(g)}\) as \[\begin{align} \label{eq:weight} w_i^{(g)} & = \frac{ \xi(H(\boldsymbol{p}_i^{(g)})) }{ \sum_{g=1}^G \xi(H(\boldsymbol{p}_i^{(g)})) } , \end{align}\tag{10}\] where \(H(\cdot)\) represents the entropy, calculated as \(H(\boldsymbol{p}_i^{(g)}) = -\sum_{k} p_{ik}^{(g)} \log(p_{ik}^{(g)})\), and \(\xi\) is a monotonically descreasing function. Intuitively, from the perspective of information theory, information entropy is used to describe the degree of randomness and uncertainty of information, so it is reasonable to apply information entropy to quantify the knowledge scale of the teacher model on the given data point. Thus, we assign higher distillation weights to teacher models with lower entropy predictions. This allows the student model to learn more strongly from the teachers that have developed greater expertise on each specific data point, leading to more effective knowledge transfer.
We now discuss the effect of \(\lambda\). In the extreme case of \(\lambda = 0\), \(f(\boldsymbol{q}\mid\boldsymbol{p}_{i}^{(g)})\) becomes the density of a symmetric Dirichlet distribution, \(Dir(\mathbf{1}_K)\). This implies that there is no prior knowledge favoring one category over another when classifying each data point. For the case of \(\lambda \to \infty\), we derive the following proposition with proof in Appendix 12,
Proposition 3. Consider the probability density function \(f(\boldsymbol{q}\mid\boldsymbol{p}_{i}^{(g)})\) as defined in Equation (8 ), as \(\lambda \to \infty\), we have \(f(\boldsymbol{q}\mid\boldsymbol{p}_{i}^{(g)}) \longrightarrow \delta(\boldsymbol{q}-\boldsymbol{p}_{i}^{(g)}) ,\) where \(\delta(\cdot)\) is the multivariate Dirac delta function.
The above proposition shows that, as \(\lambda\) approaches infinity, the student model assigns nonzero probability only when its predicted class probabilities exactly match the teacher’s, aligning with original knowledge distillation (KD) that fully trusts the teacher’s probabilities at infinitely large \(\lambda\). In summary, a large \(\lambda\) indicates strong reliance on the teacher’s predictions, while a small \(\lambda\) suggests limited teacher influence. In practice, \(\lambda\) can be determined by selecting candidate values based on domain expertise and selecting the best one through cross-validation [40].
In practice, it is very common to use the trained student model to analyze confidential data. Analyzing confidential data using pre-trained teacher models as presented in the previous section may result in information leaking [41]. In this section, we develop a method to quantify prediction uncertainty using the trained student model alone.
For a new data point \(\tilde{\mathbf{x}}\), we denote the trained student model’s prediction as \(\tilde{\boldsymbol{q}}=\boldsymbol{h}(\tilde{\mathbf{x}}; \boldsymbol{\theta})\). Suppose that the label of \(\tilde{\mathbf{x}}\) is \(\tilde{\mathbf{y}}\). Instead of estimating the predictive posterior distribution [42] of the unobserved label \(\tilde{\mathbf{y}}\), we introduce a single-valued metric to gauge our confidence in the model’s prediction.
A commonly used confidence measure for classification is the deviance [43], defined as \(\operatorname{dev}\left(\tilde{\mathbf{y}}, \tilde{\boldsymbol{q}}\right) = -2 \sum_{k=1}^K \tilde{y}_{k} \log(\tilde{q}_{k})\), where \(\tilde{y}_{k}\) and \(\tilde{q}_{k}\) are the \(k\)th elements of \(\tilde{\mathbf{y}}\) and \(\tilde{\boldsymbol{q}}\). This deviance can be interpreted as a generalized sum of squared residuals for \(\tilde{\boldsymbol{q}}\). A higher deviance value suggests greater uncertainty in the model’s predictions. To make predictive inferences about \(\operatorname{dev}\left(\tilde{\mathbf{y}}, \tilde{\boldsymbol{q}}\right)\), we use its posterior predictive distribution. Under the multinomial model in Equation (6 ) and the posterior distribution of \(\boldsymbol{\theta}\) in Equation (9 ), we have \[\label{eq:32post95pred95dev} \pi \left(\operatorname{dev}\left(\tilde{\mathbf{y}}, \tilde{\boldsymbol{q}}\right) | \mathcal{D}, \boldsymbol{p}, \lambda, \boldsymbol{h}(\cdot;\cdot) \right) = \int \pi( \boldsymbol{\theta}| \mathcal{D}, \boldsymbol{p}, \lambda, \boldsymbol{h}(\cdot;\cdot) ) p\left(\operatorname{dev}\left(\tilde{\mathbf{y}}, \tilde{\boldsymbol{q}}\right) | \boldsymbol{\theta}\right) d \boldsymbol{\theta},\tag{11}\] where \(p\left(\operatorname{dev}\left(\tilde{\mathbf{y}}, \tilde{\boldsymbol{q}}\right) | \boldsymbol{\theta}\right)\) is the conditional probability density function of \(\operatorname{dev}\left(\tilde{\mathbf{y}}, \tilde{\boldsymbol{q}}\right)\) given \(\boldsymbol{\theta}\). Finally, the posterior mean deviance is used to quantify the model’s prediction uncertainty.
For the high dimension of the parameter \(\boldsymbol{\theta}\) and the complex structure of \(\boldsymbol{h}(\cdot ; \boldsymbol{\theta})\), it may not be feasible to obtain an analytical solution for Equation (11 ). Hence, we apply the stochastic Gradient Langevin Dynamics (SGLD) [12], a variant of Langevin Monte Carlo, to efficiently sample from the complex, high-dimensional posterior distributions. SGLD achieves this by seamlessly integrating the general approach of stochastic gradient descent with Langevin dynamics.
Suppose \(\boldsymbol{\theta}^{\langle j-1 \rangle}\) is sampled in iteration \(j-1\). Then in the \(j\)-th step, given a mini-batch of \(m\) data points \(\mathcal{D}^{\langle j \rangle} = \{(\mathbf{x}_i^{\langle j \rangle}, \mathbf{y}_i^{\langle j \rangle})\}_{i=1,\cdots,m}\) and the class probability \(\{\boldsymbol{p}_i^{(g)\langle j \rangle}\}_{i=1,\cdots,m}^{g=1,\cdots,G}\) predicted by the teacher model, SGLD generates the sample from posterior using gradient updates plus Gaussian noise,
\[\begin{align} \label{eq:SGLD} &\boldsymbol{\theta}^{\langle j \rangle}\nonumber = \nonumber \boldsymbol{\theta}^{\langle j-1 \rangle}+\tau \nabla_{\boldsymbol{\theta}} log \pi\left(\boldsymbol{\theta}^{(j-1)} \mid \mathcal{D}^{\langle j \rangle},\{\boldsymbol{p}_i^{(g)\langle j \rangle}\}_{i=1,\cdots,m}^{g=1,\cdots,G} , \lambda, \boldsymbol{h}\right) + \sqrt{2\tau} \xi^{\langle j \rangle}, \end{align}\tag{12}\]
where \(\tau\) is the step size, and \(\xi^{\langle j \rangle}\) is randomly sampled from \(N(0, I)\). SGD optimizes the log-likelihood to guide the sampling process toward regions of higher probability, while Langevin dynamics introduces controlled noise into the parameter updates, ensuring convergence to the full posterior distribution rather than merely its mode. Specifically, in the limit of \(j \rightarrow \infty\) and \(\tau \to 0\), the probability distribution of \(\boldsymbol{\theta}^{\langle j \rangle}\), denoted as \(\rho^{\langle j \rangle}\), converges to the posterior distribution of \(\boldsymbol{\theta}\). By using gradient information and introducing controlled noise, SGLD becomes more efficient in handling high-dimensional data [44]. Furthermore, SGLD’s use of mini-batches for gradient computation alleviates the computational burden associated with optimizing over entire datasets, making it particularly well-suited for large-scale datasets.
Since the distribution of MC sample \(\boldsymbol{\theta}^{\langle j \rangle}\) converges to the posterior distribution of \(\boldsymbol{\theta}\) as \(j\to \infty\), it allows for a precise estimation of the characteristics of \(\boldsymbol{\theta}\). Consequently, it fosters a comprehensive analysis and understanding of our model covering various aspects. One aspect that we are particularly interested in is the model’s prediction on new data.
With the posterior predictive distribution of the deviance in Equation (11 ), we propose the posterior mean deviance as a measure of predictive uncertainty and provide the credible interval of deviance.
With the MC sample \(\{ \boldsymbol{\theta}^{\langle j \rangle} \}_{j=1}^r\), we can estimate the posterior mean of \(\operatorname{dev}\left(\tilde{\mathbf{y}}, \tilde{\boldsymbol{q}}\right)\) by,
\[\begin{align}\label{eq:mean95dev} \tilde{\Delta} &= \frac{1}{r}\sum_{j=1}^{r} E\left( \operatorname{dev}(\tilde{\mathbf{y}},\tilde{\boldsymbol{q}}) \mid \boldsymbol{\theta}^{\langle j \rangle} \right)\nonumber\\ &= -\frac{2}{r}\sum_{j=1}^{r} \sum_{k=1}^K \tilde{q}_{k}^{\langle j \rangle} \log(\tilde{q}_{k}^{\langle j \rangle}) , \end{align}\tag{13}\]
where \(\tilde{q}_{k}^{\langle j \rangle}\) is the \(k\)th element of \(\tilde{\boldsymbol{q}}^{\langle j \rangle} = \boldsymbol{h}(\tilde{\mathbf{x}};\boldsymbol{\theta}^{\langle j \rangle})\).
Now, we construct the credible interval for the deviance of the prediction. For the deviance \(\operatorname{dev}\left(\tilde{\mathbf{y}}, \tilde{\boldsymbol{q}}\right)\), its \(1-\alpha\) credible interval, denoted by \(CI = [0,\tau]\), satisfies
\[\label{eq:CI} E( \boldsymbol{1}_{\{ \operatorname{dev}\left(\tilde{\mathbf{y}}, \tilde{\boldsymbol{q}}\right) \leq \tau \}} ) = 1-\alpha ,\tag{14}\]
where \(\alpha\) is the credible level. We can empirically estimate \(\tau\) by minimizing \(L(\tau)\) where
\[\begin{align} \label{eq:empirical} L(\tau) =& \left| \frac{1}{r} \sum_{j=1}^{r} E( \boldsymbol{1}_{\{ \operatorname{dev}\left(\tilde{\mathbf{y}}, \tilde{\boldsymbol{q}}\right) \leq \tau \}} \mid \boldsymbol{\theta}^{\langle j \rangle} )-(1-\alpha) \right| \nonumber \\ =& \left| \frac{1}{r} \sum_{j=1}^{r}\sum_{k=1}^K \tilde{q}_k^{\langle j \rangle} \boldsymbol{1}_{\{ -2\log(\tilde{q}_{k}^{\langle j \rangle}) \leq \tau \}} -(1-\alpha) \right|. \end{align}\tag{15}\]
The function \(L(\tau)\) denotes the difference between the left-hand side and the right-hand side of Equation (14 ).
Constructing a credible interval on deviance, rather than directly on the prediction \(\boldsymbol{q}\), presents a notable advantage. It enables the calculation of the coverage rate of the constructed interval using the true label. This is beneficial since, in most cases, the true labels of the testing dataset are readily available for evaluation, whereas the true class probabilities are typically elusive. The MT-BKD algorithm is summarized in Algorithm 3.
Here, we develop an asymptotic convergence of our proposed MT-BKD estimate. Without the loss of generality, we only consider the binary classification problem here. Let \(\mu\) represent the predictor \(\mathbf{x}\)’s true density function in the training sample of the student model while \(\mu^{(g)}\) represents the predictor’s true density function in the training sample of the \(g\)th teacher model for \(g = 1, \ldots, G\). Suppose the true parameter for the student model defined by Equation(1 ) and (6 ) is \(\boldsymbol{\theta}_0\). To define a good estimate, we consider the excess \(\phi\)-risk with the log loss, \[\label{eq:excess32phi32risk} {\cal E}^\phi(\widehat{\boldsymbol{\theta}}_n,\boldsymbol{\theta}_0) = {\cal E}^\phi(\widehat{\boldsymbol{\theta}}_n)-{\cal E}^\phi(\boldsymbol{\theta}_0),\tag{16}\] where \({\cal E}^\phi(\boldsymbol{\theta}) = E_{\mathbf{x},\mathbf{y}}[\phi(\mathbf{y}^T \text{logit}[\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta})] )]\), and \(\phi(x)= \log(1 + \exp(−x))\).
We define the quadratic functional \(V(f) = \int_{\Omega_\mathbf{x}} f^2 \mu(\mathbf{x}) d\mathbf{x}\).
The following conditions are assumed to hold.
(A1) For any \(\mathbf{x}\in \Omega_\mathbf{x}\), \(\mu(\mathbf{x})>0\), \(\max_{ 1\leq g \leq G} \mu^{(g)}(\mathbf{x}) >0\).
(A2) The minimizer \(\boldsymbol{\theta}^*_n\) in Equation (2 ) yield a consistent estimate of \(\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta})\), i.e., the error \(\int_{\Omega_\mathbf{x}} [\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}^*_n) - \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}) ]^2 \mu(\mathbf{x}) d\mathbf{x}= O_p(C_1(n))\) and \(\lim_{n \to \infty } C_1(n) = 0\).
(A3) There exist a positive constant \(\varepsilon^{(g)}\) such that \(\int_{\Omega_\mathbf{x}} [M_t^{(g)}(\mathbf{x}) - \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}) ]^2 \mu^{(g)}(\mathbf{x}) d\mathbf{x}\leq \varepsilon^{(g)}\) for \(g = 1,\ldots,G\).
Condition (A1) essentially requires that for any \(\mathbf{x}\in \Omega_\mathbf{x}\), there exists at least one of the training samples’ population distributions corresponding to the teacher models that assign a nonzero density to it. Condition (A2) can be proven under several regularity conditions, e.g., \(\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta})\) is a \((p,C)\)-smooth generalized hierarchical interaction model [45]. Without loss of generality, we use \(C_1(n)\) to represent its convergence rate while the detailed forms under different scenarios can be found in the existing literature [45]–[48]. Condition (A3) requires that the estimation errors of the teacher models be finite.
Theorem 4 (Convergence of MT-BKD). The estimate of \(\boldsymbol{\theta}\) by maximizing Equation(5 ) yields a consistent estimate of \(\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta})\) with the optimal convergence rate upper bounded by \[\label{eq:error32rate} C_2(n) \wedge G\cdot V \left( \left\{\sum_{g=1}^G \left[\frac{\mu^{(g)}(\mathbf{x})}{\mu(\mathbf{x}) \varepsilon^{(g)} } \right]^{\frac{1}{2}} \right\}^{-1}\right) ,\tag{17}\] where \(C_2(n) \asymp C_1(n)\).
We find the estimation error of MT-BKD could be potentially reduced compared to if the second term of Equation (17 ) is smaller than \(C_2(n)\). Smaller \(\varepsilon^{(g)}\) leads to a smaller value of the second term.
Remark. If \(\mu^{(g)} = \mu\) for \(g = 1,\ldots,G\), the second term of Equation (17 ) equals to \(G\left(\sum_{g=1}^G \sqrt{1/\varepsilon^{(g)}} \right)^{-2}\).
We conduct simulation studies to evaluate the performance of the proposed MT-BKD method in classification tasks under two scenarios with varying dimensions and class counts. Model performance is assessed using classification accuracy and mean square error (MSE) on class probability \(p\), where MSE for MT-BKD compares the predicted posterior mode with the true \(p\). Beyond overall accuracy, we focus on the model’s ability to quantify uncertainty. Mean deviance, as discussed in Section 5, is used to measure prediction uncertainty, and its relationship with the original data is examined. We further analyze coverage rates across different confidence levels and sample sizes, offering a comprehensive evaluation of the method’s uncertainty quantification.
We compare MT-BKD with four benchmark methods: (1) the teacher model; (2) the original KD method; (3) a combination of original KD with a Bayesian neural network (BNN) [49]; and (4) a combination of original KD with Monte Carlo dropout (Dropout) [50]. Details on the implementation of BNN and Dropout are provided in Appendix 13. For each distillation method, we evaluate both weighted distillation across multiple teachers, as described in Section 4.2, and distillation with equal teacher weights.
In this simulation, we consider binary classification problems. Let \(p(\mathbf{x}) = \frac{\exp(\eta(\mathbf{x}))}{1+\exp(\eta(\mathbf{x}))}\), and \(Y\sim Bernoulli(p(\mathbf{x}))\) with \(\eta(\mathbf{x}) = -0.5+0.1x_{1}+0.8x_{1}^2 - 0.8x_{2}^2.\) We generate a sample of size \(20,000\) by letting \(x_{1}\sim Unif(-3,3.5)\) and \(x_{2}\sim Unif(-3,3)\). The top left panel of Fig. 4 shows the probability distribution \(p(\mathbf{x})\) mapped across \(x_1\) and \(x_2\) coordinates, with probability values indicated by color intensity.
To simulate teacher specialization, we divide the dataset into three domains \(b_1\), \(b_2\), and \(b_3\), as shown in the top left panel of Fig. 4. We have Teacher 1 primarily trains on domain \(b_1\), with limited exposure to the other domains, while Teacher 2 specializes in \(b_3\). As shown in Table 1, Teacher 1 excels in \(b_1\), whereas Teacher 2 achieves high accuracy in \(b_3\). We then train a student model on data evenly drawn from all three domains, distilling knowledge from both teachers. Details on model structure are in Appendix 13.
Accuracy and MSE. Table 1 presents the overall and domain-specific classification accuracy of all methods. We have the following findings: (1) Despite their smaller model sizes, student models trained with KD methods achieve performance comparable to or even exceeding that of their teacher models. (2) MT-BKD outperforms original KD, Dropout, and BNN methods, achieving the highest overall accuracy (0.83) and excelling particularly in domains \(b_1\) (0.87) and \(b_3\) (0.85). (3) The proposed weighted distillation strategy consistently surpasses the equal-weight approach, highlighting its effectiveness in enhancing model performance. Additionally, MSE results for class probability \(p\) demonstrate that the proposed MT-BKD method outperforms all other methods, with the weighted version achieving the lowest MSE.
| Teachers | Orig KD | MT-BKD | Dropout | BNN | ||||||
| Acc | T1 | T2 | equal | weighted | equal | weighted | equal | weighted | equal | weighted |
| Overall | 0.66 | 0.67 | 0.72 | 0.72 | 0.79 | 0.83 | 0.72 | 0.73 | 0.77 | 0.73 |
| \(b_1\) | 0.84 | 0.41 | 0.77 | 0.75 | 0.83 | 0.87 | 0.75 | 0.76 | 0.78 | 0.78 |
| \(b_2\) | 0.78 | 0.79 | 0.66 | 0.69 | 0.78 | 0.78 | 0.67 | 0.68 | 0.80 | 0.80 |
| \(b_3\) | 0.39 | 0.83 | 0.74 | 0.74 | 0.79 | 0.85 | 0.73 | 0.73 | 0.76 | 0.63 |
| MSE | 0.14 | 0.14 | 0.08 | 0.08 | 0.02 | 0.01 | 0.08 | 0.08 | 0.07 | 0.07 |
Uncertainty quantification. We evaluate the inference performance of MT-BKD and compare it to Dropout and BNN. For clarity, we also include results from the single-teacher BKD method to assess individual teacher contributions. Fig. 4 compares the uncertainty quantification performance of various methods, with the top left subplot displaying the ground truth probability distribution. In this setting, where \(Y\sim Bernoulli(p)\), the variance of each observation is highly correlated with \(p(1-p)\), making the class probability \(p\) a key determinant of uncertainty. Predictions for data points with \(p\) near 0.5 are inherently more uncertain, while predictions for \(p\) close to 0 or 1 are more certain. Consequently, the white region (\(p\approx 0.5\)) in the top left subplot should exhibit the highest uncertainty, whereas the dark blue/red regions (\(p\approx 0/1\)) should show the lowest. Among the methods, weighted MT-BKD (top middle) aligns best with this expected uncertainty pattern.
The bottom-left subplots illustrate BKD’s performance with individual teachers. The student model trained solely with Teacher 1 exhibits clear uncertainty contrasts (yellow vs. dark purple) in domain \(b_1\) but blurred boundaries in \(b_3\). In contrast, the model trained with Teacher 2 shows sharp uncertainty patterns in \(b_3\) but weak distinction in \(b_1\), reflecting each teacher’s domain expertise. The weighted MT-BKD effectively integrates both teachers’ strengths, resulting in improved uncertainty estimation across all domains.
Coverage rate. With the credible interval constructed using Equation (15 ), We evaluate the coverage rate at three standard credible levels (0.85, 0.90, 0.95). As shown in Fig. 5, the weighted MT-BKD method achieves strong performance, with empirical coverage closely matching the nominal credible levels, indicating its effectiveness in quantifying uncertainty. A strong competitor is the weighted BNN method, which achieves coverage performance comparable to the weighted MT-BKD method.
In this simulation, we consider the following multi-class scenario. Let \(Y \sim Multinomial(\boldsymbol{p}(\mathbf{x}))\), and \(\boldsymbol{p}(\mathbf{x}) = (p_1(\mathbf{x}), p_2(\mathbf{x}), p_3(\mathbf{x}), p_4(\mathbf{x}), p_5(\mathbf{x}))^T\), where \(p_i(\mathbf{x}) = \frac{f(\mathbf{x}|\boldsymbol{\mu}_i,\boldsymbol{\Sigma})}{\sum_{j=1}^5f(\mathbf{x}|\boldsymbol{\mu}_1,\boldsymbol{\Sigma})},\;i=1,\ldots,5\). We have \(\boldsymbol{\mu}_1=(0,0,0,0,0)^T\), \(\boldsymbol{\mu}_2=(3,3,3,3,3)^T\), \(\boldsymbol{\mu}_3=(0,0,1,3,2)^T\), \(\boldsymbol{\mu}_4=(2,0,1,2,1)^T\), \(\boldsymbol{\mu}_5=(2,2,1,0,1)^T\), and \(\Sigma=(0.5^{|i-j|})_{5\times5}\). We simulate \(8,000\) sample from each \(N(\boldsymbol{\mu}_i,\Sigma)\), \(i=1,\ldots,5\), and generate labels accordingly.
| Teachers | Orig KD | MT-BKD | Dropout | BNN | |||||||
| Acc | T1 | T2 | T3 | equal | weighted | equal | weighted | equal | weighted | equal | weighted |
| Overall | 0.56 | 0.61 | 0.53 | 0.85 | 0.85 | 0.85 | 0.87 | 0.86 | 0.86 | 0.84 | 0.85 |
| \(b_1\) | 0.96 | 0.00 | 0.33 | 0.84 | 0.84 | 0.91 | 0.92 | 0.83 | 0.83 | 0.77 | 0.79 |
| \(b_2\) | 0.90 | 0.92 | 0.05 | 0.90 | 0.90 | 0.89 | 0.89 | 0.92 | 0.92 | 0.91 | 0.90 |
| \(b_3\) | 0.35 | 0.91 | 0.93 | 0.90 | 0.91 | 0.86 | 0.88 | 0.91 | 0.92 | 0.90 | 0.90 |
| \(b_4\) | 0.58 | 0.53 | 0.51 | 0.79 | 0.78 | 0.86 | 0.86 | 0.81 | 0.80 | 0.81 | 0.82 |
| \(b_5\) | 0.00 | 0.67 | 0.83 | 0.85 | 0.84 | 0.76 | 0.82 | 0.82 | 0.81 | 0.81 | 0.82 |
| MSE | 0.07 | 0.07 | 0.08 | 0.02 | 0.02 | 0.01 | 0.00 | 0.02 | 0.01 | 0.02 | 0.01 |
We consider three teacher models \(T_1\), \(T_2\), \(T_3\), each specializing in data drawn from \(N(\boldsymbol{\mu}_i,\Sigma)\) for \(i=1,2,3\). Table 2 shows the overall classification accuracy, MSE, and domain-specific performance on domain \(b_i\), which contains data generated from \(N(\boldsymbol{\mu}_i,\Sigma)\), \(i=1,\cdots,5\). The results align with those from Simulation 1, demonstrating that the proposed MT-BKD method outperforms all other approaches. In particular, the weighted distillation strategy for combining multiple teachers surpasses the equal-weight strategy in accuracy. The advantages of the MT-BKD method are further highlighted in its uncertainty quantification capabilities, as shown by the coverage rates in Fig. 5 (b). MT-BKD with the weighting strategy closely matches the assigned credible levels, whereas other methods exhibit significantly larger errors, with discrepancies of up to 10%.
We evaluate the performance of MT-BKD on two tasks: eukaryotic protein subcellular localization prediction and image classification. In each task, a compact student model is distilled from teacher models, and we report the student model’s uncertainty predictions. Details of the model architecture and training procedures are provided in Appendix 14.
We apply MT-BKD to predict eukaryotic protein subcellular localization, a key factor in understanding protein function and disease mechanisms. Protein localization provides the necessary context for function, while mislocalization can lead to diseases such as cancer, neurodegenerative, and metabolic disorders [51]. Accurate localization prediction is therefore vital for advancing both basic research and therapeutic development.
| Location | Number of proteins | Sublocations |
|---|---|---|
| Nucleus | 1243 | Envelope, inner and outer membrane, matrix, lamina, chromosome, |
| nucleus speckle | ||
| Cytoplasm | 955 | Cytoplasm (cytosol and cytoskeleton) |
| Extracellular | 717 | Extracellular |
| Mitochondrion | 515 | Envelope, inner and outer membrane, matrix, intermembrane space |
| Cell membrane | 619 | Apical, apicolateral, basal, basolateral, lateral, cell membrane, cell projection |
| Endoplasmic reticulum | 334 | ER membrane and lumen, microsome, rough ER, smooth ER, |
| Sarcoplasmic reticulum | ||
| Plastid | 247 | Plastid membrane, stroma and thylakoid |
| Golgi apparatus | 142 | Golgi apparatus membrane and lumen |
| Lysosome/Vacuole | 171 | Contractile, lytic and protein storage vacuole, vacuole lumen and membrane, |
| lysosome lumen and membrane | ||
| Peroxisome | 57 | Peroxisome matrix and membrane |
We analyze a subset of the DeepLoc2 dataset [52], derived from UniProt (release 2021_03) [53], which classifies proteins into ten subcellular locations: Cytoplasm, Nucleus, Extracellular, Cell Membrane, Mitochondrion, Plastid, Endoplasmic Reticulum, Lysosome/Vacuole, Golgi Apparatus, and Peroxisome. We focus on sequences from Metazoa, Fungi, or Viridiplantae that range from 40 to 1000 amino acids, excluding proteins assigned to multiple locations. From this filtered set, we randomly select 5,000 sequences. Table 3 summarizes the subcellular locations and corresponding sequence counts. The data is then split into training, validation, and test sets in a 6:2:2 ratio. We adopt two teacher models with high-capacity protein sequence representations: ESM-2-650M [54] (650 million parameters) and ProtT5-XL-half (120 million parameters), alongside a lightweight student model, ESM-2-8M (8 million parameters). A classifier is added after the feature extraction layers of each model to adapt their pre-trained embeddings to subcellular location prediction.
We explore the performance of MT-BKD in uncertainty quantification. Fig. 7 shows the log-transformed mean deviance distribution. In subfigure (a), the first box shows results from the test set of exclusively single-label proteins, aligning with the training data. The second box displays results from a separate multi-label protein dataset for comparison. As expected, uncertainty is higher for multi-label proteins. This increased uncertainty likely comes from the greater complexity of multiple-label proteins, making them harder to predict. Moreover, since our model was trained only on single-label proteins, reduced performance on multi-label data is unsurprising. The elevated uncertainty appropriately indicates the model’s recognition of its own predictive limitations. Furthermore, we evaluate the model’s prediction confidence across different classes. In Fig. 7 (b), each box corresponds to a specific class in the test set. MT-BKD exhibits the highest confidence for Extracellular, Nuclear, and Mitochondrial, and the lowest for Golgi apparatus, Lysosomes, and Peroxisomes. Such observations align with biological evidence that locations like Extracellular, Nuclear, and Mitochondrial are typically easier to identify with high confidence [55]–[57]. This is because of their distinct, well-characterized targeting signals that are consistent across species and highly conserved, i.e., they remain very similar or unchanged across different species over long evolutionary timescales. In particular, nuclear proteins contain one or more short, positively charged nuclear localization signals (NLS) recognized by importins [56], [58], while mitochondria-bound proteins typically possess an N-terminal targeting sequence [55], and extracellular proteins are generally marked by a clear N-terminal signal peptide that directs them through the secretory pathway [57]. On the contrary, locations such as the Golgi apparatus, lysosomes, and peroxisomes present challenges due to more diverse and less conserved signals. Specifically, Golgi proteins often undergo complex modification processes, such as glycosylation, which can vary between species and even cell types [59]. Lysosomal targeting relies on tags like mannose-6-phosphate, but not all lysosomal proteins consistently carry this motif [60]. Peroxisomal proteins are similarly difficult to predict due to a varied array of peroxisomal targeting signals (PTS1 and PTS2) with structural nuances [61]. These biological facts support the biological basis for varying certainty in subcellular localization predictions given by MT-BKD.
The uncertainty scores achieved by MT-BKD can be leveraged to prioritize experimental validation of borderline predictions, guide model refinement by highlighting where additional data or improved feature engineering is needed, and ultimately enhance our understanding of how subcellular localization signals shape protein distribution.
In this study, we utilized the Digit-Five dataset, which contains digit images from multiple domains, each presenting unique characteristics [62]. There are 10 classes corresponding to digits ranging from 0 to 9 in each domain. Specifically, we focused on three colored variants of the dataset, where each color scheme corresponds to a distinct domain, thereby introducing variations in data distribution and visual representation. These colored variants aim to simulate real-world challenges in domain adaptation by altering pixel intensity patterns, backgrounds, or foreground colors, which can significantly impact model performance.
Each teacher model was trained to specialize in one of the colored datasets, effectively becoming an expert in its respective domain. Using the proposed MT-\(BKD\) approach, we trained a single student model to generalize across all three color domains, enabling it to learn from the unique expertise of each teacher. The teacher models use a ResNet architecture with approximately 25.6 million parameters, while the student model employs a convolutional neural network with around 0.1 million parameters, which is 256 times fewer than the teacher models. Additional details about the data and models can be found in Appendix 14.
In Fig. 8, the images with the lowest and highest uncertainty are shown on the left and right panels, respectively. The images with higher uncertainty are noticeably more difficult to classify. The clear association between high uncertainty and the inherent difficulty in image recognition validates the performance of MT-\(BKD\) in uncertainty quantification.
In this work, we develop a Multiple-teacher Bayesian Knowledge Distillation (MT-\(BKD\)) framework aimed at compressing the teacher model and providing a suite of Bayesian inference tools for the student model’s uncertainty quantification. MT-\(BKD\) constructs a teacher-informed prior by integrating external knowledge from multiple teacher models with task-specific training data for the student model’s parameters, from which a posterior is derived. We show that minimizing KD loss is related to estimating the posterior mode in MT-\(BKD\), offering a clearer interpretation of how KD operates and suggesting new ways to enhance KD methods. The developed MT-BKD framework contributes to the growing demand for reliable, privacy-aware, and resource-efficient machine learning models in real-world applications.
We evaluate the performance of MT-\(BKD\) using both synthetic and real datasets, focusing on classification accuracy and uncertainty quantification ability. The results demonstrate that MT-\(BKD\) achieves performance comparable to the teacher model with the advantage of a significantly smaller size. A key strength of MT-\(BKD\) is its ability to precisely quantify prediction uncertainty, as demonstrated by three main findings. First, we confirm that mean deviance is a reliable metric for uncertainty by examining its relationship to the true class probability \(p\) and the data’s intrinsic structure (see Fig. 4). Second, an analysis of coverage rates highlights the robustness of MT-BKD. Third, in experiments on real datasets, visualizations and distribution analyses of uncertainty estimates across classes further affirm the reliability of MT-BKD’s uncertainty measures.
MT-\(BKD\) could be extended to regression and data generation tasks by imposing suitable teacher-informed priors. For regression, one can adopt a Gaussian prior whose mean aligns with the teacher models’ predictions. For score-based generative models, similarly, a Gaussian prior on the score function can be set to have its mean match the teacher-estimated scores.
APPENDIX
Here is an outline of the Appendix.
Appendix 10 summarizes the notations used in the main text.
Appendix 11 discusses the single teacher situation.
Appendix 12 gives the proof of theoretical results.
Appendix 13 presents supplementary details for simulation experiments.
Appendix 14 presents supplementary details for real data analysis.
In Table 4, we summarize the notations in the main text.
| Notation | Interpretation |
|---|---|
| \(\mathcal{F}\) | function classes |
| \(\mathbb{P}\) | population distribution |
| \(\mathcal{D}, \mathcal{T}\) | dataset |
| \(\mathcal{P}\) | distribution class |
| \(\boldsymbol{h}\) | function representing the student neural network |
| \(\boldsymbol{\theta}\) | parameters of the student neural network |
| \(\boldsymbol{\theta}^*\) | estimate of \(\boldsymbol{\theta}\) |
| \(\widehat\boldsymbol{\theta}\) | Monte Carlo sample of \(\boldsymbol{\theta}\) |
| \(\mathbf{x}\) | vector representing \(m\)-dimensional feature |
| \(y\) | class label |
| \(\mathbf{y}\) | one-hot encoding of class label \(y\) |
| \(\boldsymbol{p}\) | output of the tea neural network, representing the predicted class probability |
| \(\boldsymbol{q}\) | output of the student neural network, representing the predicted class probability |
| \(N\) | sample size of the training dataset \(\mathcal{D}\) |
| \(n\) | sample size of the testing dataset \(\mathcal{T}\) |
| \(l\) | likelihood function |
| \(\nabla\) | gradient operator |
| \(\lambda\) | weight parameter |
| \(\operatorname{dev}\) | deviance |
In the case where \(G=1\), meaning there is only a single teacher model, we present the following theorem, which establishes the connection between our established Bayesian posterior and the original KD framework.
Theorem 5. Assume the prior distributions are defined in Equation (7 ) and (8 ). When there is only one teacher model, represented by \(G = 1\), the posterior mode of \(\boldsymbol{\theta}\) is the minimizer of KD, i.e., \(\boldsymbol{\theta}_{KD}^*\).
Proof. The negative log-transformed posterior distribution of \(\boldsymbol{\theta}\) is
\[\begin{align} \label{eq:ll95q} -l(\boldsymbol{\theta}\mid\mathcal{D}, \boldsymbol{p}, \lambda, \boldsymbol{h}(\cdot;\cdot) ) =& -\sum_{i=1}^N \sum_{k=1}^K y_{ik}\log\left(h_k(\mathbf{x}_i;\boldsymbol{\theta})\right) \nonumber\\ &- \lambda \sum_{i=1}^N \sum_{k=1}^K p_{ik}\log \left(h_k(\mathbf{x}_i;\boldsymbol{\theta})\right) + c,\nonumber\\ =&\; \mathcal{L}^{\mathrm{KD}}( \boldsymbol{h}(\cdot , \boldsymbol{\theta})\mid \mathcal{D}, \boldsymbol{p}, \lambda) +c, \end{align}\tag{18}\]
where \(c\) is a constant. It is easy to verify that the posterior mode of \(\boldsymbol{\theta}\) is the minimizer of the KD loss in Equation (4 ). ◻
In other words, optimizing the KD loss to search for parameters \(\boldsymbol{\theta}\) in the student model is tantamount to achieving the maximum likelihood estimation of the parameters \(\boldsymbol{\theta}\).
Proposition 1. Consider the probability density function \(f(\boldsymbol{q}\mid\boldsymbol{p}_{i}^{(g)})\) as defined in Equation (8 ) with a constant \(\lambda > 0\). Assuming that the parameters of the student model lie in a compact space, then \(\pi(\boldsymbol{\theta}\mid\{\boldsymbol{p}_{i}^{(g)}\}_{i=1,\cdots,N}^{g=1,\cdots,G})\) is a proper prior.
Proof. According to Equation (7 ) and (8 ), for some constant \(C>0\), we have \[\begin{align} \pi(\boldsymbol{\theta}\mid\{\boldsymbol{p}_{i}^{(g)}\}_{i=1,\cdots,N}^{g=1,\cdots,G}) &= C\;\prod_{i=1}^N \sum_{g=1}^G f( \boldsymbol{h}(\mathbf{x}_i; \boldsymbol{\theta})\mid\boldsymbol{p}_{i}^{(g)}) \nonumber \\ &= C\;\prod_{k=1}^K \sum_{g=1}^G b_{ij} \prod_{i=1}^N \left(\boldsymbol{h}(\mathbf{x}_i; \boldsymbol{\theta}) \right) ^{\lambda p_{i k}^{(g)}}, \end{align}\] where \(b_{ig} \triangleq \frac{1}{B(\mathbf{1}_K+\lambda \boldsymbol{p}_{i}^{(g)})}>0\) is a constant. Taking integration, we get \[\label{eq:intg} \int \pi(\boldsymbol{\theta}\mid\{\boldsymbol{p}_{i}^{(g)}\}_{i=1,\cdots,N}^{g=1,\cdots,G}) d\boldsymbol{\theta} = C\int \prod_{k=1}^K \sum_{g=1}^G b_{ig} \prod_{i=1}^N \left(\boldsymbol{h}(\mathbf{x}_i; \boldsymbol{\theta}) \right)^{\lambda p_{i k}^{(g)}} d\boldsymbol{\theta}.\tag{19}\] Since \(0 \leq \boldsymbol{h}(\mathbf{x}_i; \boldsymbol{\theta}) \leq 1\), we have \(0 \leq \prod_{k=1}^K \sum_{g=1}^G b_{ig} \prod_{i=1}^N\left(\boldsymbol{h}(\mathbf{x}_i; \boldsymbol{\theta}) \right)^{\lambda p_{i k}^{(g)}} \leq 1\) for \(1\leq k \leq K\), \(1\leq g \leq G\) and \(1\leq i\leq N\). Given that \(\boldsymbol{\theta}\) lies in a compact space and \(\prod_{k=1}^K \sum_{g=1}^G b_{ig} \prod_{i=1}^N\left(\boldsymbol{h}(\mathbf{x}_i; \boldsymbol{\theta}) \right)^{\lambda p_{i k}^{(g)}}\) is bounded, we have \[\int \prod_{k=1}^K \sum_{g=1}^G b_{ig} \prod_{i=1}^N \left(\boldsymbol{h}(\mathbf{x}_i; \boldsymbol{\theta}) \right)^{\lambda p_{i k}^{(g)}} d\boldsymbol{\theta}< \infty.\] Thus, we can easily verify that \(\pi(\boldsymbol{\theta}\mid\{\boldsymbol{p}_{i}^{(g)}\}_{i=1,\cdots,N}^{g=1,\cdots,G})\) is a proper prior. ◻
Proposition 3. Consider the probability density function \(f(\boldsymbol{q}\mid\boldsymbol{p}_{i}^{(g)})\) as defined in Equation (8 ), as \(\lambda \to \infty\), we have \[f(\boldsymbol{q}\mid\boldsymbol{p}_{i}^{(g)}) \longrightarrow \delta(\boldsymbol{q}-\boldsymbol{p}_{i}^{(g)}) ,\] where \(\delta(\cdot)\) is the multivariate Dirac delta function.
Proof. For the sake of brevity, we consider the general formula \(f(\boldsymbol{q}\mid\boldsymbol{p},\lambda)\), where \(\boldsymbol{q}=(q_1,q_2,\cdots,p_K)^T\) and \(\boldsymbol{p}=(p_1,p_2,\cdots,p_K)^T\). WLOG, we consider the situation where \(p_k \neq 0\) for \(k \in 1,\cdots,K\). We have \[f(\boldsymbol{q}\mid\boldsymbol{p},\lambda) = \frac{1}{B(\mathbf{1}_K+\lambda \boldsymbol{p})}\Pi_{k=1}^K (q_{k})^{\lambda p_{k}} = \frac{\Gamma(\lambda + K)}{\prod_{k=1}^K\left(\Gamma(\lambda p_k+1)\right) }\Pi_{k=1}^K (q_{k})^{\lambda p_{k}}.\]
We start with the situation when \(\lambda p_k, k=1, \cdots, K\), are all integers. We have \[\label{suppeq:pdf} f(\boldsymbol{q}\mid\boldsymbol{p},\lambda)=\frac{(\lambda+K-1) !}{\prod_{k=1}^K\left(\lambda p_k\right) !} \cdot \prod_{k=1}^K\left(q_k\right)^{\lambda p_k}.\tag{20}\] According to the Stirling formula that \(n ! \sim \sqrt{2 \pi n}\left(\frac{n}{e}\right)^n\), we have \[\begin{align} f(\boldsymbol{q}\mid\boldsymbol{p},\lambda) &\sim \frac{\lambda^{K-1} \cdot \sqrt{\lambda}\left(\frac{\lambda}{e}\right)^\lambda}{\prod_{k=1}^K \sqrt{\lambda p_k}\left(\frac{\lambda p_k}{e}\right)^{\lambda p_k}} \cdot \prod_{k=1}^K\left(q_k\right)^{\lambda p_k} \\ &\sim \lambda^{\frac{K-1}{2}} \cdot \prod_{k=1}^K\left(\frac{q_k}{p_k}\right)^{\lambda p_k}. \end{align}\] Write \(g(\boldsymbol{q})=\prod_{k=1}^K\left(\frac{q_k}{p_k}\right)^{p_k}\), we now get the maximization of \(g(\boldsymbol{q})\) w.r.t. \(\boldsymbol{q}\). Since we have the constraint that \(\sum_{k=1}^K q_k=1\), we use the Lagrange multiplier to find the maximum of \(g(\boldsymbol{q})\). The Lagrangian function is defined as \[\mathcal{L}\left(q_1, \ldots, q_K, c\right) \equiv \log g+c\left(1-\sum_{k=1}^K q_k\right) ,\] where \(\log g=\sum_{k=1}^K p_k\left(\log q_k-\log p_k\right)\).
To solve \[\nabla_{q_1, \cdots, q_K, c} g\left(q_1, \cdots, q_K, c\right)=0,\] we have \[\left\{\begin{array}{l}\frac{p_k}{q_k}-c=0, \qquad \text{for } k=1, \ldots, K \\ \sum_{k=1}^K q_k=1\end{array}\right.\] since \(\sum_{k=1}^K p_k=1\), we can easily get that the solution is \[\label{suppeq:sol} \begin{align} & \left\{\begin{array}{l} q_k=p_k, \qquad \text{for } k=1, \cdots, K \\ c=1. \end{array}\right. \end{align}\tag{21}\] That is, we have \(g_{\text{max}}(\boldsymbol{q})=1\) when Equation (21 ) holds, and this leads to \(f(\boldsymbol{q}\mid\boldsymbol{p},\lambda) \sim \lambda^{\frac{K-1}{2}} \cdot 1^{\lambda} \stackrel{\lambda \rightarrow \infty}{\longrightarrow} \infty\).
For \(\forall \boldsymbol{q}\neq \boldsymbol{p}\), we have \(g(\boldsymbol{q}) < 1\), which leads to \(f(\boldsymbol{q}\mid\boldsymbol{p},\lambda) \sim \lambda^{\frac{K-1}{2}} \cdot \left(g(\boldsymbol{q})\right)^{\lambda} \stackrel{\lambda \rightarrow \infty}{\longrightarrow} 0\). Thus, as \(\lambda \rightarrow \infty\), we have \[\begin{align} & \lim _{\lambda \rightarrow \infty} \int f(\boldsymbol{q}\mid\boldsymbol{p},\lambda) \mathrm{d} \boldsymbol{q}= 1 , \\ & f(\boldsymbol{q}\mid\boldsymbol{p},\lambda) = 0 \text{ if } \boldsymbol{q}-\boldsymbol{p}\neq 0 . \end{align}\] That is, \(f(\boldsymbol{q}\mid\boldsymbol{p},\lambda) \longrightarrow \delta(\boldsymbol{q}-\boldsymbol{p})\) as \(\lambda \rightarrow \infty\).
If \(\exists \;k\) s.t. \(\lambda p_k\) is not an integer. According to the Stirling formula for the gamma function that \(\Gamma(z) \sim \sqrt{\frac{2 \pi}{z}}\left(\frac{z}{e}\right)^z\), following the same analysis, we can also get \(f(\boldsymbol{q}\mid\boldsymbol{p},\lambda) \longrightarrow \delta(\boldsymbol{q}-\boldsymbol{p})\) as \(\lambda \rightarrow \infty\). ◻
Theorem 2. Given the prior distribution defined in Equation (7 ) and (8 ), the traditional multi-teacher KD strategy in Equation (5 ) is equivalent to deriving the posterior distributions for individual teacher models separately and combining them using logarithmic opinion pooling.
Proof. Based on a single \(g^{\text{th}}\) teacher model, we have the prior distribution \(\pi_{(g)}(\boldsymbol{\theta}; \{\boldsymbol{p}_i^{(g)}\}_{i=1,\cdots,N}) \propto \prod_{i=1}^N f(\boldsymbol{h}(\mathbf{x}_i; \boldsymbol{\theta}); \boldsymbol{p}_i^{(g)})\), where \(g = 1, \dots, G\). The posterior density is given by
\[\begin{align} \pi_{(g)}(\boldsymbol{\theta}\mid\mathcal{D}, \boldsymbol{p}^{(g)}, \lambda, \boldsymbol{h}(\cdot,\cdot) ) & \propto \prod_{i=1}^N \prod_{k=1}^K h_k\left(\mathbf{x}_i; \boldsymbol{\theta}\right)^{y_{i k}} \cdot \prod_{i=1}^N \frac{1}{B\left(\mathbf{1}_k + \lambda\boldsymbol{p}_i^{(q)}\right)} \prod_{k=1}^K\left(h_k\left(\mathbf{x}_i; \boldsymbol{\theta}\right)\right)^{\lambda p_{i k}{ }^{(g)}}. \end{align}\]
Combining the posterior distributions based on each teacher model through a logarithmic pooling strategy, we have the combined posterior
\[\begin{align} \pi_{\text{pooled}}(\boldsymbol{\theta}\mid\mathcal{D}, \boldsymbol{p}, \lambda, \boldsymbol{h}(\cdot,\cdot)) \propto \Pi_{g=1}^G \pi_{(g)}(\boldsymbol{\theta}\mid\mathcal{D}, \boldsymbol{p}^{(g)}, \lambda, \boldsymbol{h}(\cdot,\cdot))^{w^{(g)}}, \end{align}\]
where \(w^{(g)}\) is the weight for the \(g^{\text{th}}\) teacher model. Thus, the negative log-transformed posterior distribution of \(\boldsymbol{\theta}\) is
\[\begin{align} -l(\boldsymbol{\theta}\mid\mathcal{D}, \boldsymbol{p}, \lambda, \boldsymbol{h}(\cdot,\cdot) ) =& -\sum_{i=1}^N \sum_{k=1}^K y_{ik}\log\left(h_k(\mathbf{x}_i;\boldsymbol{\theta})\right) \nonumber\\ &- \lambda \sum_{g=1}^G w^{(g)} \sum_{i=1}^N \sum_{k=1}^K p_{ik}^{(g)}\log \left(h_k(\mathbf{x}_i;\boldsymbol{\theta})\right) + c,\nonumber\\ =&\; \mathcal{L}^{\mathrm{Multi-KD}}( \boldsymbol{h}(\cdot , \boldsymbol{\theta})\mid \mathcal{D}, \boldsymbol{p}, \lambda) +c, \end{align}\]
where \(c\) is a constant. That is, finding the mode of the combined posterior, obtained by merging the posterior distributions from each teacher model using a logarithmic pooling strategy, is equivalent to minimizing the multi-teacher KD loss described in Equation (5 ). ◻
Theorem 4. The estimate of \(\boldsymbol{\theta}\) by maximizing Equation (5 ) yields a consistent estimate of \(\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta})\) with the optimal convergence rate upper bounded by \[C_2(n) \wedge G\cdot V \left( \left\{\sum_{g=1}^G \left[\frac{\mu^{(g)}(\mathbf{x})}{\mu(\mathbf{x}) \varepsilon^{(g)} } \right]^{\frac{1}{2}} \right\}^{-1}\right) ,\] where \(C_2(n) \asymp C_1(n)\).
(A1) For any \(\mathbf{x}\in \Omega_\mathbf{x}\), \(\mu(\mathbf{x})>0\), \(\max_{ 1\leq g \leq G} \mu^{(g)}(\mathbf{x}) >0\).
(A2) The minimizer \(\widehat{\boldsymbol{\theta}}_n\) in Equation (2 ) yield a consistent estimate of \(\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta})\), i.e., the error \(\int_{\Omega_\mathbf{x}} [\boldsymbol{h}(\mathbf{x};\widehat{\boldsymbol{\theta}}_n) - \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}) ]^2 \mu(\mathbf{x}) d\mathbf{x}= O_p(C_1(n))\) and \(\lim_{n \to \infty } C_1(n) = 0\).
(A3) There exist a positive constant \(\varepsilon^{(g)}\) such that \(\int_{\Omega_\mathbf{x}} [M_t^{(g)}(\mathbf{x}) - \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}) ]^2 \mu^{(g)}(\mathbf{x}) d\mathbf{x}\leq \varepsilon^{(g)}\) for \(g = 1,\ldots,G\).
Proof. Let \(\widehat{\boldsymbol{\theta}}_n\) be the maximize of Equation (5 ). We define \[\begin{align} {\cal E}_n^{\varphi}(\boldsymbol{\theta})=\frac{1}{N} \sum_{i=1}^N \varphi\left(\mathbf{y}_i^T \text{logit}[\boldsymbol{h}(\mathbf{x}_i;\boldsymbol{\theta})]\right) \end{align}\] Notice that \({\cal E}_n^{\varphi}(\boldsymbol{\theta}) = {\cal L}(\boldsymbol{\theta}|{\cal D})\) and \[\begin{align} {\cal E}^{\varphi}(\boldsymbol{\theta}) &=& E_\mathbf{x}(KL[\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0) ,\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}) ])\\ &=& E_{\mathbf{x},\mathbf{y}}(CE[\mathbf{y},\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}) ]) - E_{\mathbf{x}}(CE[\boldsymbol{h}(\mathbf{x}_0;\boldsymbol{\theta}),\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0)]) \end{align}\]
For \(g\)th teacher \[\begin{align} CE[\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0),\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta})]) &\leq& CE[M_t^{(g)}(\mathbf{x}),\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta})] + \|M_t^{(g)}(\mathbf{x}) - \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0) \|_2\|\log \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}) \|_2 \\ &\leq& CE[M_t^{(g)}(\mathbf{x}),\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta})] + D_1 \|M_t^{(g)}(\mathbf{x}) - \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0) \|_2 \end{align}\]
Without loss of generality, we consider the function \[\begin{align} J(\boldsymbol{\theta};\mathbf{x}) = \log \left[\sum_{g=1}^G \tilde{w}^{(g)}(\mathbf{x}) \exp(\lambda CE[M^{(g)}_t(\mathbf{x}),\boldsymbol{h}(\mathbf{x};\theta) ] )\right], \end{align}\] where \(\tilde{w}^{(g)}(\mathbf{x})= w^{(g)}(\mathbf{x})/{B\left(\mathbf{1}_k+\lambda M^{(g)}_t(\mathbf{x})\right)}\).
suppose we set \(\tilde{w}^{(g)}(\mathbf{x}) \propto 1/\exp(\lambda CE[M^{(g)}_t(\mathbf{x}),\boldsymbol{h}(\mathbf{x};\theta_0) ] )\).Let \(W(\mathbf{x}) = \sum_{g=1}^G \tilde{w}^{(g)}\). Since \(\sum_{g=1}^G w^{(g)} = 1\) and \(B\left(\mathbf{1}_k+\lambda M^{(g)}_t(\mathbf{x})\right)\) is bounded, we may assume \(\frac{1}{W(\mathbf{x})} < D_2\) for a constant \(D_2>0\). By Jensen’s inequality, we have \[\begin{align} J(\boldsymbol{\theta}_0;\mathbf{x}) &=& -\log\left[\frac{W(\mathbf{x})}{G} \sum_{g=1}^G \exp(-\lambda CE[M^{(g)}_t(\mathbf{x}),\boldsymbol{h}(\mathbf{x};\theta_0) ] ) \right]\\ &\leq& \lambda \frac{W(\mathbf{x})}{G} \sum_{g=1}^G CE[M^{(g)}_t(\mathbf{x}),\boldsymbol{h}(\mathbf{x};\theta_0) ] \end{align}\]
By Hölder’s inequality, we have \[\begin{align} && \sum_{g=1}^G \tilde{w}^{(g)}(\mathbf{x}) \exp(CE[\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0),\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta})]) \\ &=& \sum_{g=1}^G \tilde{w}^{(g)}(\mathbf{x}) \exp(CE[M^{(g)}_t(\mathbf{x}),\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta})]) \exp( [M^{(g)}_t(\mathbf{x})-\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0)]^T\log \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}) ) \\ &\leq & \left[\sum_{g=1}^G \tilde{w}^{(g)}(\mathbf{x}) \exp(\lambda CE[M^{(g)}_t(\mathbf{x}),\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta})]) \right]^{\frac{1}{\lambda}} \left[\sum_{g=1}^G \tilde{w}^{(g)}(\mathbf{x}) \exp( \lambda[M^{(g)}_t(\mathbf{x})-\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0)]^T\log \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}) ) \right]^{\frac{1}{\lambda}} \\ &\leq & \exp{J(\boldsymbol{\theta};\mathbf{x})} \left[\sum_{g=1}^G \tilde{w}^{(g)}(\mathbf{x}) \exp( [M^{(g)}_t(\mathbf{x})-\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0)]^T\log \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}) ) \right] \\ &\leq& \exp{J(\boldsymbol{\theta};\mathbf{x})} \exp(T_1(\mathbf{x},\boldsymbol{\theta})) \\ \end{align}\]
\[\begin{align} {\cal E}^{\varphi}(\boldsymbol{\theta}) &\leq& E_\mathbf{x}\left[\frac{1}{W(\mathbf{x})}\sum_{g=1}^G \tilde{w}^{(g)}(\mathbf{x}) CE[M^{(g)}_t(\mathbf{x}),\boldsymbol{h}(\mathbf{x};\theta) ] + \sum_{g=1}^G \frac{\tilde{w}^{(g)}(\mathbf{x})}{W(\mathbf{x})} \|M_t^{(g)}(\mathbf{x}) - \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0) \|_2 \right] \\ &\leq& D_2 E_\mathbf{x}[J(\boldsymbol{\theta};\mathbf{x})]+ \sum_{g=1}^G \int\frac{\tilde{w}^{(g)}(\mathbf{x})}{W(\mathbf{x})} \|M_t^{(g)}(\mathbf{x}) - \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0) \|_2 \mu(\mathbf{x})d\mathbf{x}\\ &\leq& D_2 E_\mathbf{x}[J(\boldsymbol{\theta};\mathbf{x})]+ \sum_{g=1}^G \sqrt{\int \frac{[\tilde{w}^{(g)}(\mathbf{x})\mu(\mathbf{x})/W(\mathbf{x})]^2}{\mu^{(g)}(\mathbf{x})} d\mathbf{x}\int \|M_t^{(g)}(\mathbf{x}) - \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0) \|_2^2 \mu^{(g)}(\mathbf{x})d\mathbf{x}} \\ &\leq& D_2 E_\mathbf{x}[J(\boldsymbol{\theta};\mathbf{x})]+ \sum_{g=1}^G \sqrt{\varepsilon^{(g)} \int \frac{[\tilde{w}^{(g)}(\mathbf{x})\mu(\mathbf{x})/W(\mathbf{x})]^2}{\mu^{(g)}(\mathbf{x})} d\mathbf{x}}, \end{align}\]
We have \[\begin{align} && \log \left[\sum_{g=1}^G \tilde{w}^{(g)}(\mathbf{x}) \exp( [M^{(g)}_t(\mathbf{x})-\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0)]^T\log \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}) ) \right] \\ &\leq& \log \left[\sum_{g=1}^G \tilde{w}^{(g)}(\mathbf{x}) \exp( \|M^{(g)}_t(\mathbf{x})-\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0)\|_2^2/2 + \|\log \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}) \|_2^2/2 )\right] \\ &\leq& D_1 \log \left[\sum_{g=1}^G \tilde{w}^{(g)}(\mathbf{x}) \exp( \|M^{(g)}_t(\mathbf{x})-\boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0)\|_2^2/2 )\right] \end{align}\]
By Jensen’s inequality, we have \[\begin{align} J(\boldsymbol{\theta};\mathbf{x}) &\geq& \sum_{g=1}^G \tilde{w}^{(g)}(\mathbf{x}) CE[M^{(g)}_t(\mathbf{x}),\boldsymbol{h}(\mathbf{x};\theta) ]. \end{align}\]
Let \(W(\mathbf{x}) = \sum_{g=1}^G \tilde{w}^{(g)}\). Since \(\sum_{g=1}^G w^{(g)} = 1\) and \(B\left(\mathbf{1}_k+\lambda M^{(g)}_t(\mathbf{x})\right)\) is bounded, we may assume \(\frac{1}{W(\mathbf{x})} < D_2\) for a constant \(D_2>0\). We have
\[\begin{align} {\cal E}^{\varphi}(\boldsymbol{\theta}) &\leq& E_\mathbf{x}\left[\frac{1}{W(\mathbf{x})}\sum_{g=1}^G \tilde{w}^{(g)}(\mathbf{x}) CE[M^{(g)}_t(\mathbf{x}),\boldsymbol{h}(\mathbf{x};\theta) ] + \sum_{g=1}^G \frac{\tilde{w}^{(g)}(\mathbf{x})}{W(\mathbf{x})} \|M_t^{(g)}(\mathbf{x}) - \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0) \|_2 \right] \\ &\leq& D_2 E_\mathbf{x}[J(\boldsymbol{\theta};\mathbf{x})]+ \sum_{g=1}^G \int\frac{\tilde{w}^{(g)}(\mathbf{x})}{W(\mathbf{x})} \|M_t^{(g)}(\mathbf{x}) - \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0) \|_2 \mu(\mathbf{x})d\mathbf{x}\\ &\leq& D_2 E_\mathbf{x}[J(\boldsymbol{\theta};\mathbf{x})]+ \sum_{g=1}^G \sqrt{\int \frac{[\tilde{w}^{(g)}(\mathbf{x})\mu(\mathbf{x})/W(\mathbf{x})]^2}{\mu^{(g)}(\mathbf{x})} d\mathbf{x}\int \|M_t^{(g)}(\mathbf{x}) - \boldsymbol{h}(\mathbf{x};\boldsymbol{\theta}_0) \|_2^2 \mu^{(g)}(\mathbf{x})d\mathbf{x}} \\ &\leq& D_2 E_\mathbf{x}[J(\boldsymbol{\theta};\mathbf{x})]+ \sum_{g=1}^G \sqrt{\varepsilon^{(g)} \int \frac{[\tilde{w}^{(g)}(\mathbf{x})\mu(\mathbf{x})/W(\mathbf{x})]^2}{\mu^{(g)}(\mathbf{x})} d\mathbf{x}}, \end{align}\] Notice that if \(\boldsymbol{\theta}= \boldsymbol{\theta}_0\), we have \({\cal E}^{\varphi}(\boldsymbol{\theta}_0) \leq D_2 E_\mathbf{x}[J(\boldsymbol{\theta}_0;\mathbf{x})]\).
Suppose we set \(\tilde{w}^{(g)}(\mathbf{x}) \propto \sqrt{\mu^{(g)}(\mathbf{x})/\varepsilon^{(g)}}\), we have \[\begin{align} {\cal E}^{\varphi}(\boldsymbol{\theta}) &\leq& D_2 E_\mathbf{x}[J(\boldsymbol{\theta};\mathbf{x})]+ G\cdot V( \frac{ \sqrt{\mu(\mathbf{x})} }{ \sum_{g=1}^G \sqrt{\mu^{(g)}(\mathbf{x}) /\varepsilon^{(g)}}} ) \end{align}\]
Furthermore, we have \[\begin{align} E_\mathbf{x}[J(\boldsymbol{\theta}_0;\mathbf{x}) ] &\leq& \log \sum_{g=1}^G E_\mathbf{x}\left[ \tilde{w}^{(g)}(\mathbf{x}) \exp(\lambda CE[M^{(g)}_t(\mathbf{x}),\boldsymbol{h}(\mathbf{x};\theta_0) ] )\right], \end{align}\]
For convenience, we denote \(\frac{1}{N} \sum_{i=1}^N \log \left[\sum_{g=1}^G w_i^{(g)} \frac{1}{B\left(\mathbf{1}_k+\lambda \boldsymbol{p}_i^{(g)}\right)} \prod_{k=1}^K\left(h_k\left(\mathbf{x}_i; \boldsymbol{\theta}\right)\right)^{\lambda p_{ik}^{(g)}}\right]\) by \(\widehat{E}_\mathbf{x}[J(\boldsymbol{\theta};\mathbf{x})|{\cal D}]\). Let \(T_1 = V( \frac{ \sqrt{\mu(\mathbf{x})} }{ \sum_{g=1}^G \sqrt{\mu^{(g)}(\mathbf{x}) /\varepsilon^{(g)}}} )\)
Hence, we have \[\begin{align} (1+\lambda){\cal E}^{\varphi}(\widehat{\boldsymbol{\theta}}_n,\boldsymbol{\theta}_0) &\leq& {\cal E}^{\varphi}(\widehat{\boldsymbol{\theta}}_n,\boldsymbol{\theta}_0) + \lambda E_\mathbf{x}[J(\widehat{\boldsymbol{\theta}}_n;\mathbf{x})]+\lambda T_1 - \lambda {\cal E}^{\varphi}(\boldsymbol{\theta}_0) \\ &=& {\cal E}^{\varphi}(\widehat{\boldsymbol{\theta}}_n,\boldsymbol{\theta}_0) - {\cal E}_n^{\varphi}(\widehat{\boldsymbol{\theta}}_n,\boldsymbol{\theta}_0) + {\cal L}^{Multi-KD}(\widehat{\boldsymbol{\theta}}_n) - {\cal L}(\boldsymbol{\theta}_0|{\cal D}) \\ && + \lambda E_\mathbf{x}[J(\boldsymbol{\theta};\mathbf{x})]-\lambda \widehat{E}_\mathbf{x}[J(\boldsymbol{\theta};\mathbf{x})|{\cal D}] +\lambda T_1 - \lambda {\cal E}^{\varphi}(\boldsymbol{\theta}_0) \\ &\leq& \left|{\cal E}^{\varphi}(\widehat{\boldsymbol{\theta}}_n,\boldsymbol{\theta}_0) - {\cal E}_n^{\varphi}(\widehat{\boldsymbol{\theta}}_n,\boldsymbol{\theta}_0)\right| + {\cal L}^{Multi-KD}(\boldsymbol{\theta}_0) - {\cal L}(\boldsymbol{\theta}_0|{\cal D}) \\ && + \lambda \left|E_\mathbf{x}[J(\boldsymbol{\theta}_0;\mathbf{x})]- \widehat{E}_\mathbf{x}[J(\boldsymbol{\theta};\mathbf{x})|{\cal D}]\right| +\lambda T_1 - \lambda {\cal E}^{\varphi}(\boldsymbol{\theta}_0) \\ &\leq& \left|{\cal E}^{\varphi}(\widehat{\boldsymbol{\theta}}_n,\boldsymbol{\theta}_0) - {\cal E}_n^{\varphi}(\widehat{\boldsymbol{\theta}}_n,\boldsymbol{\theta}_0)\right| + \lambda \widehat{E}_\mathbf{x}[J(\boldsymbol{\theta}_0;\mathbf{x})|{\cal D}] - \lambda E_\mathbf{x}[J(\boldsymbol{\theta}_0;\mathbf{x})] \\ && + \lambda \left|E_\mathbf{x}[J(\boldsymbol{\theta}_0;\mathbf{x})]- \widehat{E}_\mathbf{x}[J(\boldsymbol{\theta};\mathbf{x})|{\cal D}]\right| +\lambda T_1 + \lambda E_\mathbf{x}[J(\boldsymbol{\theta}_0;\mathbf{x})] - \lambda {\cal E}^{\varphi}(\boldsymbol{\theta}_0) \end{align}\] ◻
Simulation 1: The teacher models employ a Multilayer Perceptron (MLP) architecture consisting of five hidden layers. These layers have 7, 10, 12, 10, and 5 nodes respectively. The model uses the ReLU activation function. Different dropout rates are applied for two teacher models. The student model employs an MLP architecture consisting of one hidden layer with 5 nodes.
Simulation 2: The teacher models employ the MLP architecture consisting of five hidden layers with 7, 15, 12, 10, and 7 nodes respectively. The model uses the ReLU activation function and incorporates a dropout rate of 0.2. The student model utilizes an MLP architecture with a single hidden layer comprising 10 nodes, activated by the ReLU function.
| Teacher models | Student Model | |||
| \(M_T^{(1)}\) | \(M_T^{(2)}\) | \(M_T^{(3)}\) | \(M_s\) | |
| Simulation 1 | MLP-1 (430) | MLP-1 (430) | - | MLP-S1 (27) |
| Simulation 2 | MLP-2 (600) | MLP-2 (600) | MLP-2 (600) | MLP-S2 (115) |
| Protein | ESM-2-650M (650M) | ProtT5-XL-half (120M) | - | ESM-2-8M (8M) |
| Digit-5 | CNN-L (2M) | Resnet-50 (26M) | Resnet-50 (26M) | CNN-S (0.1M) |
To implement the method integrating the original KD and BNN, we set the BNN model as the student. We train this model using variational inference with details described in [63], and incorporate the KD loss to distill knowledge from the teacher model. An important hyperparameter when training the BNN model is the weight of the CE loss. We perform a grid search over the values \(\{0.01, 0.05, 0.1, 0.2\}\) to find the parameter that yields the highest accuracy. For the CIFAR-10 and CIFAR-100 datasets, due to computational constraints, we add the Bayesian layers only after the feature extraction blocks.
To implement the method integrating the original KD and Monte Carlo dropout, we enable dropout during the testing/inference stage to obtain multiple stochastic predictions.
ESM-2 model: The ESM-2 model (Evolutionary Scale Modeling v2) [54], developed by Meta AI, is a state-of-the-art protein language model trained on large-scale protein sequence datasets, including UniRef50 [64]. Using a Transformer architecture, it captures intricate sequence patterns and relationships, excelling in tasks such as structure prediction, mutation effect prediction, and protein design. We use the ESM-2-650M variant, which contains 650 million parameters and offers improved performance and scalability over its predecessor, as teacher model 1. The ESM-2-8M variant, which contains 8 million parameters, is used as the student model.
ProtT5 model: The ProtT5 model is a protein language model based on the T5 architecture and trained on UniRef50. It captures both local and global sequence features, excelling in tasks such as structure prediction, function annotation, and sequence classification. We use the ProtT5-XL-half model as teacher model 2, which is a half-precision encoder-only variant optimized for efficient and accurate protein sequence embeddings.
During training, the feature extraction layers of these teacher models are frozen, and only the classifier layers are fine-tuned using our dataset.
We focused on three colored variants of Digit-5: SVHN, MNIST-M, and SynthDigits. SVHN (Street View House Numbers) is a dataset of digit images obtained from real-world house number signs in Google Street View [65]. It includes a training set of 73,257 examples and a test set of 26,032 examples. Each example is a 32 × 32 colored image associated with a label of 10 classes, representing the digits 0 through 9. MNIST-M is a modified version of the original MNIST dataset, where grayscale digit images are superimposed on randomly chosen colored background patches [62]. The dataset retains the original MNIST structure, with a training set of 55,000 examples and a test set of 10,000 examples. Each example is a 28 × 28 colored image associated with a label of 10 classes. SynthDigits is a synthetic dataset consisting of 25,000 training images and 10,000 test images, where each example is a 32 × 32 colored image of a digit from 0 to 9 [66]. The images are generated with diverse fonts, colors, and background patterns, providing a visually rich and challenging dataset for digit recognition tasks.
Model Structure: The teacher model employs the 50-layer ResNet architecture as described by [67]. The ResNet50 model has 25.6 million parameters. The student model employs a CNN architecture with two convolutional layers and a subsequent fully connected layer. The first convolutional layer has 16 output channels, batch normalization, ReLU activation, and max pooling. The second convolutional layer follows a similar pattern but with 8 output channels. The convolutional layer is then followed by fully connected layers, consisting of a hidden layer with 60 nodes and an output layer with 10 nodes. The student model has approximately 0.1 million parameters, which is 256 times fewer than the teacher models.