Scaling Ensemble Distribution Distillation to Many Classes with Proxy Targets

Abstract

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. \emph{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 \emph{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.

References in corpus (9)

• Distilling the Knowledge in a Neural Network
• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• Can You Trust Your Model's Uncertainty? Evaluating Predictive Uncertainty Under Dataset Shift
• BatchBALD: Efficient and Diverse Batch Acquisition for Deep Bayesian Active Learning
• ReZero is All You Need: Fast Convergence at Large Depth
• Reverse KL-Divergence Training of Prior Networks: Improved Uncertainty and Adversarial Robustness
• Hydra: Preserving Ensemble Diversity for Model Distillation
• Uncertainty Estimation in Autoregressive Structured Prediction
• Regression Prior Networks

Cited by in corpus (3)

• Shifts: A Dataset of Real Distributional Shift Across Multiple Large-Scale Tasks
• A Confidence-based Acquisition Model for Self-supervised Active Learning and Label Correction
• Uncertainty Measures in Neural Belief Tracking and the Effects on Dialogue Policy Performance