Hydra: Preserving Ensemble Diversity for Model Distillation

TL;DR

Hydra is a novel distillation method that uses a multi-headed neural network to preserve ensemble diversity and uncertainty estimates in a compact model, improving over traditional averaging approaches.

Contribution

The paper introduces Hydra, a multi-headed neural network for ensemble distillation that maintains diversity and uncertainty information.

Findings

Hydra outperforms traditional distillation methods in classification and regression.

Hydra better captures uncertainty in both in-domain and out-of-distribution tasks.

Slight increase in parameters yields significant performance gains.

Abstract

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.