Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 8: Top 50% of accepted NeurIPS papers, clear accept
review_content: 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.




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.

That 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.

It 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.

There are a couple notable missing references:

* 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).

* 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.

** Update after author response ** The authors have addressed several of these concerns with additional experiments, so I am revising my recommendation upwards.

Minor issues

* L185: missing "for"
limitations_and_societal_impact: Yes.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.
main_review: Strengths
1. 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. 

2. The experimental results are quite extensive. The paper provides classification results for Imagenet and large-scale machine translation tasks. 

3. The paper is well written and easy to follow.

Weaknesses
1. 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.

2. 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.

3. 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.

Questions
1. 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.

2. 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.

3. 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?

4. 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?

5. Line 195, maybe "do not even begin" rather than "do ven begin"?





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_title: null

final_score: rating: 7: Good paper, accept
review_content: 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.
main_review: _Originality & Significance_

This 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.

_Quality & Clarity_

Overall, 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.

_Comments_

* 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?
* 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.
* 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.
* Can you discuss potential limitations and future directions?

Some typos:
* l.35 "this reduces to KL-divergence" -> this reduces to minimizing
* l.113 "it necessary" -> it is necessary
* l.462 missing reference

[1] Malinin, A., Mlodozeniec, B. and Gales, M., 2019. Ensemble distribution distillation. arXiv preprint arXiv:1905.00076.
[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.
[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.
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_title: null

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)