Scaling Ensemble Distribution Distillation to Many Classes with Proxy Targets
Year: 2022

Abstract:
Ensembles of machine learning models yield improved system performance as well as robust and interpretable uncertainty estimates; however, their inference costs can be prohibitively high. Ensemble Distribution Distillation (EnD2) is an approach that allows a single model to efficiently capture both the predictive performance and uncertainty estimates of an ensemble. For classification, this is achieved by training a Dirichlet distribution over the ensemble members’ output distributions via the maximum likelihood criterion. Although theoretically principled, this work shows that the criterion exhibits poor convergence when applied to large-scale tasks where the number of classes is very high. Specifically, we show that for the Dirichlet log-likelihood criterion classes with low probability induce larger gradients than high-probability classes. Hence during training the model focuses on the distribution of the ensemble tail-class probabilities rather than the probability of the correct and closely related classes. We propose a new training objective which minimizes the reverse KL-divergence to a Proxy-Dirichlet target derived from the ensemble. This loss resolves the gradient issues of EnD2, as we demonstrate both theoretically and empirically on the ImageNet, LibriSpeech, and WMT17 En-De datasets containing 1000, 5000, and 40,000 classes, respectively.

heading: 1 Introduction
text: As machine learning models are being deployed ever more widely into production, especially in safetycritical applications such as self-driving cars or medical diagnostics, obtaining reliable estimates of uncertainty in their predictions becomes increasingly important to guaranteeing their safety and reliability [1]. Ensemble methods are known to yield both improved predictive performance and reliable uncertainty estimates [2, 3, 4, 5, 6]. Furthermore, ensembles allow separating out total uncertainty into data and knowledge uncertainty2. The former is the intrinsic uncertainty due to class overlap and noise inherent in the data, while the latter is the model’s uncertainty due to lack of understanding of the test data [7]. Estimates of knowledge uncertainty are often used to detect anomalous and unfamiliar inputs [8, 9, 10, 7] as well as in active learning [11].
Despite their numerous advantages, using ensembles can be computationally prohibitive. Obtaining predictions in real time is often expensive even for a single model, and the hardware requirements for serving an ensemble scale linearly with size. As a result, over the past several years the area of ensemble distillation has gained increasing attention. Broadly speaking, distillation methods aim to train a single model which can approximate the behavior of an ensemble, but at low computational cost. In the simplest and most frequently used form of distillation [12], the student model is trained to capture the average prediction of the ensemble: for example, in case of classification this reduces
⇤Equal contribution, order decided by coin toss. 2Data and Knowledge Uncertainty are also known as Aleatoric and Epistemic uncertainty.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
to minimizing the KL-divergence between the model and the ensemble mean. While this method allows the student to obtain predictive performance comparable to that of the original ensemble, the information about its distributional properties (in other words, its diversity) is lost in the process. As ensemble-based uncertainty estimation utilizes the information about disagreement between its members, such distillation approaches fail to preserve valuable information from the ensemble.
Recently, several works have proposed distillation procedures that capture information about both the mean and distribution of ensemble predictions within a single model [13, 14, 15, 16]. These can be broadly split into methods where the student has multiple output heads - one per ensemble member [14, 15], and those which yield a higher-order distribution over the ensemble’s outputs [13, 16]. The former are conceptually simpler, but more computationally expensive and may model the ensemble in excessive detail; the latter approaches are more light-weight, but also complex and, as we show in this work, suffer from optimization issues. This latter class of approaches, which we refer to as Ensemble Distribution Distillation (EnD2), is the focus of this work.
Ensemble Distribution Distillation views outputs of each ensemble member as samples from a higher-order Dirichlet or Normal-Wishart distribution, and the student model learns the parameters of that distribution. Typically, EnD2 is done by maximizing the likelihood of the ensemble’s output distributions under the Dirichlet or Normal-Wishart prior. While this criterion is theoretically sound, we show that for large-scale classification tasks with many classes, gradient-based optimization of it is highly problematic, which limits its usefulness in real-life production scenarios. In this work, we investigate the poor convergence of models trained with Ensemble Distribution Distillation at scale. We analyze the the Dirichlet log-likelihood criterion and show that it leads to high gradient norm values that affect optimization. Specifically, if the ensemble members’ output distributions have most probability mass allocated to a few classes, with the remainder spread among a long tail of exponentially less likely classes, then the gradients associated with the tail-classes will be significantly larger than those associated with high-probability classes. As a result, the optimization procedure focuses on modelling this distribution of tail-class probabilities, preventing convergence.
To solve this, we propose a two-step solution. First, we transform the empirical distribution of ensemble member predictions into a Proxy-target Dirichlet distribution with the same statistics and to use it as the target during distribution distillation. Second, we show that it is crucial to minimize the reverse KL-divergence between the model and the Proxy-Dirichlet, as minimization of the forward KL-divergence exacerbates the optimization issues. The proposed training procedure allows the model to converge, mitigating the issue of gradient explosion. We demonstrate this by distribution-distilling ensembles of models trained on the ImageNet classification, LibriSpeech speech recognition and WMT17 English-German language translation datasets, where there are 1000, 5000 and classes, respectively. To accomplish the latter two, we extend EnD2 to autoregressive models and refer to this approach as Sequence Ensemble Distribution Distillation (SenD2). On both datasets the distribution-distilled models outperform the baselines and yield uncertainty estimates competitive this those of the original ensemble.
Thus, our contributions are as follows:
• We draw attention to and analyze the convergence issues of Dirichlet likelihood optimization when applied to tasks with large numbers of classes;
• We propose an alternative two-step training approach for Ensemble Distribution Distillation which overcomes the aforementioned optimization issues;
• We adapt Ensemble Distribution Distillation to auto-regressive models and propose Sequence Ensemble-Distribution Distillation (SEnD2).

heading: 2 Preliminaries: Ensembles and Distillation
text: We view ensembles within a Bayesian framework, where the model parameters ✓ are random variables over which a prior p(✓) is placed. The posterior distribution p(✓|D) is obtained via Bayes’rule:
p(✓|D) = p(D|✓)p(✓)
p(D) / p(D|✓)p(✓) (1)
Consider an ensemble of models {P(y|x⇤,✓(m))}Mm=1 sampled from the posterior: P(y|x,✓(m))
M m=1 ! P(y|⇡(m)) M m=1 , ⇡(m) = f(x;✓(m)), ✓(m) ⇠ p(✓|D) (2)
where ⇡ are the parameters of a categorical distribution [P(y = !1), · · · , P(y = !K)]T. The predictive distribution, or predictive posterior, for a test input x⇤ is obtained by marginalization:
P(y|x⇤,D) = Ep(✓|D) ⇥ P(y|x⇤,✓) ⇤ ⇡ 1
M
MX
m=1
P(y|x⇤,✓(m)) (3)
In practice this is intractable and we approximate via Monte-Carlo sampling. Given the ensemble, the entropy of the predictive posterior is a measure of total uncertainty. Knowledge uncertainty can be assessed via measures of the spread, or ‘disagreement’, of the ensemble such as mutual information:
I[y,✓|x⇤,D]| {z } Knowledge Uncertainty
= H ⇥ Ep(✓|D)[P(y|x⇤,✓)] ⇤ | {z }
Total Uncertainty
Ep(✓|D) ⇥ H[P(y|x⇤,✓)] ⇤ | {z }
Expected Data Uncertainty (4)
An alternative measure of spread, which is particularly useful for ensemble of autoregressive sequence models, is the reverse-mutual information (RMI) [17]:
M ⇥ y,✓|x,D ⇤ = Ep(✓|D) h EP(y|x,D) h ln P(y|x,D)
P(y|x,✓)
ii (5)
While ensembles yield improved predictive performance and meaningful uncertainty estimates, they are computationally expensive. To reduce inference-time cost, an ensemble can be distilled into a single model P(y|x; ) by minimizing the KL-divergence to the predictive posterior of the ensemble:
L EnD( ,Dens) =Ep̂(x) h KL ⇥ P(y|x,D) || P(y|x; ) ⇤i (6)
This approach has been applied to tasks such as image classification and machine translation [12, 18]. While distillation allows a single model to capture the ensemble’s predictive quality and estimates of total uncertainty at low cost, information about ensemble diversity is lost. Consequently, we cannot obtain estimates of knowledge uncertainty, which are particularly useful for anomaly detection [7, 13].
Recently, a distillation technique called Ensemble Distribution Distillation (EnD2) was proposed [13]. Here, both the mean and diversity of an ensemble are captured within a single model via distillation into a Prior Network [19], which is a model that parameterizes the Dirichlet distribution:
p(⇡|x; ) =Dir(⇡|↵),↵ = ez, z = f(x; ), ↵c > 0, ↵0 = KX
c=1
↵c (7)
Distribution distillation is done in two steps. Firstly, a transfer dataset Dens = {x(i),⇡(i,1:M)}Ni=1 ⇠ p̂(x,⇡) is composed of the inputs x(i) from the training set D = {x(i), y(i)}Ni=1 and the ensemble’s output distributions {⇡(i,1:M)}Ni=1 for each input. Secondly, given this transfer set, the model p(⇡|x; ) is trained by minimizing the negative log-likelihood of each categorical distribution ⇡(im):
L EnD2( ,Dens) = Ep̂(x) ⇥ Ep̂(⇡|x)[ln p(⇡|x; )] ⇤ (8)
Given a distribution-distilled Prior Network, the predictive distribution is given by the expected categorical distribution ⇡̂ under the Dirichlet prior:
P(y = !c|x ⇤; ) = Ep(⇡|x⇤; )[P(y = !c|⇡)] = ⇡c = ↵cPK k=1 ↵k = ezc PK k=1 e zk (9)
Measures of total and knowledge uncertainty are obtained by considering the mutual information between the prediction y and the parameters of ⇡ of the categorical:
I[y,⇡|x⇤; ]| {z } Knowledge Uncertainty
= H ⇥ Ep(⇡|x⇤; )[P(y|⇡)] ⇤ | {z }
Total Uncertainty
Ep(⇡|x⇤; ) ⇥ H[P(y|⇡)] ⇤ | {z }
Expected Data Uncertainty (10)
Alternative measures, such as RMI, can also be obtained, as detailed in appendix C.
An alternative approach to capture ensemble diversity is to distill an ensemble into a mixture model which yields a separate output head for each ensemble member [14, 15]. However, the principal downside of this approach is that it attempts to model the ensemble in excessive detail, as it models the behaviour of each ensemble member, rather than bulk distributional properties. This requires more flexible and powerful models. As a result, for good performance, it is necessary to split the model into multiple heads at an early stage, which significantly increases the computational and memory complexity. In contrast, EnD2 has a fixed computational and memory cost of one model, regardless of the size of the original ensemble, as it models the ensemble’s bulk distributional properties.

heading: 3 Analysis of first-order gradients
text: The previous section described how EnD2 can be done by maximising the log-likelihood of the ensemble’s outputs under the Dirichlet prior. However, we empirically observed significant convergence issues when optimizing this criterion on tasks with large numbers of classes. As neural networks are predominantly trained using gradient-based optimization, we examine the gradients of the Dirichlet negative log-likelihood loss and propose an alternative training approach.
Our analysis considers the following three situations. First, a Prior Network initialization which yields a uniform Dirichlet distribution (↵init = 1). The target categorical distribution whose probability is being maximized is a sparse K-length vector of probabilities:
⇡tgt = ⇥ 1 ✏, ✏/(K 1), ✏/(K 1), · · · ⇤T , ✏ = 1e-4 (11)
Second, a Prior Network near convergence with the following output distribution:
↵conv = ⇡conv · ↵0, ↵0 = 9e4, ⇡conv = ⇥ 1 5✏, 5✏
K 1 ,
5✏
K 1 , · · ·
⇤T (12)
Third, a Prior Network which has made a strong misclassification — a situation which could occur in the middle of training, far from convergence:
↵misc = ⇡misc · ↵0, ↵0 = 9e4, ⇡misc = ⇥ 5✏ K 1 , 5✏ K 1 , · · · , 1 5✏ ⇤T (13)
We consider the following training criteria in this analysis: categorical cross entropy3, Dirichlet negative log-likelihood, Dirichlet Forward and Reverse KL-divergence.
L Cat-KL =
KX
k=1
⇡̂k ln ↵k ↵0 , LDir-NLL=
KX
k=1
(↵k) (↵k 1) MX
m=1
ln⇡(m)k M
(↵0)
L Dir-KL =
KX
k=1
(↵k) KX
k=1
( k) + ( 0) (↵0) + KX
k=1
( k ↵k) ⇣ ( k) ( 0) ⌘
L Dir-RKL =
KX
k=1
( k) KX
k=1
(↵k) + (↵0) ( 0) + KX
k=1
(↵k k) ⇣ (↵k) (↵0) ⌘
(14)
The first loss is used for standard distillation [12], which is empirically known to converge well. The second is the loss function originally proposed for EnD2 [13]. The third and fourth loss functions, the Dirichlet forward and reverse KL-divergence, are analyzed as they relate to our proposed solution in the latter part of this section. The target Dirichlet distribution for these two loss functions has parameters = ⇡tgt · 0 where 0 = 1e5. Full derivations are available in appendix A.
Let’s examine the gradients of each criterion presented in equation 14 with respect to the logits.4 We introduce the ratio ⇢ of the gradient with respect to the logit z1, the high-probability class of interest, to the average of the absolute values of gradients with respect the logits z2:K of the low-probability tail-classes. The ratio ⇢ represents the relative contribution of the gradients with respect to the class we are interested in modelling to the long tail.
⇢ = K 1
K ·
@L
@z1
.⇣ KX
k=2
@L
@zk
⌘
(15)
We normalize by number of classes K, as this improves interpretability by making ⇢ of LCat-KL plateau. This makes it easier to compare to the behaviour of ⇢ for other criteria.
Figure 1 shows that, at initialization, as the number of classes is increased, the categorical crossentropy loss primarily focuses on the high probability class and ignores the long tail. In contrast, the Dirichlet NLL displays a diminishing contribution. As the loss is very sensitive to minor perturbations of tail-class probabilities, it will focus on modelling the high-probability classes only after it perfectly
3Which is equivalent to categorical Forward KL-divergence. 4We remind the reader that the ↵k = ezk , where zk are the logits.
models the long tail. This means that on complex tasks the model is perpetually stuck modelling the probabilities of tail classes. Note that even near convergence, the ratio ⇢ is far smaller for LDir-NLL than for categorical cross-entropy. Finally, if a significant error is made, ⇢ remains large for cross-entropy, but becomes very small for Dirichlet NLL as the number of classes increases. This shows that a desirable property of the loss which ensures good convergence is that the ratio ⇢ is high and either constant or increasing as the number of classes grows, otherwise the model focuses on modelling the distribution of tail-class probabilities across the ensemble.
Figure 1 also displays the ratio ⇢ for both the forward and reverse Dirichlet KL-divergence losses. Clearly, the forward KL-divergence displays the same issues as the Dirichlet NLL loss and ⇢ continues to decrease as the number of classes in increased. This is unsurprising, as the negative log-likelihood is equivalent to the KL-divergence in the limit. However, the reverse KL-divergence displays the desirable properly that ⇢ grows and stabilizes as the number of classes is increased in all situations. This suggests that if we were to minimize the reverse KL-divergence to an appropriately chosen target Dirichlet distribution, then we could avoid convergence issues.

heading: 4 Distribution Distillation via Proxy-Dirichlet target distribution
text: The results of the above analysis suggest the reverse KL-divergence criterion has many of the same attributes of the categorical cross-entropy loss — specifically that it focuses on modelling the distribution of probabilities of high-probability classes. Thus, if we could specify an appropriate target Proxy Dirichlet distribution which represents the ensemble, then we could accomplish distributiondistillation without optimization difficulties.
To correctly construct a Proxy-Dirichlet target which represents the ensemble, we need to obtain a mean ⇡̂k and precision ̂0 derived from the ensemble. While obtaining the ensemble mean is trivial, obtaining an estimate of the precision is more challenging. One approach based on Stirling’s approximation is described in [20] and proposes the following estimate:
⇡̂k(x) = 1
M
MX
m=1
P(y = !k|x,✓ (m)), ̂k(x) =
⇡̂k · (K 1)
2 PK
k=1 ⇡̂k(ln ⇡̂k 1 M PM m=1 ln⇡
(m) k )| {z }
Stirling’s Approximation
(16)
Thus, the Proxy Dirichlet has both the same mean and bulk-diversity properties of the ensemble, such that the measures of diversity derived from the Proxy Dirichlet are similar to those of the original
ensemble. To prevent the Dirichlet distributions of both the target and model from entering highly non-linear regimes impacting the stability of the training process, both the model parameters, ↵k, and target parameters, k, are constrained to be greater than 1. Figure 1 shows that adding this constraint to the concentration parameters of the Proxy-Dirichlet and the model improves ⇢, especially near convergence. Altogether, this yields the following criterion:
L PD-EnD2=KL[p(⇡|x, )kp(⇡|̂+1)] / Ep(⇡|x, )
⇥
KX
k=1
⇡̂k ln⇡k ⇤
| {z } Reconstruction term
+ KL[p(⇡|x, )kp(⇡|1)]
̂0| {z } Prior
(17)
By dividing through by ̂0 the above criterion is made to have the same form as an ELBO5. This serves two purposes. First, it becomes apparent that the reconstruction term is an upper bound on categorical cross entropy, which explains why they share similar properties. Second, normalizing by ̂0 allows this criterion to have the same dynamic range the KL-divergence loss used for standard distillation, and therefore use the same hyperparameters. The proposed approach is evaluated on synthetic data in appendix B.
We want to stress that while this criterion is similar to the one proposed in [10], the underlying reason for its use is entirely different. Here, the issue is due to large gradients from tail classes, while [10] uses a reverse-KL criterion to avoid inducing in expectation a multi-modal target Dirichlet distribution.

heading: 5 Experiments
text: In this section, we evaluate EnD2 via minimization of the reverse KL-divergence to Proxy Dirichlet target, which we refer to as PD-EnD2. We apply PD-EnD2 to ensembles of convolutional networks trained on the ImageNet [21] dataset, ensembles of VGG-Transformer [22] ASR models trained on LibrSpeech [23] and ensembles of Transformer-big [24] models trained on WMT’17 En-De, which feature 1000, 5000 and 40,000 classes, respectively. Our goal is to demonstrate that given an ensemble, we can successfully distribution-distill it into a single model on tasks with many classes, which was previously impossible. We do not provide results for EnD2 accomplished by optimizing Dirichlet negative log-likelihood or Dirichlet forward KL-divergence because these approaches do not even begin to converge on the tasks considered. Finally, we would like to highlight that providing over-tuned SOTA predictive performance or OOD detection performance on the chosen tasks is not the focus of this research 6.
Setup We consider three large-scale tasks involving classification: image classification, speech-to-text transduction and text-to-text transduction. For each task, we first train a Deep Ensemble [4] of regular models and then distribution-distill it with PD-EnD2. We consider the following baselines:
• Single refers to the average performance across 10 individual, independently evaluated models. • Ensemble refers to the performance of a Deep Ensemble, which is currently considered to be both
the best and the simplest ensemble method [4, 5, 6]. • EnD or SEnD refer to ensemble distillation via minimizing the KL divergence between the student
and the ensemble’s mean for image classification [12] and machine translation [18], respectively.
We do not use multi-head approaches [14, 15, 25], for distilling each separate ensemble member, as the computational overhead necessary to individually model each ensemble member with high fidelity on these tasks is significant and it is unclear how to extend them to structured tasks. In all experiments with PD-EnD2, we add 1 both to the predicted parameters of the Dirichlet distribution and the Dirichlet proxy parameters, as discussed in the previous section.
To do distribution-distillation for NMT, which is an autoregressive structured prediction task, we extend Prior Networks and EnD2 to such models and develop Sequence Ensemble Distribution Distillation (SEnD2). Models trained with Proxy-Dirichlet targets are referred to as PD-SEnD2.
5A full derivation is available in appendix A 6Code and training configurations are available at https://github.com/yandex-research/
proxy-dirichlet-distillation.
This extension involves modifications to both the model and derivation of associated sequence-level uncertainty measures [17]. While conceptually straightforward, this modification is extensive and mathematically involved, and therefore instead fully detailed in appendix C.
We evaluate the proposed approach and the baseline in two ways. First, we evaluate their predictive performance and calibration. This assesses how useful the resulting models are to the task at hand. We expect that a successfully distribution-distilled model should have comparable predictive performance to standard. ensemble distillation. Second, we evaluate the measures of uncertainty which the models yield on the the task of out-of-distribution detection. Here we consider a range of distributionally shifted, or out-of-distribution, datasets on which we expect the models to either perform poorly, or not to be able to perform at all. The goal is to be able to discriminate between the in-domain and out-of-distribution datasets using measures of uncertainty. We expect that a successfully distribution-distilled model will yield similar performance to the ensemble, especially using measures of knowledge uncertainty.
Large-scale image classification We begin by examining PD-EnD2 on the ImageNet [21] dataset, which features 1000 classes. This represents an increase in the scale of the task in in terms of dataset size, image size and, most importantly, number of classes, relative to the tasks considered in the original work on Ensemble Distribution Distillation [13]. Here, we distribution-distill an ensemble of 10 ResNet-50 [26] models. To train each ResNet-50 model, we use the standard training setup outlined in [27, 28]. Specifically, we train for 90 epochs using stochastic gradient descent with a momentum of 0.9 and a learning rate of 0.1 ⇥ B/256, where B is the per-device batch size multiplied by the number of GPUs. In our experiments, we use a single-GPU batch size of 256 and 8 NVIDIA V100 GPUs. The learning rate is divided by 10 every 30 epochs. For data augmentations, we use a combination of random resized crops and horizontal flips implemented in the Albumentations library [29]. In all experiments, we found it beneficial to initialize the last batch normalization in each residual branch to zero, which agrees with previous results [28, 30, 31]. Additional results are provided in appendix E.
To evaluate the predictive performance of the proposed method, we measure the classification accuracy and Expected Calibration Error (ECE) on the original ImageNet validation subset [21] and on a range of distributionally shifted datasets. Specifically, we use natural adversarial examples from ImageNet-A [32], corrupted and perturbed versions of the ImageNet validation data from ImageNetC [33], and artistic renditions from ImageNet-R [34]. These datasets feature the same classes as ImageNet, but they are mismatched in different ways. This assesses how robust the proposed methods are to distributional shift [35], or mismatch, between training and evaluation data. The metrics on ImageNet-C are averaged over all degrees of corruption.
Results in Table 1 show that EnD2 is capable of accurate emulation of the ensemble in terms of classification performance. In terms of accuracy, the method displays results on par or slightly better than standard distillation, and typically better than a single model. The only exception is on the ImageNet-A dataset, where single models achieve best performance. It is likely that the natural adversarial attacks are such that the models in the ensemble become correlated and the error is reinforced. In terms of calibration, the distribution distilled models are closest to the Ensemble. Furthermore, on the in-domain set they outperform the ensemble. These results clearly demonstrate that PD-EnD2 is successfully able to produce models which are comparable to quality to ensemble distillation, but are better calibrated.
We evaluate the measures of uncertainty produced by the proposed approach and the baselines on the task of out-of-distribution detection. As discussed above, the goal is to discriminate between indomain and distributionally shifted (OOD) data using measures of uncertainty. We consider ImageNet
A,C and R as well as ImageNet-O [32], which features altogether new classes not seen in ImageNet, as out-of-distribution data. We use the entropy of the predictive posterior and mutual information (equation 10) as estimates of total and knowledge uncertainty, respectively. Out-of-distribution detection performance is assessed in terms of Receiver Operating Characteristic area under curve (ROC AUC). We additionally verify how well the Proxy Dirichlet target approximates the ensemble by assessing how well the measures of uncertainty which it yields perform on this task.
Results in Table 2 show several trends. Firstly, the Proxy-Dirichlet distillation can closely match the performance of the ensemble on all datasets except ImageNet-O, which seems to be particularly difficult. Secondly, in all cases the distribution-distilled model is able to obtain performance close to that of the ensemble and generally yields similar behaviour. This clearly demonstrates that PD-EnD2 can successfully train a model to emulate an ensemble.
In addition to the tables above, we provide Figure 3, which shows a detailed breakdown by corruption strength of the predictive and OOD detection performance on ImageNet-C. Results show that as the degree of corruption increases, all models suffer a drop in accuracy and calibration quality; notably, EnD and PD-EnD2 have the same calibration performance on original data, but PD-EnD2 has lower calibration errors for the highest degrees of corruption. Unsurprisingly, the greater the distributional shift from the original training images, the better the models are at out-of-distribution detection. Overall, the distribution-distilled model clearly follows the behaviour of the ensemble in all metrics.
Speech Recognition
We base our training setup on the description given in [22], training for 80 epochs on 8 GPUs with a batch size of 8,000 tokens with the AdaDelta [36] algorithm and a constant learning rate. The dataset is processed with a vocabulary of 5000 sentence-piece tokens. Checkpoints over the last 30 epochs are averaged together. The only significant difference is that we train SEnD2 with the Adam [37] optimizer, using the same learning rate schedule as in the machine translation experiments.
We use four datasets for evaluation: LibriSpeech test-clean and test-other, AMI meeting transcription dataset [38], and French speech recordings from Common Voice [39]. This setup is identical to the one considered in [17]. LibriSpeech test-other represents a noisy dataset, while AMI and Common-Voice French represent significant distributional shifts.
First, we compare the methods in terms of their predictive performance using the Word Error Rate (WER)7 and test-set negative log-likelihood on the LibriSpeech test-clean datset (LTC). Results in Table 3 show that both distillation and distribution distillation are challenging for ASR: neither model is able to outperform a single model in terms of WER. However, in terms of test-set NLL, the distribution-distilled model (PD-SEnD2) is significantly better (lower NLL) than both a single model and the standard distilled model. Coupled with the fact that the WER is worst for PD-SEnD2, this suggest that the model is underfit, and a different training setup, with longer training and less regularization, should be considered. Nevertheless, these preliminary results show that comparable performance is attainable. We remind the reader that Naive EnD2 does not even begin to converge on the tasks which we consider here. Secondly, we assess the uncertainty estimates via OOD detection. As the in-domain set, we use Librispeech test-clean, and as the OOD sets, we consider LibriSpeech test-other (LTO), AMI and CV French. For out-of distribution detection, we use the same measures of total and knowledge uncertainty as in Section 5, which are detailed in appendix C. Results in table 3 show that the distribution-distilled model is able to yield measures of total uncertainty which outperform both a single model and a distilled model. However, measures of knowledge uncertainty do not match the ensemble well, which suggests that the model is either misfit or underfit. In the light of the predictive performance results, the latter is more likely.
Machine translation Now we evaluate PD-EnD on the WMT’17 English-German machine translation dataset with a vocabulary of 40,000 Byte-Pair Encoding tokens [40]. This represents a more challenging task of greater scale than ImageNet. For distillation and distribution distillation we use an ensemble of 10 Transformer-big [41] trained with the setup described in [42]. Specifically, ensemble members are trained for 193,000 steps with Adam [37] on 8 NVIDIA V100 GPUs with a batch size of 4096 tokens per GPU. For both distillation and distribution distillation, we train all models for 20,000 steps with a large batch size of 32K tokens. As our implementation requires fitting all 10 ensemble members in GPU memory, we reduce the immediate batch size for each step to 1024, but compensate for it with gradient accumulation over 32 steps. We use beam search with the beam size equal to 5 for inference.
7Computed using NIST SCLITE tool.
First, we assess the proposed approach in terms of translation quality on the newstest14 EnglishGerman test set, for which we use the BLEU score [43] computed with SacreBLEU [44]. Results in Table 4 confirm that PD-EnD2 is capable of matching regular distillation in terms of translation quality. This shows that the proposed training approach does not suffer from any optimization issues and yields fully converged models of the same quality as standard ensemble distillation.
Second, we assess the calibration performance in terms of NLL both on the reference transcriptions and on the ensemble’s transcription (ENS-Pred) of test data. The results show that the ensemble has better (lower) NLL than a single model on reference transcriptions, but both distilled models have worse (higher) NLL than both the ensemble and single model on reference transcriptions. On ensemble transcriptions, both distilled models have better NLLs than the single model. This is expected, as these models were never trained on the reference data — they were always trained on transcriptions provided by the ensemble. Furthermore, note that the NLLs of all models are better (lower) on ensemble predictions than on the references. This is consistent with the observations made in [45] that reference transcriptions are more complex and multimodal, more surprising, therefore would have higher NLLs, while predictions of NMT models are more consistent in style and therefore simpler — they are less ’surprising’, and therefore have lower NLL. These results show that while all distilled models consider the references to be more surprising, as their decoders were never exposed to the reference text, the distilled models are more closely matched to the ensemble, and thus yield better performance.
Third, we assess the quality of the uncertainty estimates which our distribution-distilled model yields. For out-of-distribution detection, we use sequence-level entropy and reverse mutual information as measures of total and knowledge uncertainty, which have been shown to perform well on this task [17].8 The in-domain test set is newstest’14 and OOD datasets are as follows: ‘permuted’ — sentences with permuted tokens in the input, ‘speech’ — LibriSpeech [23] test-clean English speech transcriptions; ‘German’ and ‘French’ — source sentences from newstest’14 in German and French languages respectively. This is the same setup and datasets as were considered by [17]. Results of both distillation methods are averaged over 5 random seeds.
Table 4 clearly shows that in all cases, the PD-SEnD2 model is able to yield measures of uncertainty which outperform SEnD and single models and are highly competitive with the original ensemble. However, the distribution-distilled model doesn’t follow the behaviour of the ensemble as closely, which is not surprising, as the task is far more complex than image classification and the output space is exponentially vast.

heading: 6 Conclusion
text: This work has examined applying Ensemble Distribution Distillation (EnD2) to large-scale tasks where the number of classes is very high. Training these models is shown to be challenging with often poor convergence properties. We provided an analysis which shows that the standard training criterion, the Dirichlet log-likelihood yields gradients associated with low probability classes that dominate those of the higher probability classes. Thus, during training the model will focus on the distribution of these ensemble tail-class probabilities. To address this problem, we proposed a new training objective which minimizes the reverse KL-divergence to a Proxy-Dirichlet target derived from the ensemble. This new criterion resolves the gradient issues of EnD2, which we demonstrate empirically on the ImageNet, LibriSpeech and WMT17 En-De datasets, which have 1000, 5000 and 40,000 classes, respectively. Thus, this work enables Ensemble Distribution Distillation to be applied to tasks with arbitrary numbers of classes and complexity, dramatically broadening the range of applications that can benefit from rapid inference and efficient uncertainty estimation.
8Further details and derivations are available in appendix C.

Review summaries:
This paper considers uncertainty estimation for neural networks, specifically predictive uncertainty and knowledge uncertainty. For these problems, deep ensembles are considered state-of-the-art, but are expensive at training time and, more importantly, inference time. To address the latter problem, this paper explores Ensemble Distribution Distillation (EnD^2), an approach previously proposed by Malinin et al (2020) which consists of distilling the ensemble into a single model that captures both the mean and diversity of the original ensemble. However, this paper notes that EnD^2 fails to converge for problems where there is a large number of classes, which is attributed to the use of the forward KL divergence as the distillation criterion. Instead, the reverse-KL divergence is proposed using a proxy Dirichlet distribution to represent the ensemble. The precision of the proxy distribution is estimated using Stirling's approximation, and some constraints are imposed on the model parameters to ensure stability. Experiments on ImageNet (several variants) and NMT show that the proposed approach successfully retrains some of the accuracy and calibration benefits of the corresponding deep ensemble.    

This paper proposes an algorithm for distilling deep ensembles into a single neural network for large-scale classification tasks. The most common baseline in this context would be simply distilling the ensemble models' mean to a single student, but this would throw away some important aspects of ensemble teachers, especially the diversities in predictions. To resolve this, the ensemble distribution distillation (EnD^2) was proposed, where the ensemble predictions are treated as samples from an underlying distribution of distribution (chosen to be Dirichlet), and the student model tries to learn this distribution via maximum likelihood. This paper remedies an important drawback of EnD^2; EnD^2 does not train well for classification tasks with a large number of classes. The authors empirically identify the source of this problem to be the imbalanced gradient norms between the low-probability classes (long-tail classes) and high-probability classes and propose to instead use reverse-KL minimization as objective. To define reverse-KL divergence, the proposed approach constructs a proxy Dirichlet distribution whose parameter is estimated from the ensemble predictions. The proposed EnD^2 variant with proxy Dirichlet target distribution is empirically validated on large-scale image classification and machine translation tasks.

The authors investigate ensemble distribution distillation for classification models with a high number of classes. This work is based on the work of Malinin et al.[1] which proposed EnD2 to use a Dirichlet prior over categorical predictive distributions to distill ensembles for classification tasks. The authors show through a synthetic example that the gradient norm ratio is very small for increasing number of data classes for the categorical cross-entropy (used in knowledge distillation) and Dirichlet KL (used in EnD2) loss. In contrast, the gradient norm ratio for the reverse KL for Dirichlet distributions maintains similarly for the increasing number of classes. Based on these findings, the author proposes to use a Proxy-Dirichlet target (by calculating the mean and approximating the precision of the ensemble) and use that for the reverse KL objective for ensemble distillation. They evaluate the approach on image classification (ImageNet) and sequence-to-sequence modeling (WMT'17 En-DE). In both settings, the model achieves competitive or better results for performance (Acc, Bleu) and calibration.