PaperlayerView exact revision on arXiv

Source-linked AI summary

Hydra: Preserving Ensemble Diversity for Model Distillation

Linh Tran, Bastiaan S. Veeling, Kevin Roth, Jakub Swiatkowski, Joshua V. Dillon, Jasper Snoek, Stephan Mandt, Tim Salimans, Sebastian Nowozin, Rodolphe Jenatton

arXiv:2001.04694v2cs.LGstat.ML

TL;DR

Ensemble distillation often matches average predictions, leaving member-level diversity unrepresented even though it matters for uncertainty estimation. Hydra uses a shared-body, multi-headed student to distill individual members, and the paper reports improved predictive performance and uncertainty behavior across classification, regression, and in- and out-of-distribution settings.

  • Problem

    Average-prediction distillation does not retain ensemble-member prediction diversity, which matters for uncertainty estimation, particularly in out-of-distribution settings.

    Source: §1 · pp. 1–2

  • Method

    Hydra distills each ensemble member into a separate head of a multi-headed student whose heads share a common body.

    Source: §2 · p. 3

  • Results

    Hydra outperforms existing distillation methods on classification and regression predictive performance and comes closest to ensemble behavior for in- and out-of-distribution data.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Hydra preserves ensemble-member diversity without strong parametric assumptions, including for regression settings.

    Source: §1 · p. 2

Abstract

from arXiv · show

Ensembles of models have been empirically shown to improve predictive performance and to yield robust measures of uncertainty. However, they are expensive in computation and memory. Therefore, recent research has focused on distilling ensembles into a single compact model, reducing the computational and memory burden of the ensemble while trying to preserve its predictive behavior. Most existing distillation formulations summarize the ensemble by capturing its average predictions. As a result, the diversity of the ensemble predictions, stemming from each member, is lost. Thus, the distilled model cannot provide a measure of uncertainty comparable to that of the original ensemble. To retain more faithfully the diversity of the ensemble, we propose a distillation method based on a single multi-headed neural network, which we refer to as Hydra. The shared body network learns a joint feature representation that enables each head to capture the predictive behavior of each ensemble member. We demonstrate that with a slight increase in parameter count, Hydra improves distillation performance on classification and regression settings while capturing the uncertainty behavior of the original ensemble over both in-domain and out-of-distribution tasks.

1 INTRODUCTION

Ensembles can improve predictive performance and uncertainty estimates but are costly at evaluation; average-prediction distillation does not retain member-level diversity. Hydra uses a shared body and separate heads to distill individual ensemble members, and the paper reports gains across classification, regression, and uncertainty evaluation.

1 INTRODUCTION

  • Average-prediction distillation loses ensemble members’ prediction diversity, limiting uncertainty estimation, especially for out-of-distribution tasks.The introduction notes this diversity is important when uncertainty outputs matter.Source: §1 · pp. 1–2
  • Hydra assigns each ensemble member a lightweight head that mimics its predictions, while a shared body supplies common features and amortizes computation.The body and head choices trade computational and memory efficiency against how faithfully ensemble diversity is retained.Source: §1 · p. 2, Figure 1 · p. 2
  • Hydra outperforms existing distillation methods on classification and regression predictive test performance and comes closest to ensemble behavior on in-domain and out-of-distribution data.The authors identify this as an experimental contribution.Source: §1 · p. 2
  • Hydra is presented as the first multi-headed architecture for ensemble distillation and avoids strong parametric assumptions, including assumptions that do not apply to regression settings.The authors describe it as simple to implement and requiring few modifications to the distilled ensemble model.Source: §1 · p. 2

2 HYDRA: A MULTI-HEADED APPROACH

Hydra extends knowledge distillation with a shared-body, multi-headed student whose heads match corresponding ensemble members. The method uses average KL-divergence objectives, with task-specific predictive distributions and a two-phase growth procedure described for training.

2 HYDRA: A MULTI-HEADED APPROACH

  • The offline distillation objective trains a student to match the predictive distribution of an ensemble of teacher networks.The paper frames this as matching the teachers’ predictive distribution on dataset instances.Source: §2 · p. 2
  • Hydra uses one shared body and one head per ensemble member, with each head trained to mimic its corresponding member’s predictions.The body provides a common feature representation; the heads are intended to be much lighter than the shared component.Source: §2 · p. 3
  • The shared body uses N − k layers of the original member architecture, with the original final layer or layers duplicated as heads.The setup assumes the heads are much lighter than the shared component so distillation remains meaningful.Source: §2 · p. 3
  • For classification, Hydra matches each head’s temperature-scaled soft probabilities to the corresponding teacher using average KL divergence.The temperature increases common support during training; evaluation uses T = 1.Source: §2 · p. 3
  • For heteroscedastic regression, each Hydra head outputs a Gaussian mean and variance and minimizes KL divergence against its corresponding ensemble member.Traditional distillation instead matches a single Gaussian to the ensemble’s Gaussian mixture; each Hydra KL term has an analytical solution.Source: §2 · pp. 3–4
  • Hydra training first converges a single Hinton head on average ensemble predictions, then expands and trains M heads against individual members.The authors report that this initialization can help convergence and typically leads to quicker convergence than other initializations.Source: §2 · p. 4

3 EVALUATION

The evaluation tests predictive performance and uncertainty matching across classification and regression, including in-domain and shifted data. Hydra generally matches ensemble behavior more closely than existing distillation methods, with dataset- and metric-specific differences.

3 EVALUATION

  • The evaluation uses MNIST, CIFAR-10, and UCI regression datasets, assessing test performance and uncertainty; classification tests include shifted and corrupted data and SVHN.Models are selected by validation loss, and comparisons include Knowledge Distillation and Prior Networks.Source: §3 · pp. 4–5

3.1 UNCERTAINTY QUANTIFICATION

  • On the spiral dataset, Hydra models the ensemble’s uncertainty decomposition with a slight decrease in scale, showing low model uncertainty at class overlap and high model uncertainty near the in-/out-of-distribution boundary.Figure 2 presents the ensemble and Hydra alongside the original dataset, with colors identifying classes.Source: §3.1 · pp. 5–6, Figure 2 · p. 5
  • Hydra better matches ensemble uncertainty on MNIST test and shifted data, but trails Prior Networks on several CIFAR-10 datasets; on cyclic translation, it matches accuracy, Brier score, and NLL best as shift increases.The CIFAR-10 comparison also examines a larger Hydra head; Prior Networks’ larger uncertainty scale contributes to better overall model-uncertainty scores despite Hydra better mimicking the uncertainty decomposition.Source: §3.1 · p. 6, Figure 3 · p. 6

3.2 TEST PERFORMANCE ON COMMON CLASSIFICATION AND REGRESSION BENCHMARKS

  • On MNIST, Hydra nearly matches ensemble NLL (0.0465 versus 0.0439) and Brier score (-0.9776 versus -0.9780); on CIFAR-10, it has the smallest accuracy gap and lower NLL (0.3179) than Prior Networks (0.4392).All distillation methods approach but do not match the CIFAR-10 ensemble accuracy of 0.9226.Source: §3.2 · pp. 6–7, Table 2 · p. 7, Table 3 · p. 7
  • On UCI regression datasets, Hydra outperforms knowledge distillation in NLL, producing a Gaussian mixture with one component per head rather than a single Gaussian.Prior Networks are not applicable to probabilistic regression because distributions cannot be averaged.Source: §3.2 · p. 7

4 RELATED WORK

The related work contrasts average-prediction distillation and online co-distillation with Hydra’s offline approach, which uses separate heads to match members of a predefined ensemble. It also places Hydra among multi-headed architectures while distinguishing its purpose from prior uses.

Distillation

  • Traditional distillation trains a student to match the teacher’s average predictions, while co-distillation trains student and teacher networks simultaneously as a regularization mechanism.For probabilistic regression, averaging Gaussian densities generally does not produce a Gaussian.Source: §4 · pp. 7–8

Closest to our approach

  • Hydra uses multiple heads to match individual members of a predefined ensemble offline, unlike Lan et al. (2018), whose student models jointly induce an ensemble teacher.Hydra’s objective averages the KL divergence between each head and its corresponding teacher, without a learned gating mechanism.Source: §4 · p. 8

Multi-headed architectures

  • Multi-headed architectures can share parameters for lower memory use, and Hydra uses heads to mirror individual members of a predefined teacher ensemble.The work focuses on simple architectures whose heads are the original network’s last layer or small extensions; it primarily distills predictive probabilities rather than feature representations.Source: §4 · p. 8

Distillation

  • Prior work also studies distillation theoretically and proposes feature-representation strategies, while this section identifies Hydra’s focus as predictive-probability supervision.The cited theoretical analyses are Lopez-Paz et al. (2015) and Phuong & Lampert (2019); feature-based approaches include Ba & Caruana (2014) and Romero et al. (2014).Source: §4 · p. 8

5 CONCLUSION

Hydra preserves diversity in ensemble predictions, and the paper reports improved performance and uncertainty quantification on standard models. The authors identify meta-learning and domain adaptation as possible directions for further improving distillation.

5 CONCLUSION

  • Hydra preserves ensemble-prediction diversity and improves performance and uncertainty quantification on standard models.The authors suggest meta-learning and domain adaptation as possible ways to further improve distillation performance.Source: §5 · p. 9

Related papers

Loading 2001.04694v2…