# Scaling Ensemble Distribution Distillation to Many Classes with Proxy Targets

arXiv: https://arxiv.org/abs/2105.06987 · HTML: https://mlsanity.com/paper/2105.06987/

## Key facts

- Cited by 19 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 43 references that are in the map.
- First posted on arXiv 14 May 2021.

## Builds on (most cited first, 30 of 33)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Neural Machine Translation of Rare Words with Subword Units](https://mlsanity.com/paper/1508.07909/) (2015)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [The Many Faces of Robustness: A Critical Analysis of Out-of-Distribution Generalization](https://mlsanity.com/paper/2006.16241/) (2020)
- [A Call for Clarity in Reporting BLEU Scores](https://mlsanity.com/paper/1804.08771/) (2018)
- [Concrete Problems in AI Safety](https://mlsanity.com/paper/1606.06565/) (2016)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [Natural Adversarial Examples](https://mlsanity.com/paper/1907.07174/) (2019)
- [Common Voice: A Massively-Multilingual Speech Corpus](https://mlsanity.com/paper/1912.06670/) (2019)
- [Can You Trust Your Model's Uncertainty? Evaluating Predictive Uncertainty Under Dataset Shift](https://mlsanity.com/paper/1906.02530/) (2019)
- [Sequence-Level Knowledge Distillation](https://mlsanity.com/paper/1606.07947/) (2016)
- [Albumentations: fast and flexible image augmentations](https://mlsanity.com/paper/1809.06839/) (2018)
- [Predictive Uncertainty Estimation via Prior Networks](https://mlsanity.com/paper/1802.10501/) (2018)
- [Adversarial Examples Are Not Easily Detected: Bypassing Ten Detection Methods](https://mlsanity.com/paper/1705.07263/) (2017)
- [BatchBALD: Efficient and Diverse Batch Acquisition for Deep Bayesian Active Learning](https://mlsanity.com/paper/1906.08158/) (2019)
- [ReZero is All You Need: Fast Convergence at Large Depth](https://mlsanity.com/paper/2003.04887/) (2020)
- [Pitfalls of In-Domain Uncertainty Estimation and Ensembling in Deep Learning](https://mlsanity.com/paper/2002.06470/) (2020)
- [Fixing the train-test resolution discrepancy](https://mlsanity.com/paper/1906.06423/) (2019)
- [Scaling Neural Machine Translation](https://mlsanity.com/paper/1806.00187/) (2018)
- [Understanding Measures of Uncertainty for Adversarial Example Detection](https://mlsanity.com/paper/1803.08533/) (2018)
- [Training independent subnetworks for robust prediction](https://mlsanity.com/paper/2010.06610/) (2020)
- [Ensemble Distribution Distillation](https://mlsanity.com/paper/1905.00076/) (2019)
- [Understanding Knowledge Distillation in Non-autoregressive Machine Translation](https://mlsanity.com/paper/1911.02727/) (2019)
- [Reverse KL-Divergence Training of Prior Networks: Improved Uncertainty and Adversarial Robustness](https://mlsanity.com/paper/1905.13472/) (2019)
- [Transformers with convolutional context for ASR](https://mlsanity.com/paper/1904.11660/) (2019)

## Built on by (most cited first, 19 of 19)

- [Shifts: A Dataset of Real Distributional Shift Across Multiple Large-Scale Tasks](https://mlsanity.com/paper/2107.07455/) (2021)
- [Shifts 2.0: Extending The Dataset of Real Distributional Shifts](https://mlsanity.com/paper/2206.15407/) (2022)
- [Distill n' Explain: explaining graph neural networks using simple surrogates](https://mlsanity.com/paper/2303.10139/) (2023)
- [Window-Based Early-Exit Cascades for Uncertainty Estimation: When Deep Ensembles are More Efficient than Single Models](https://mlsanity.com/paper/2303.08010/) (2023)
- [On the Usefulness of Deep Ensemble Diversity for Out-of-Distribution Detection](https://mlsanity.com/paper/2207.07517/) (2022)
- [Decoupled Training for Long-Tailed Classification With Stochastic Representations](https://mlsanity.com/paper/2304.09426/) (2023)
- [Self-Distribution Distillation: Efficient Uncertainty Estimation](https://mlsanity.com/paper/2203.08295/) (2022)
- [Conformal Prediction Sets for Deep Generative Models via Reduction to Conformal Regression](https://mlsanity.com/paper/2503.10512/) (2025)
- [Uncertainty Measures in Neural Belief Tracking and the Effects on Dialogue Policy Performance](https://mlsanity.com/paper/2109.04349/) (2021)
- [Applications of Knowledge Distillation in Remote Sensing: A Survey](https://mlsanity.com/paper/2409.12111/) (2024)
- [Advancing Deep Active Learning & Data Subset Selection: Unifying Principles with Information-Theory Intuitions](https://mlsanity.com/paper/2401.04305/) (2024)
- [Improving Ensemble Distillation With Weight Averaging and Diversifying Perturbation](https://mlsanity.com/paper/2206.15047/) (2022)
- [Logit-Based Ensemble Distribution Distillation for Robust Autoregressive Sequence Uncertainties](https://mlsanity.com/paper/2305.10384/) (2023)
- [Who Needs Decoders? Efficient Estimation of Sequence-level Attributes](https://mlsanity.com/paper/2305.05098/) (2023)
- [Accuracy-Preserving Calibration via Statistical Modeling on Probability Simplex](https://mlsanity.com/paper/2402.13765/) (2024)
- [A Confidence-based Acquisition Model for Self-supervised Active Learning and Label Correction](https://mlsanity.com/paper/2310.08944/) (2023)
- [Fast Ensembling with Diffusion Schr"odinger Bridge](https://mlsanity.com/paper/2404.15814/) (2024)
- [Knowledge Distillation of Uncertainty using Deep Latent Factor Model](https://mlsanity.com/paper/2510.19290/) (2025)
- [Unknown-Traffic Detection, Calibration and Shortcut Reliance in Distilled Encrypted-Traffic Classifiers over One Year](https://mlsanity.com/paper/2609.31141/) (2026)

## Abstract (arXiv)

Ensembles of machine learning models yield improved system performance as well as robust and interpretable uncertainty estimates; however, their inference costs may often be prohibitively high. Ensemble Distribution Distillation 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 criterion exhibits poor convergence when applied to large-scale tasks where the number of classes is very high. In our work, we analyze this effect and show that the Dirichlet log-likelihood criterion classes with low probability induce larger gradients than high-probability classes. This forces the model to focus on the distribution of the ensemble tail-class probabilities. We propose a new training objective that minimizes the reverse KL-divergence to a Proxy-Dirichlet target derived from the ensemble. This loss resolves the gradient issues of Ensemble Distribution Distillation, as we demonstrate both theoretically and empirically on the ImageNet and WMT17 En-De datasets containing 1000 and 40,000 classes, respectively.