PaperlayerView exact revision on arXiv

Source-linked AI summary

Ensemble Distribution Distillation

Andrey Malinin, Bruno Mlodozeniec, Mark Gales

arXiv:1905.00076v3stat.MLcs.LG

TL;DR

Ensembles improve accuracy and uncertainty estimation but are expensive, and standard distillation loses their predictive diversity. This paper introduces EnD2, which uses Prior Networks to distill the distribution of ensemble predictions into one model; across artificial and image tasks, it approaches ensemble classification performance and improves misclassification and OOD detection over standard DNNs and Ensemble Distillation.

  • Problem

    Standard ensemble distillation reduces computation but collapses ensemble diversity, limiting the preservation of information useful for distinguishing uncertainty sources.

  • Method

    EnD2 uses a Prior Network to model the distribution over ensemble predictions as a conditional distribution over categorical distributions.

  • Results

    Across artificial data and CIFAR-10, CIFAR-100, and TinyImageNet, EnD2 approaches ensemble classification performance and outperforms standard DNNs and Ensemble Distillation for misclassification and OOD detection.

  • Takeaways & Limitations

    A single model can retain ensemble classification performance together with information about ensemble diversity for uncertainty estimation.

  • Takeaways & Limitations

    The approach can fail to capture knowledge uncertainty far from training data, and distilling out-of-distribution behavior from in-domain observations alone is challenging.

Abstract

from arXiv · show

Ensembles of models often yield improvements in system performance. These ensemble approaches have also been empirically shown to yield robust measures of uncertainty, and are capable of distinguishing between different \emph{forms} of uncertainty. However, ensembles come at a computational and memory cost which may be prohibitive for many applications. There has been significant work done on the distillation of an ensemble into a single model. Such approaches decrease computational cost and allow a single model to achieve an accuracy comparable to that of an ensemble. However, information about the \emph{diversity} of the ensemble, which can yield estimates of different forms of uncertainty, is lost. This work considers the novel task of \emph{Ensemble Distribution Distillation} (EnD$^2$) --- distilling the distribution of the predictions from an ensemble, rather than just the average prediction, into a single model. EnD$^2$ enables a single model to retain both the improved classification performance of ensemble distillation as well as information about the diversity of the ensemble, which is useful for uncertainty estimation. A solution for EnD$^2$ based on Prior Networks, a class of models which allow a single neural network to explicitly model a distribution over output distributions, is proposed in this work. The properties of EnD$^2$ are investigated on both an artificial dataset, and on the CIFAR-10, CIFAR-100 and TinyImageNet datasets, where it is shown that EnD$^2$ can approach the classification performance of an ensemble, and outperforms both standard DNNs and Ensemble Distillation on the tasks of misclassification and out-of-distribution input detection.

1 INTRODUCTION

Neural networks perform well but can be over-confident and lack uncertainty estimates, while ensembles improve accuracy and characterize distinct uncertainty sources at higher computational cost. EnD2 addresses the loss of ensemble diversity during distillation by modeling the distribution of ensemble predictions in one model.

  • Motivation: Neural networks achieve strong supervised performance but tend to make over-confident predictions and historically provided limited uncertainty estimates.These limitations matter as neural networks are applied to safety-critical tasks.
  • Motivation: Ensembles improve accuracy, support uncertainty estimation, and can distinguish knowledge uncertainty from data uncertainty.Knowledge uncertainty concerns limited model knowledge, whereas data uncertainty arises from complexity, multimodality, and noise in the data.
  • Problem: Ensemble inference and training can cost many times more than using a single model, motivating ensemble distillation.Standard distillation retains the ensemble mean but loses information about prediction diversity.
  • Proposed task: Ensemble Distribution Distillation (EnD2) distills the distribution of ensemble predictions rather than only their mean into a single neural network.The task aims to retain both ensemble distributional information and improved classification performance.
  • Proposed task: The proposed EnD2 solution uses Prior Networks to model a conditional distribution over categorical output distributions with a Dirichlet distribution.This allows one model to emulate the complete ensemble.
  • Evaluation: On artificial and image datasets, EnD2 approaches ensemble classification performance and outperforms standard DNNs and Ensemble Distillation for misclassification and OOD detection.The image evaluations use CIFAR-10, CIFAR-100, and TinyImageNet.

2 ENSEMBLES

The paper views ensembles as samples from uncertainty over model parameters and uses their predictive spread to separate total uncertainty into knowledge and data components. Desired ensemble behavior includes consistent in-domain predictions and disagreement where model uncertainty is present.

  • Bayesian viewpoint: A Bayesian ensemble treats model parameters as random variables and represents their uncertainty through a posterior distribution p(θ|D).The paper also notes that ensembles can be constructed using non-Bayesian approaches such as random initialization or bootstrap methods.
  • Predictive distributions: Each ensemble member produces a categorical predictive distribution, and the ensemble predictive distribution is obtained by averaging over model uncertainty.The categorical distribution is parameterized by class probabilities π.
  • Desired behaviors: For in-domain inputs, the desired ensemble produces consistent predictions with little spread across its members.This behavior is presented as one of the target behaviors on the simplex of categorical probabilities.
  • Uncertainty decomposition: Total uncertainty is measured by predictive entropy, while disagreement among ensemble members captures knowledge uncertainty through mutual information.The decomposition separates expected data uncertainty from uncertainty associated with model disagreement.
  • Uncertainty decomposition: In class-overlap regions, ensemble members can each be uncertain without disagreeing, whereas disagreement becomes non-zero when their predictions differ.Thus predictive entropy can be high without a corresponding increase in the disagreement component.

3 ENSEMBLE DISTRIBUTION DISTILLATION

Ensemble Distribution Distillation distills an ensemble’s distribution over categorical predictions into one Prior Network, preserving diversity information alongside predictive performance. The model is trained on ensemble-derived distributions and supports decomposable uncertainty measures.

  • Motivation: Ensemble distillation traditionally matches the ensemble’s expected predictive distribution, collapsing its conditional distributions into a single point estimate and losing ensemble diversity.This loss prevents decomposing total uncertainty into knowledge and data uncertainty via mutual information.
  • Prior Network solution: Prior Networks parameterize a conditional distribution over categorical output distributions using one neural network, with a Dirichlet distribution specified by concentration parameters.The concentration parameters are produced from network outputs through an exponential function, and the Dirichlet is chosen for tractable uncertainty calculations.
  • Task definition: EnD2 instead models the distribution over ensemble predictions with a single neural network, aiming to preserve distributional information and ensemble classification performance.The task treats an ensemble as samples from an implicit distribution over output distributions.
  • Training: Training constructs a transfer dataset containing original inputs and ensemble-generated categorical distributions, then minimizes their negative log-likelihood.This negative-log-likelihood objective is equivalent to minimizing KL-divergence between the model and the empirical transfer distribution.
  • Prediction and uncertainty: The predictive distribution is the expected categorical distribution under the learned Dirichlet prior, while total uncertainty can be decomposed into data uncertainty and knowledge uncertainty.If EnD2 succeeds, these uncertainty measures should match those obtained from the original ensemble.

4 EXPERIMENTS ON ARTIFICIAL DATA

Artificial-data experiments use a noisy, nonlinearly separable three-spiral dataset to compare ensemble distillation approaches and their uncertainty estimates. EnD2 recovers much of ensemble classification performance and captures data uncertainty, but its out-of-distribution knowledge-uncertainty behavior remains imperfect without auxiliary data.

  • Dataset: The three-spiral dataset has three classes, increasing noise and arm separation, and overlapping regions requiring nonlinear decision boundaries.The experiment uses 1000 examples per class.
  • Classification performance: An ensemble of 10 models clearly outperforms the mean individual-model performance, while 100 models provide a smaller additional gain.Models are trained from different random initializations and compared with individual networks, EnD, and EnD2.
  • Classification performance: EnD recovers ensemble classification performance with minor degradation, while EnD2 recovers most of the gain with slightly greater degradation.The additional degradation is attributed to learning both the ensemble mean and its surrounding distribution, which may require greater model capacity.
  • Uncertainty estimation: EnD2 captures data uncertainty and decomposes total uncertainty into knowledge and data uncertainty, but produces low knowledge uncertainty in some regions far from training data.These regions appear as dark holes in figure 3f.
  • Auxiliary data: Auxiliary training data sampled in a ring far from the training data causes a minor performance drop while preserving performance above the average individual-DNN level.The auxiliary data are labeled using ensemble predictions.
  • Auxiliary data: Distilling an ensemble’s out-of-distribution behavior from in-domain observations alone is challenging and may require additional training data.Ensemble diversity is typically smaller on seen training data than on a held-out test set.

5 EXPERIMENTS ON IMAGE DATA

Image experiments evaluate whether EnD2 preserves ensemble classification, rejection, calibration, and OOD-detection behavior while modelling predictive diversity. Across CIFAR-10, CIFAR-100, and TinyImageNet, EnD2 generally retains ensemble classification and rejection performance, but Dirichlet mismatch affects calibration and uncertainty distributions.

  • Experimental setup: Ensembles of 100 models are distilled on CIFAR-10, CIFAR-100, and TinyImageNet, with additional transfer sets incorporating auxiliary data.The transfer set uses ensemble logits, and temperature annealing was essential for training EnD2 effectively.
  • Classification and rejection: EnD and EnD2 retain the ensemble’s improved classification and prediction-rejection performance relative to individual maximum-likelihood models across all datasets.EnD2 yields marginally better classification than EnD on CIFAR-100 and TinyImageNet and matches or outperforms EnD in PRR.
  • Calibration: EnD2 tends to yield better calibration performance than EnD, while both provide only similarly minor gains in test-set NLL over a single model.Auxiliary data provides limited EnD2 gains in ECE and NLL and sometimes degrades both metrics.
  • OOD detection: Ensemble uncertainty measures outperform those from a single neural network for OOD detection, while total uncertainty can sometimes outperform knowledge uncertainty on image datasets.The latter behavior is partly attributed to the image datasets containing relatively little data uncertainty.
  • Uncertainty distributions: On in-domain CIFAR-10 data, EnD2 emulates ensemble uncertainty reasonably well, but on OOD data it consistently produces higher uncertainty estimates.Auxiliary data increases total and data uncertainty; knowledge-uncertainty distributions become more ensemble-like but shift toward higher values.
  • Uncertainty distributions: The ensemble’s diversity differs from a Dirichlet distribution, causing EnD2 to overestimate empirical support and entropy while preserving or improving ID/OOD separation.This mismatch prevents EnD2 from fully retaining ensemble calibration quality; mixtures of Dirichlets or Logistic-normal distributions are suggested alternatives.

6 CONCLUSION

The conclusion presents EnD2 as a method for distilling ensemble classification performance and predictive diversity into one model. Experiments support uncertainty-type discrimination and strong image-task performance, while auxiliary data may be needed for suitable out-of-domain behavior.

  • Contribution: Ensemble Distribution Distillation distills an ensemble into one model while retaining improved classification performance and information about ensemble diversity.The method uses Prior Network models.
  • Findings: On artificial and image-classification tasks, distribution-distilled models retain ensemble classification performance and distinguish between data and knowledge uncertainty.The conclusion reports this behavior on the artificial dataset and image tasks.
  • Scope: EnD2 models can require auxiliary training data on which the ensemble is more diverse to learn appropriate out-of-domain behavior.This is stated as a condition for learning suitable OOD behavior.

APPENDIX A DATASETS, MODEL ARCHITECTURE AND TRAINING

The appendix describes the datasets, model architecture, optimization, and transfer-set construction used for the experiments. Models use a VGG16 variant in PyTorch, with different losses for individual, ensemble-distilled, and distribution-distilled training.

  • Architecture and optimization: All models are implemented in PyTorch with a VGG16 variant for image classification.The architecture is described as a variant of VGG16.
  • Architecture and optimization: DNNs, EnD models, and EnD2 models are trained with label negative log-likelihood, mean ensemble-prediction loss, and ensemble-output categorical-distribution negative log-likelihood, respectively.All models use the Adam optimizer.
  • Training configuration: Training configurations report initial learning rate, initial temperature, temperature annealing, and batch size.Table 6 defines η0, T0, and the annealing indicator.
  • Transfer-set construction: The transfer set uses ensemble predictions on unaugmented CIFAR training examples while distillation trains on corresponding augmented examples.This procedure applies to both EnD and EnD2.

A.1 NUMERICAL ISSUES WITH MAXIMUM LIKELIHOOD DIRICHLET TRAINING

The appendix states that likelihood-based EnD2 training can encounter numerical issues when ensemble probability terms cannot be computed. Central smoothing is applied to avoid this problem while preserving normalized ensemble predictions.

  • Training objective: Likelihood maximisation for Prior Network ensemble distribution distillation is equivalent to minimising the stated loss function.The loss is introduced as the objective corresponding to equation 8.
  • Numerical issue: Numerical precision can prevent computation of an ensemble member’s probability term, so central smoothing is applied to ensemble predictions.The smoothing addresses the computation of π(im)c terms.

A.2 TEMPERATURE ANNEALING SCHEDULE

The Ensemble Distillation baseline used a fixed temperature of 2.5, while annealing was excluded after reducing classification performance.

  • 2.5 was used as the fixed Ensemble Distillation temperature because it was recommended and performed best among {1, 2.5, 5, 10}.
  • Temperature annealing was not used for Ensemble Distillation because it resulted in worse classification performance.

APPENDIX B ASSESSING MISCLASSIFICATION DETECTION PERFORMANCE

This appendix defines evaluation procedures for misclassification detection using uncertainty-based rejection, ROC/PR analysis, and rejection metrics. It also identifies the performance measures reported for these evaluations.

  • Misclassification detection treats incorrect predictions as outliers and uses uncertainty thresholds to decide whether predictions should be rejected.
  • AUROC summarizes ROC trade-offs, but PR curves are preferable when negatives greatly outnumber positives because AUROC can over-estimate performance.
  • Misclassification detection results are reported using mean AUPR across C10, C100, and TIM models with uncertainty measures summarized in the accompanying table.
  • Rejected predictions can be replaced by oracle-provided predictions, with rejection curves showing the effect of ordering replacements by estimated uncertainty.
  • The prediction rejection ratio compares uncertainty-based rejection with random and oracle rejection, ranging from 1.0 for optimal rejection to 0.0 for random rejection.

APPENDIX C APPROPRIATENESS OF DIRICHLET DISTRIBUTION

Across CIFAR-100 and TinyImageNet, the appendix finds that EnD2 generally follows the CIFAR-10 uncertainty trends but produces higher uncertainty than the original ensemble. EnD models match in-domain uncertainty yet often underestimate OOD uncertainty.

  • EnD2 consistently yields higher uncertainties than the original ensemble on CIFAR-100 and TinyImageNet, likely reflecting limitations of the Dirichlet distribution.
  • All EnD models match ensemble uncertainty on in-domain data across the evaluated datasets, while EnD generally underestimates uncertainty on OOD data.
  • EnD+AUX is the exception on CIFAR-100 OOD data, where it matches the ensemble’s predictions well and nearly matches its OOD detection performance.

APPENDIX D ABLATION STUDIES

The ablation studies examine ensemble size and temperature choices for EnD2. They find that 20 models outperform 5 without conclusive gains beyond 20, while successful distribution distillation requires sufficiently high temperature and annealing.

  • Ensemble size: 20-model ensembles outperform 5-model ensembles, but the ablations show no conclusive EnD2 gains from using more than 20 models.
  • Ensemble size: Using fewer than 20 ensemble models may reduce the computation and memory costs of EnD2 when models parameterize the Dirichlet distribution.
  • Ensemble size: A more flexible distribution, such as a Mixture of Dirichlets, might allow larger ensembles to provide further gains.
  • Temperature: An initial temperature of at least 5 is necessary for successful distribution distillation, while increasing it from 10 to 20 gives no significant further performance increase.
  • Temperature: The ablations identify temperature annealing as important for EnD2 to work well.
Loading 1905.00076v3…