Scaling Ensemble Distribution Distillation to Many Classes with Proxy Targets
Authors: Max Ryabinin
Andrey Malinin
Mark Gales
Venue: NeurIPS 2021 Poster (2021)
Keywords: Distribution Distillation
Uncertainty
Scaling
Large-Scale Tasks

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 (EnD$^2$) 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 \emph{Proxy-Dirichlet} target derived from the ensemble. This loss resolves the gradient issues of EnD$^2$, as we demonstrate both theoretically and empirically on the ImageNet, LibriSpeech, and WMT17 En-De datasets containing 1000, 5000, and 40,000 classes, respectively.

Scaling Ensemble Distribution Distillation to Many Classes with Proxy Targets
Max Ryabinin' Yandex, HSE University
Andrey Malinin Yandex, HSE University Moscow, Russia am969ttaandex-team.r
Mark Gales Cambridge-ALTA Institute Cambridge, United Kingdom mjf ggeng cam ac uk
Moscow, Russia
mryabininO@gmail.cmmm
Abstract
Ensembles of machine learning models yield improved system performance as well robust and interpretable uncertainty estimates; however, their inference costs be prohibitively high. Ensemble Distribution Distillation (EnD?) is an approach that allows single model efficiently capture both the predictive performance and uncertainty estimates of an ensemble. For classification, this is achieved by training 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 large-scale tasks where the number of classes is very high. Specifically, we show that for the Dirichlet log-likelln log-likelihood criterion classes with low probability induce larger gradients than high-pre probbbbbbbbblity 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 new training objective which target derived from the
minimizes the reverse KL-divergence to a Froxy-Dirichhe ensemble. This loss resolves the gradient issues of EnD?, as we demonstrate both theoretically and empirically on the ImageNet, LibriSpeech, and WMT17 En-De datasets containing 1000, 5000, and 40,000 classes, respectively.
Introduction
As machine learning models are being deployed ever more widely into production, especially in safety- critical 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 [ Ensemble methods are known to yield both improved predictive performance and reliable uncertainty estimates [2] R 4 5 Kl. Furthermore, ensembles allow separating out total uncertainty into data and knowledge uncertainty The former 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 detect well in active learning Α|
anomalous and unfamiliar inputs 8
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 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
"quual contribution, order decided by coin toss.
2 Data and Knowledge Uncertainty are also known as Alcatoric and Epistemic uncertainty.
35th Conference on Neural Information Processing Systems (NcurIPS 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 [ 14,
15 [ó]. These
can be broadly split into methods where the student has multiple output heads one per ensemble member [114 111.. and those which yield higher-order distribution over the ensemble's 's outputs [I3 [6]. 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 (EnD²), is the focus of this work.
Ensemble Distribution Distillation views outputs of each ensemble member as samples from higher-order Dirichlet or Normal-Wishart distribution, and the student model learns the parameters of that distribution. Typically, EnD is done by maximizing the likelihood of the ensemble's output distributions under the Dirichlet or Normal-Wishhe we show that for 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 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 result, the optimization procedure focuses on modelling this distribution of tail-class probabilities, preventing convergence.
prior. While this criterion is theoretically sound,
To solve this, we propose a two-step solution. First, we transform the empirical distribution of ensemble member predictions into Proxy-targer 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-ddvergeecc 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 iistribution-distillin 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 EnD to autoregressive models and refer to this approach as Sequence Ensemble Distribution Distillation (SenD²). On both datasets the citriiuttiondisstille 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 Se- quence Ensemble-Distribution Distillation (SEnD?).
2 Preliminaries: Ensembles and Distillation
We view ensembles within a Bayesian framework, where the model parameters @ are random variables over which prior p(0) is placed. The posterior distribution (D) is obtained via Bayes' rule:
p P( p(D)0)p(0)
P(D) p(D(0)p(0)
(1)
Consider an ensemble of models {P(y)æ>, ))
M sampled from the posterior:
{P(y|æ, 2r, 0 )
+(P(y)> (m))) de m ा {(m)_ f(x; ((m)), g(m) ((ODD
(2)
2

---

where π are the parameters of a categorical distribution [P(y wi).... P(y = wk]].. The predictive distribution, or predictive posterior, for a test imput x" is obtained by marginalization:
M
P(ylar a D) ppper (Pyyk=" 0) R3
Σ
P(ylæ, ((m))
(3)
M m==
In practice this is intractable and we approximate via Monte-Care sampling. Given the ensemble, the entropy of the predictive posterior a measure of total uncertainty. Knowledge uncertainty can be assessed via measures of the spread, or "disagreement" of the ensemble such as mutual information:
Ily 0 .D) DJ
H. (LEEEOO
[P(y/æ" .0]]]
(p(orD) [H[P[y) ((()(y)"",))) ,0]]
Knowledge Uncertainty
Total Uncertainty
Expected Data Uncertainty
An alternative measure of spread, which particularly useful for ensemble of autoregressive sequence models, is the reverse-me muteral information (RMI) [17
P(y|æ, D) P( xr 0)
M 0 x,, D] = Eppopd) (Pyy)a
In
(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 Pay|æ; o) by minimizing the KL-divergence to the predictive posterior of the ensemble:
( Dens) -Rp(a) KL (Lvlaa, a,, D) || P(y la; <)]
This approach has been applied to tasks such as image classification and machine translation 12 [8
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 [3]
Recently, a distillation technique called Ensemble Distribution Distillation (EnD²) was proposed B Here, both the mean and diversity of an ensemble are captured within a single model via distillation into a Prior Network [9. which a model that parameterizes the Dirichlet distribution:
K
br; Φ) -Dir(( a = =f:; Φ), ae > O, dn
Σ
O
Distribution distillation done two steps. Firstly, a transfer dataset Dana = (ce(1), T (1.1:M) p(ær, πT) composed of the inputs (:()) from the training set D (ax(1) ())))) Jim and the ensemble's output distributions (i.1::))) ji=z for each input. Secondly, given this transfer set, the model p(T [z; Φ) trained by minimizing the negative log-likelihood of each categorical distribution (im);
j
[EDD² ( Dene) = Ep(w)) [In az; Φ)]]
Given a crabbbttion-isstllle Prior Network, the predictive distribution is given by the expected categorical distribution á under the Dirichlet prior:
Py wel we|m" :=: Φ) pppmm/= (P(y = we(r)) =
Oe
Σk= 1= Ok
k==
E
Measures of total and knowledge uncertainty obtained by considering the mutual information between the prediction and the parameters of π of the categorical:
Tly π | Φ) = H
E (p(m(a))))) [M(((]]
(10)
Knowlodge Uncertainty
Total Uncertuinty
Expected Data Unvertainty
Alternttive measures, such as RMI, can also be obtained, as detailed in appendix C
An alternative approach to capture ensemble diversity
distill an ensemble into a mixture model
which yields a separate output head for each ensemble member [14] [5] However, the principal downside of this approach is that it attempts to model the ensemble in excessive detail, as 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, EnD has a fixed computational and memory cost of one model, regardless of the size of the original ensemble, as it models the ensemble's 's bulk distributional properties.
3

---

3 Analysis of first-order gradients
The previous section described how EnD can be done by maximising the log- legliikli of the en- semble's 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-basse based optimization, we examine the gradients of the Dirichlet negative ligeliiid loss and propose an alternative training approach.
Our analysis considers the following three situations. First, a Prior Network initialization which yields uniform Dirichlet distribution (a.lmit 1). The target categorical distribution whose probability is being maximized is a sparse K-length vector of probabilities:
# [1 - e, - 1), e/((K 1)....I"
le-4
(11)
Second, Prior Network near convergence with the following output distribution:
5c
conV
7 comY
Ccg. Crg = 9c4,
π cany
[1 5c,
(12)
K
K
Third, a Prior Network which has made a strong misclassification a situation which could occur in the middle of training, far from convergence:
5e K
mise
π
misc
Cku, Cryy 9e4, π mise
5c]*
(13)
K
We consider the following training criteria in this analysis: categorical cross entropy negative lielliielhood, Dirichlet Forward and Reverse KL--divergence
Dirichlet
CU-KLL
Σ
K
In
OCey
DDENNLL
K I((on)
(ax
Σ m=
M
In
Tg
(m)
r(ao)
M
k=1
k=1
K
Dir-KI
I(ax)
T(Bx) +T(B)) T(ao) +
(3x axx ((((3)) (b(B3)-wol v((w))
(14)
km1
kml
km
K
K
CDi-RKL_
Σ r(aa)
Σ +(ax)+(((o))) I(Bo)
aa Bxx Ba) (v((a)) rb(w))
k=1
A=1
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 EnD? [33. The third and fourth loss functions, the Dirichlet forward and reverse KI KI--divergence, in the latter part of this section. The targer Dirichlet distribution for these two loss functions has parameters s fo where fn- 1e5. Full derivations available in appendíx A
analyzed they relate our proposed solution
Let's examine the gradients of each criterion presented in equation 14 with respect to the logits We introduce the ratio of the gradient with respect to the logit 21, the high-probability class of interest, to the average of the absolute values gradients with respect the logits 22:K the tow-probability tail-classes. The ratio p represents the relative contribution of the gradients with respect the class we are interested in modelling to the long tail.
K K
OL Oz
a azy
(15)
2=
We normalize by number of classes K, this improves interpretability by making p of Ka plateau. This makes it casier to compare to the behaviour of p for other criteria.
Figure T shows that, at initialization, the number of classes is increased, the categorical cross- entropy Ioss primarily focuses on the high probability class and ignores the long tail. 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-brobability classes only after it perfectly
A Which is equivalent to categorical Forward -ivveregeee
Wee remind the reader that the on =a. where 24 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 i far smaller for D than for categorical cross-entropy. Finally, if significant enror is made, / remains large for cross-entropy, but becomes very small for Dirichlet NL.L as the number of classes increases. This shows that 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.
16
10
Gradiemt Norm Ratio
10"
CarKI
Cat-K
Kl.
10" Di-NLLL
Gradien Norz
DiNNLL
10
Di-NLL
10
to Di-AKL.
Di-RKL
Dr-RKL
10 D2-RXL a-1 P+]
Di -KL 0+1 P-)
10
Di XXL 0v1 (-)
10
10'
10'
10' Cf
Number Clussss
(a) Initialization
(b) Near Convergence
(c) Misclassification
Figure 1: Gradient ratio analysis of distillation and distribution distillation criteria.
Figure D also displays the ratio p 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 p continues decrease as the number classes in increased. This unsurprising, as the negative log-likellng likelihhodd is equivalent to the KL-divergence in the limit. However, the reverse KL-divergence displays the desirable properly that o grows and stabilizes as the number of classes is increased 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.
4 Distribution Distillation via Proxy-Dirichlet target distribution
Max Likelinood
Estimate
Ensemble
Model
Ensemble
Proxy Dirchlet
Model
(a) Original EnD?
(b) Proposed
EnD?
Figure 2: Schematic of Original and Proposed Distribution Distillation Approaches
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-probabilitt classes. Thus, if we could specify an appropriate target Proxy Dirichlet distribution which represents the ensemble, then could accomplish distribution- distillation without optimization difficulties.
To correctly construct a Proxy-Dirichlet target which represents the ensemble, we need to obtain mean ीg and precision 3o 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 200 and proposes the following estimate:
fk (K
+x)
Σ P(y wa/ze, O(m)), Ba(x))
M
2 Σ K
कN (In A DE EM In o.
(m)
(16)
Stirling's ' Approximation
Thus, the Proxy Dirichlet has both the same mean and bulk-diversity properties the ensemble, such that the measures of diversity derived from the Proxy Dirichlet are similar to those of the original
5

---

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, ag. and target parameters, SR. are constrained to be greater than 1. Figure | shows that adding this constraint to the concentration parameters of the Proxy-Dirichlet and the model improves p. especially near convergence. Altice...Algge Altogether, this yields the following criterion:
C.D-EaD" -K[ppTT/a ()]l[(+|(++]) x Bp(mla, tu In
KL[p( (1, p(P(((
80
(17)
Reconstruction cerm
By dividing through by 9o the above criterion is made to have the same form as ELB This serves two purposes. First, becomes apparent that the reconstruction term is an upper bound on categorical cross entropy, which explains why they share similar properties. Second, normalizing by 3n 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 similar to the one proposed I0 the underlying reason for its use is entirely different. Here, the issue is due to large gradients from tail classes, while 100 uses a reverse-KI. criterion avoid inducing expectation multi-mode target Dirichlet distribution.
5 Experiments
In this section, we evaluate EnD² via minimization of the reverse KL-divergence to Proxy Dirichlet target, which we refer to as PD-EnD2. We apply PD-EnD² ensembles of convolutional networks trained on the ImageNet [21 dataset, ensembles of VGG-Transformer [22] ASR models trained on LibrSpeech [23] and ensembles Transformer-big A models trained on WMT' 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 previously impossible. We do not provide results for EnD? accomplished by optimizing Dirichlet negative log-likelihood or Dirichlet forward KL-divergence because these approaches do not even begin to converge the tasks considered. Finally, 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
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 Deep Ensemble [4] of regular models and then distribution-distilll with PD-EnD We consider the following baselines:
Single refers to the average performance across 10 individual, independently evaluated models.
Ensemble refers the performance Deep Ensemble, which is currently considered to both
the best and the simplest ensemble method 5 a
EnD or SEnD refer to ensemble distillation via minimizing the KI. divergence between the student and the ensemble's s mean for image classification 12 and machine translation [8] respectively.
We do not use multi-head approaches
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 unclear how to extend them to structured tasks. In all experiments with PD-EnD we add both to the predicted parameters of the Dirichlet distribution and the Dirichlet proxy parameters, as discussed in the previous section.
To do distribution- distribution-diitillation for NMT, which is an autoregressive structured prediction task, we distillation extend Prior Networks and EnD² such models and develop Sequence Ensemble Distribution Distillation (SEnD?). Models trained with Proxy-Dirichlet targets are referred to as PD-SEnD2
$ A full derivation available appendix CCode and training configurations are available at http:://gthub..co//aaaeeeeerr

---

This extension involves modifications to both the model and derivation of associated sequence-level uncertainty measures [77. While conceptually straightforward. this modification is extensive and mathematically involved, and therefore instead fully detailed in appendix
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 range of distributionally shifted, or out-of-distribun perform pootly, or not to be able perform at all. The goal is be able to discriminate between the in-domain and out-of-distribus successfully ietrribttionddstillle model will yield similar performance to the ensemble, especially using measures of knowledge uncertainty.
datasets on which we expect the models to either
datasets using measures of uncertainty. We expect that a
Large-scale image classification
We begin by examining PD-EnD² on the ImageNet [2]] dataset, which features 1000 classes. This represents 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 [B. Here, we distribution-disiie an ensemble of 10 ResNet-50 26 models. To train each ResNet-50 model, we use the standard training setup outlined in 22 28. Specifically, train for 90 epochs using stochastic gradient descent with momentum of 0.9 and learning rate of 0.1 x B/256, where B is the per-device batch size multiplied by the number of GPUs. In our experiments, we use single-GPP 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
In all experiments, we
found it beneficial to initialize the last batch normalization y in each residual branch to zero, which agrees with previous results 28 30 31. Additional results are provided in appendix
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 [2]] and on a range of distributionally shifted datasets. Specifically, we use natural adversarial examples from ImageNet-A [2), corrupted and perturbed versions of the ImageNet validation data from ImageNet- C (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 show that EnD? is capable of accurate emulation of the ensemble in terms of classification performance. In terms of accuracy, the method displays results on par 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 closest to the Ensemble. Furthermore, on the in-domain set they outperform the ensemble. These results clearly demonstrate that PD-EnD2 is successfully able produce models which are comparable to quality to ensemble distillation, but are better calibrated.
Table 1: Prediction quality results for image classification.
ImageNet-val
ImageNet-A
ImageNet-C.
ImageNet-R
Model
Acc
ECE
Acc
ECE
Acc
BCE
Acc
FCE
Single
75.
4.8
4.420.2
51.1
39.1
11.3
35.0+0.2
21.3
EnD
77.0_a.1
1.620.1
3.720.1
46.8_n.2
40.6+n4
5.8
37.0+0.3
15.8
2D-EnD2
77.1
1.640.1
3.740.2
42.9-m.2
40.6+0.5
4.7
36.9
12.1
Ensemble
79.0ANA
2.3
3.9
42.0aNA
43.5
4.5
38.8
9.8
We evaluate the measures of uncertainty produced by the proposed approach and the baselines on the task of outtt--distiibtttin detection. As discussed above, the goal is to discriminate between in- domain and distributionally shifted (OOD) data using measures of uncertainty. We consider ImageNet
7

---

80
Single
Single
15
Ensemble
Ensemble
Accuracy
60
EnD
BCB
EnD
10
PD-EnD?
PD-EnD
40
5
2
5
0
4
5
Corruption degree
Corruption degree
(a) Accuracy
(b) Expected Calibration Error
ROC AUC (%)
90
Single
ROC ROC AUC (%) (%))
90
Ensemble
Ensemble
Proxy
Proxy
PD-EnD²
80
EnD
PD-EnD2
70
70
2
5
0
4
5
Corruption degree
Corruption degree
(c) ROC-AUC w/ Total Unc.
(d) ROC-AUC w/ Knowledge Unc.
Figure 3: Performance of image classification models depending the level of ImageNet-C corrup tion. No corruption corresponds to the original ImageNet validation data.
A,C and R as well as ImageNet-O (32,, which features altogether new classes not seen in ImageNet, as
data. We use the entropy of the predictive posterior and mutual information (equation (T] as estimates of total and knowledge uncertainty, respectively. 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 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 able obtain performance close to that of the ensemble and generally yields similar behaviour. This clearly demonstrates that PD-EnD2 can successfully train model emulate an ensemble.
Table 2: Out-of-distribution detection results for image classification.
ImageNet-O
ImageNet-A
ImageNet-C
ImageNet-R
Total
Knowledge
Total
Knowledge
Total
Knowledge
Total
Knowledge
Single
50.7.10.3
85.8
79.9
83 13.0
EnD
48.810.6
87.2
80.810.3
83 1..2
2D-EnD2
52.1
53.2_g.3
86.8
84.7
80.2
77.4
83.7 83.740.1
81.9
Ensemble
54.64NA
61.2
88.8
87.2
82.0
78.1
86.1
84.4
Proxy
54.6
56.2
89.0
86.5
82.0=NA
81.1
86.1
86.3
In addition to the tables above, we provide Figure which shows detailed breakdown by corruption strength of the predictive and OOD detection performance ImageNet-C. Results show that as the degree of corruption increases, all models suffer a drop accuracy and calibration quality; notably, EnD and 2D-EnD2 have the same calibration performance original data, but PD-EnD2 has lower calibration errors for the highest degrees corruption. Unsurprisingly, the greater the distributional shift from the original training images, the better the models out-of-distribttion detection. Overall, the distribution-diittlledd model clearly follows the behaviour the ensemble in all metrics.
Speech Recognition

---

Table 3: Predictive performance results for speech recognition in terms of WER and NLL.
LTC
ITTO (ROC-AUC %)
AMI (ROC-AUC %)
C-FR (ROC-AUC %)
Model
Total
Knowledge
Total H
Knowledge
Total H
Knowledge
WER
NLL
M
I
K
M
I
IC
M
I
K.
Single
6.1
0.35
72.8
94.4
99.8
SEnD
6.2
0.36
72.8
90.2
99.8
PD-SEnD2
7.1
0.25
76.6
75.5
50.0
75.7
94.4
75.2
50.0
76.8
99.9
98.8
50.0
98.9
Ensemble
4.2
0.20
76.9
77.0
76.6
77.0
96.5
94.8
94.9
94.9
99.9
99.9
99.9
99.9
Now let's consider applying PD-EnD² to speech recognition. We distill an ensemble of 6 VGG- Transformer models [22] trained on the English speech recordings from LibriSpeech [23]. Au- toregressive speech recognition systems are challenging to train even in a standard setting the VGG-Transformer architecture tends to be unstable in training. Furthermore, speech models tend to yield extremely sharp, over-confident predictions, which both exacerbates the issues with tail-class gradients dominating the gradients from high- probb-bbbbbiltty classes (when training with naive EnD²), ) and also causes numerical issues when distilling using Proxy Dirichlet. This makes ASR interesting "stress-test" task for assessing the proposed solution for Ensemble Distribution Distillation.
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 constant learning rate. The dataset is processed with vocabulary of 5000 sentence-piece tokens. Checkpoints over the last 30 epochs are averaged together. The only significant difference is that we train SEnD² with the Adam K7 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 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 and test-set negative log- ligliiiel 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-diistlleed model (PD-SEnD²)) 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-SEnD?, this suggest that the model underfit, and 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 EnD? 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-distilledd model is able 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 the WMT'17 English-German German machine translation dataset with 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 4ill trained with the setup described in [2)] Specifically, ensemble members are trained for 193,000 steps with Adam si 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 for inference.
'Computed using NIST SCLITE tool.

---

Table 4: BLEU and out-of-distribttion detection results for machine translation.
Model
BLEF
NLL
NLL Ens.
Permuted
Speech
German
French
Total
Know.
Total
Know.
Total
Know,
Total
Know.
Single
28.8
1.46
0.42
80.7a1.5
73.7
32.8-za
27.1asa
Ensemble
1.33
0.39
83.74XA
97.4.A
67.84NA
73.7-M
39 SANA
82.4
25. 25.0
73.6.NA
SEnD
29.4
1.52
0.4
79.521.1
75.9
25.4-16
15.6
PD-SEnD"
29.5
1.52
0.4
78.3
97. 17.s
77.040.a
28.5
38.321.
20.9
15.9
60. 12s.
First, we assess the proposed approach in terms of translation quality on the newstest| English- German test set, for which we use the BLEU score A computed with SacreBLEU 4. Results
in Table Table confirm that PD-EnD² 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 '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 more complex and multimodal, more surprising, therefore
in 4S that reference transcriptions would have higher NLLs, while predictions of NMT models are more consistent in style and therefore simpler -hhy 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-distilee model yields. For
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 [7] The in-domain test set is newstest' 14 and OOD datasets are as follows: 'permuted' sentences with permuted tokens in the input, 'speech' - LibriSpeech [X] 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 were considered by [17 Results of both distillation methods are averaged over random seeds.
Table 4 clearly shows that in all cases, the PD-SEnD² 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-din distilled model doesn't 't follow the behaviour of the ensemble closely, which is not surprising, as the task is far more complex than image classification and the output space is exponentially vast.
6 Conclusion
This work has examined applying Ensemble Distribution Distillation (EnD²) to large-scale tasks where the number of classes is very high. Training these models 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 from the ensemble. This new criterion resolves the gradient issues of EnD?, 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.
Proxy-Dirichler target derived
Further details and derivations available in appendix
10

---

References
[1] Dario Amodei, Chris Olah, Jacob Steinhardt, Paul F. Christiano, John Schulman, and Dan Mané, "Concrete problems in Al safety," 1606.06565.
06565 2016, arXiv:
[2] Thomas G. Dietterich, "Ensemble methods in machine learning." " in Proceedings the First International Workshop on Multiple Classifter Systems, Berlin, Heidelberg, 2000, MCS '00, 1-15, Springer-Verlag.
[3] Yarin Gal and Zoubin Ghahramani, "Dropout as Bayesian Approximation: Representing Model Uncertainty in Deep Learning," in Proc. 33rd International Conference on Machine Learning (ICML-16), 2016.
[4] B. Lakshminarayanan, A. Pritzel, and C. Blundell, "Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles," in Proc. Conference on Neural Information Processing
Systems (NIPS), 2017.
[5] Arsenii Ashukha, Alexander Lyzhov, Dmitry Molchanov, and Dmitry Vetrov, "Pitfalls of in- domain uncertainty estimation and ensembling in deep learning." in International Conference on Learning Representations, 2020.
[6] Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, D Sculley, Sebastian Nowozin, Joshua V Dillon, Balaji Lakshminarayanan, and Jasper Snoek, "Can you trust your model's uncer- tainty? evaluating predictive uncertainty under dataset shift," Advances in Neural Information Processing Systems, 2019.
[7] Andrey Malinin, Uncertainty Estimation in Deep Learning with application Spoken Language Assessment, Ph.D. thesis, University of Cambridge, 2019.
[8] L. Smith and Y. Gal, "Understanding Measures of Uncertainty for Adversarial Example Detection," in UAI, 2018.
[9] Nicholas Carlini and David A. Wagner, "Adversarial examples are not easily detected: Bypassing ten detection methods," " CoRR, 2017.
[10] Andrey Malinin and Mark Gales, "Reverse kl-divergence training of prior networks: Im- proved uncertainty and adversarial robustness," 2019.
[1]] Andreas Kirsch, Joost van Amersfoort, and Yarin Gal, "Batchbald: Efficient and diverse batch acquisition for deep bayesian active learning," 2019.
[12] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean, "Distilling the knowledge in neural network." 2015, arXiv:1503.0255
[13] Andrey Malinin, Bruno Mlodozeniec, and Mark JF Gales, "Ensemble distribution distillation." in International Conference Learning Representations, 2020.
[14] Linh Tran, Bastiaan S. Veeling, Kevin Roth, Jakub Šwiętkowski, Joshua V. Dillon, Jasper Snoek, Stephan Mandt, Tim Salimans, Sebastian Nowozin, and Rodolphe Jenatton, "Hydra: Preserving ensemble diversity for model distillation," 2020.
[15] Xixin Wu, Kate M Knill, Mark JF Gales, and Andrey Malinin, "Ensemble approaches for uncertainty in spoken language assessment," Proc. Interspeech 2020, pp. 3860-3864, 2020.
[16] Andrey Malinin, Sergey Chervontsev, Ivan Provilkov, and Mark Gales, "Regression prior networks," arXiv preprint 2nXiv:2006.11590 2020.
[17] Andrey Malinin and Mark Gales, "Uncertainty in structured prediction," arXiv preprint arXiv:2002.0m 07650, 2020.
[18] Yoon Kim and Alexander M Rush. "Sequence-level knowledge distillation," in Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, 2016, pp. 1317-1327.
[19] Andrey Malinin and Mark Gales, "Predictive uncertainty estimation via prior networks." in Advances in Neural Information Processing Systems, 2018, pp. 7047-7058.
[20] Thomas Minka, "Estimating a dirichlet distribution," " 2000.
[21] Deng, W. Dong, R. Socher, L.-J. Li, Li, and I.. Fei-Fei, "ImageNet: Large-Scale Hierarchical Image Database," in CVPR09, 2009.

---

[22]
AbdeIrahman Mohamed, Dmytro Okhonko, and Luke Zettlemoyer, "Transformers with convo- lutional context for asr," arXiv preprint 1rX4v:1904. 1666, 2019.
[23] Vassil Panayotov, Guoguo Chen, Daniel Povey, and Sanjeev Khudanpur, "Librispeech: an ast corpus based on public domain audio books." " in 2015 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE, 2015, pp. 5206-5210.
[24] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin, "Attention is all you need," in Advances in neural information processing systems, 2017. pp. 5998-6008.
[25] Marton Havasi, Rodolphe Jenatton, Stanislav Fort, Jeremiah Zhe Liu, Jasper Snoek, Balaji Lakshminarayanan, Andrew M Dai, and Dustin Tran, "Training independent subnetworks for robust prediction," arXiv preprint arXiv:2010.06610) 2020.
[26] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun, "Deep residual learning for image recognition." " in Proceedings of the IEEE conference on computer vision and paltern recognition, 2016, pp. 770-778.
[27] Hugo Touvron, Andrea Vedaldi, Matthijs Douze, and Hervé Jégou, "Fixing the train-test resolution discrepancy," in Advances in Neural Information Processing Systems (NeurIPS), 2019.
[28] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jía, and Kaiming He, "Accurate, large minibatch sgd: Training imagenet in 1 hour," 2018.
[29] A. Buslaev, A. Parinov, E. Khvedchenya, V. I. Iglovikov, and A. A. Kalinin, "A "Albumentations: fast and flexible image augmentations," ArXiv e-prints, 2018.
[30] Hongyi Zhang, Yann N. Dauphin, and Tengyu Ma, "Residual learning without normalization via better initialization," in International Conference on Learning Representations, 2019.
[31] Thomas Bachlechner, Huanru Henry Majumder, Bodhisattwa Prasad Mao, Garrison W. Cottrell, and Julian McAuley, "Rezero is all you need: Fast convergence large depth." " in arXiv, 2020.
[32] Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song, "Natural adversarial examples," CVPR, 2021.
[33] Dan Hendrycks and Thomas Dietterich, "Benchmarking neural network robustness to common corruptions and perturbations," Proceedings of the International Conference on Learning Representations, 2019.
[34] Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, Dawn Song, Jacob Steinhardt, and Justin Gilmer, "The many faces of robustness: A critical analysis of dui-o--dist generalization," arXiv preprint arXiv:2006. 16241, 2020.
[35] Joaquin Culhonero-Candele Dataset Shift in Machine Learning, The MIT Press, 2009.
[36] Matthew D Zeiler, "Adadelta: an adaptive learning rate method," arXiv preprint
12115:121 2012.
[37] Diederik P. Kingma and Jimmy Ba, "Adam: A Method for Stochastic Optimization," in Proc. 3rd International Conference on Learning Representations (ICLR), 2015.
[38] Wessel Kraaij, Thomas Hain, Mike Lincoln, and Wilfried Post, "The ami meeting corpus," 2005.
[39] Rosana Ardila, Megan Branson, Kelly Davis, Michael Henretty, Michael Kohler, Josh Meyer, Reuben Morais, Lindsay Saunders, Francis M Tyers, and Gregor Weber, "Common voice: A massively-muttiiigguue multilingual speech corpus," arXiv preprint arXiv: /91X:v::912.0 19119...6660. 2019.
[40] Rico Sennrich, Barry Haddow, and Alexandra Birch, "Neural machine translation of rare words with subword units," in Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), Berlin, Germany, Aug. 2016, pp. 1715-1725, Association for Computational Linguistics.
[41] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin, "Attention is all you need," 2017.
[42] Myle Oit, Sergey Edunov, David Grangier, and Michael Auli, "Scaling neural machine transla- tion." arXiv preprint arXiv: 00t.:0000600.77.. 2018.
12

---

[43] Kishore Papineni, Salim Roukos, Todd Ward. and Wei-Jing Zhu, "Bleu: method for automatic evaluation of machine translation." in Proceedings of the 40th annual meeting of the Association for Computational Linguistics, 2002, pp. 311 311-318.
[44] Matt Post, "A call for clarity in reporting BLEU scores," in Proceedings of the Third Confer- on Machine Translation: Research Papers, Belgium, Brussels, Oct. 2018, pp. 186-191, Association for Computational Linguistics.
[45] Chunting Zhou, Graham Neubig, and Jiatao Gu, "Understanding knowledge distillation in 2019.
von-auvoregressivve machine translation," arXiv preprint
13


Reviews:
[{"review_id": "gxbcNFS0pNf", "reviewer": "Reviewer_vef7", "summary": "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "_Originality & Significance_\n\nThis work scales ensemble distribution distillation to a large number of classes which is highly relevant for large-scale modeling. This is very impactful as there are not many works that successfully distilled ensembles of large models and preserves uncertainty. Especially in safety-critical applications, uncertainty estimates are necessary to have some measures of reliability. The large-scale evaluation on ImageNet and WMT show that this approach scales to these datasets/models. This advancement is great for the (research) community.  From the details provided and looking at the supplementary materials, the ensemble distillation for image classification can be reproduced. However, I have doubts about the reproducibility of the auto-regressive model as code has not been provided. There is no discussion about limitations and potential negative impacts which I believe the author should add to make the work complete.\n\n_Quality & Clarity_\n\nOverall, this work is well executed and well-written. The authors give justification why existing knowledge distillation objectives do not scale to data with a large number of classes and propose the reverse KL, which works for these problems. The evaluation is described in detail and results are weighted and discussed. This feels like work very close to completion, and there is not much that I can complain about.\n\n_Comments_\n\n* Can you discuss whether there are alternatives to Sterling's approximation for the Dirichlet precision, how close the approximation comes to the true precision, and whether the approximation might hurt overall distillation?\n* Section 3: This is more of an editorial opinion. I found the title \"Theoretical analysis\" a bit of a stretch. You derive the gradients and the objectives and show through a synthetic example how the gradient norm ratio behaves. I am not sure if this qualifies as theoretical analysis.\n* Concurrently, Beyer et al.[2,3] also investigates ensemble knowledge distillation from different points of view. As these two works are relevant and related, it might be worth briefly discussing their works in the part of the related work.\n* Can you discuss potential limitations and future directions?\n\nSome typos:\n* l.35 \"this reduces to KL-divergence\" -> this reduces to minimizing\n* l.113 \"it necessary\" -> it is necessary\n* l.462 missing reference\n\n[1] Malinin, A., Mlodozeniec, B. and Gales, M., 2019. Ensemble distribution distillation. arXiv preprint arXiv:1905.00076.\n[2] Beyer, L., Zhai, X., Royer, A., Markeeva, L., Anil, R. and Kolesnikov, A., 2021. Knowledge distillation: A good teacher is patient and consistent. arXiv preprint arXiv:2106.05237.\n[3] Stanton, S., Izmailov, P., Kirichenko, P., Alemi, A.A. and Wilson, A.G., 2021. Does Knowledge Distillation Really Work?. arXiv preprint arXiv:2106.05945.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "This is one of the things that the authors can improve on because they did not discuss limitations and societal impact at all.", "review_text": "", "extra_scores": {"needs_ethics_review": "Yes", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "Y4Q15Lszwkv", "reviewer": "Reviewer_eKFR", "summary": "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.\n\n\n\n", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, I found the paper to be well-written, with a reasonable motivation for the main contribution, PD-EnD^2, and a convincing analysis and explanation of the proposed approach. To the extent that vanilla EnD^2 *does not work* on the class of problems considered in this work, and with distillation a requirement for deep ensembles to be practical in many settings, this paper fills an important gap in the literature.\n\nThat being said, the experimental comparisons were somewhat lacking, as discussed further below. Notably, MIMO has all of the benefits of the proposed approach (albeit with several disadvantages), which make it an interesting baseline. In addition, there is related work using (vanilla) ensemble distillation for NMT that may warrant experimental comparison, given that the sequence-level extension of PD-EnD2 is a non-trivial contribution of this work.\n\nIt is stated that \"we do not provide results for EnD2\" because \"these approaches do even begin to converge on the tasks considered.\" I think it would have improved the submission to characterize the point at which vanilla EnD2 starts to fail, for example by varying the number of classes, or even using a simple synthetic example. It would be interesting to know, for problems with fewer classes, whether PD-EnD^2 always outperforms vanilla EnD^2.\n\nThere are a couple notable missing references:\n\n* The paper mentions the alternate approach of capturing ensemble diversity using separate output heads for each ensemble member, but doesn't cite MIMO (https://arxiv.org/pdf/2010.06610.pdf). Notably, MIMO has similar computational cost for small ensembles as a single network, and each network-specific head only introduces a small amount of parameters. Therefore, I don't think the objections raised apply to MIMO, and it would have been nice to see experimental comparisons against MIMO on ImageNet (https://github.com/google/uncertainty-baselines).\n\n* Ensemble distillation (vanilla, not EnD2) has been applied to NMT (https://arxiv.org/abs/2010.06721), where a truncation is proposed to avoid modeling the long-tail. This has the practical implication of allowing for caching of the predictive distribution of the ensemble, obviating the need to hold all ensemble members in memory as is done in this work (which compensates for it with 32 steps of gradient accumulation). In the context of the EnD2 approach used here, the truncation could have another benefit by helping the Dirichlet NLL objective avoid trying to perfectly model the long-tail, and therefore may merit some consideration / discussion. In addition, the reported results could be compared to those in that work.\n\n** Update after author response ** The authors have addressed several of these concerns with additional experiments, so I am revising my recommendation upwards.\n\nMinor issues\n\n* L185: missing \"for\"", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "2aAD0XInQo", "reviewer": "Reviewer_io27", "summary": "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n1. Distilling ensemble models is an important task. As authors pointed out, many baselines either try to improve student networks by increasing the number of parameters (e.g., introducing multi-headed branches). I personally think that the EnD^2 framework is conceptually and theoretically attractive, but empirically found it to be less scalable. The solution provided in this paper can make the EnD^2 framework scale. \n\n2. The experimental results are quite extensive. The paper provides classification results for Imagenet and large-scale machine translation tasks. \n\n3. The paper is well written and easy to follow.\n\nWeaknesses\n1. The novelty is somewhat limited. I agree with the author's statement that the proposed reverse KL objective is derived for different purposes (dealing with long-tail classes) compared to the previous work (improving prior networks with reverse-KL objective). Still, the solution provided in this paper is an incremental application of the reverse KL objective (+ proxy Dirichlet ) to the EnD^2 framework.\n\n2. Although the title of section 3 reads as \"Theoretical analysis\", the actual argument is mainly empirical. The authors analyze the gradient norm ratios between a high-probability class and the others by empirically computing them for three different situations (at initialization, near convergence, and far from convergence) with different losses. Although this analysis has provided enough empirical evidence, it does not provide a deeper understanding of what is going under the hood; for instance, as Figure 1 shows, Dir-RKL with alpha +1 and beta+1 behaves better than the others. Why is this the case? Also, if the gradient ratio rho is affected by the value of the Dirichlet hyperparameters alpha and beta, can Dir-KL or Dir-NLL also be improved by different alpha or beta settings? What about the value of epsilon? In current form, the paper suggests using the reverse KL because empirically it behaved the best under the specific setting considered in section 3. It would be great to have a minimal theoretical analysis (at least with a simplified model) to justify the bad gradient ratio behavior of the Dir-NLL loss and provide a rationale to choose reverse KL instead.\n\n3. The experimental results are extensive but can be strengthened.  Personally, I find evaluating models solely based on ECE to be unreliable., so it would be good to compare the models using various uncertainty assessment metrics such as negative log-likelihoods, MCE, or Brier scores.  Also, for the baselines (especially ensembles or EnD), the calibration errors can further be improved via temperature scaling. I'm not sure whether the temperature scaling is applicable to the proposed method, but if possible it would be best to compare the calibration errors of the baselines and the proposed method under optimally scaled temperatures.\n\nQuestions\n1. Have you considered using different prior distributions for the classification distribution? For instance, the stick-breaking approximation of the Pitman-Yor process (which is known to better model the long-tail classes than Dirichlet processes), logistic normal distribution, or logistic transform of heavy-tailed distributions. I'm quite skeptical though because judging from the paper EnD^2 suffers from the instability issue because it exaggerates gradients for long-tail classes while the alternative I suggested generally put more mass on long-tail parts of the distributions.\n\n2. In experiments, ensembles of 10 ResNet-50 models are considered for the ImageNet task. How does the proposed method behave with different numbers of ensemble models?  As far as I know, people usually consider five or fewer models for ensembles for ImageNet tasks, because the number of ensembles beyond it is prohibitively expensive. I wonder if the proxy Dirichlet distribution is well constructed with fewer ensemble members.\n\n3. Comparing forward and reverse KL objectives, it is generally known that the forward KL ($\\mathrm{KL}[p_\\text{target}\\Vert p_\\text{model}]$) tends to capture the overall mass of the target distribution while the reverse KL ($\\mathrm{KL}[p_\\text{model}\\Vert p_\\text{target}]$) tends to capture specific modes. In terms of ensemble distillation, isn't this mode-capturing property of reverse KL bad for capturing diversity?\n\n4. Line 150 states that \"if a significant error is made, $\\rho$ becomes very large for cross-entropy\", but in Figure 1 (c) the $\\rho$ values for Cat-KL is similar ($10^0$) with the values for other cases. Can you elaborate on this?\n\n5. Line 195, maybe \"do not even begin\" rather than \"do ven begin\"?\n\n\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The checklist says that the limitations are described but I find is less apparent. I agree that the content discussed in the paper is unlikely to raise any negative societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers are all for accepting the work. It's easy to follow and has quite extensive experiments, especially with the help of the MIMO baseline in the rebuttal (please cite and add the additional experiments into the paper). The loss function involves several moving parts, as well as pipelines like distilling from a trained model with much longer training than the baselines they compare to. So the idea isn't necessarily easy to implement. With that said, I think the work's ideas and its empirical evaluation will be of interest to the community.
Decision: Accept (Poster)