Hydra: Preserving Ensemble Diversity for Model Distillation
Authors: Linh Tran
Bastiaan S. Veeling
Kevin Roth
Jakub Świątkowski
Joshua V. Dillon
Jasper Snoek
Stephan Mandt
Tim Salimans
Sebastian Nowozin
Rodolphe Jenatton
Venue:  (2020)
Keywords: model distillation
ensemble models

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 individual 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 behaviour of the original ensemble over both in-domain and out-of-distribution tasks.

Under review as a conference paper at ICLR 2020
HYDRA:
PRESERVING ENSEMBLE DIVERSITY FOR MODEL DISTILLATION
Anonymous authors Paper under double-blind review
ABSTRACT
Ensembles of models have been empirically shown to improve predictive perfor- mance 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 ex- isting distillation formulations summarize the ensemble by capturing its average predictions. As result, the diversity of the ensemble predictions, stemming from each individual member, is lost. Thus, the distilled model cannot provide a mea- sure 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 cap- ture 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 behaviour of the original ensemble over both in-domain and disttiiburttitoionnnnn tasks.
INTRODUCTION
Deep neural networks have achieved impressive performance, however, they tend to make over-
confident predictions and poorly quantify uncertainty (Lakshminarayanan et al. (017)). It has been demonstrated that ensembles of models improve predictive performance and offer higher quality uncertainty quantification (Dittterich] 2000, Lakshminarayanan et al. 2017] Ovadia et al., 2019). fundamental limitation of ensembles is the cost of computation and memory at evaluation time. popular solution is to distill ensemble of models into a single compact network by attempting to match the average predictions the original ensemble. This idea goes back to the foundational work of Hinton et al. (2015). itself inspired by earlier ideas developed by Buciiu et al. (2006). While this process has led to simple and well-performing algorithms. it fails to take into account the intrinsic diversity of the predictions of the ensemble, as represented by the individual predictions of each of its members. In particular, this diversity is all the more important in tasks that hinge on the uncertainty output of the ensemble, e.g., in out-of-distribution scenarios (Lakshminarayanan al. 2017 Ovadia ct al. 2019). Similarly, by losing the diversity of the ensemble, this simple form of distillation makes impossible to estimate measures of uncertainty such as model uncertainty (Depeweg al. 2017 Malinin et al. 2019). Proper uncertainty quantification is especially crucial for ralate- tasks and applications. To overcome this limitation, Malinin al. (2019) proposed to model the entire distribution of an ensemble using a Dirichlet distribution parametrized by neural network, referred to as a prior network (Malinin & Gales, 2018). However, this imposes a strong parametric assumption the distillation process.
Inspired by multi-headed architectures already widely applied in various applications (Szegedy et al. 2015 Sercu al.) 2016, Osband et al.) 2016 [ong & Chail 2018). we propose multi- headed model to distill ensembles. Our multi-headed approach which we name Hydra- Mlydr--can can be seen as interpolation between the full ensemble of models and the knowledge distillation proposed by |Hinton et al. (2015). Our distillation model is comprised of (1) single body and (2) as many heads as there members in the original ensemble. Each head is assigned to ensemble member

---

Under review as a conference paper at ICLR 2020
/
f:
20
3/2
Large
Sprod
f
gi
Spred
f
RT.
DKL
f
20
KL
PKL
S1
Ssmall
Apred
Femall
Spred
Abody
(a) Knowledge distillation
(b) Distillation of ensembles
(c) Hydra distillation
Existing distillation methods compared to Hydra. Knowledge distillation, (Hinton et al, 2015), trains a distillation network to imitate the prediction of a larger network. Applying knowledge distillation to ensemble models Hinton et al. (2013) train a network to imitate the average ensemble prediction. Hydra instead learns distill the individual predictions of each ensemble member into separate light-weight head models while amortizing the computation through a shared heavy-weight body network. This retains the diversity of ensemble member predictions which is otherwise lost in knowledge distillation.
Figure 1:
and tries to mimic the individual predictions of this ensemble member, as illustrated in Figure I The heads share the same body network whose role is provide common feature representation. The design of the body and the heads makes it possible to trade off the computational and memory efficiency against the fidelity with which the diversity of the ensemble is retained. An illustration of the common knowledge distillation and ensemble distillation as well as Hydra is shown in Figure and a detailed methodology description is found in Section 2 While the choices of the body and head architectures may appear like complex new hyperparameters we introduce, we will see in the experiments that we get good results by simply taking the k. ensemble members for the body and duplicating the layer for the heads.
< N, layers of the original
Summary of contributions. Firstly, we present multi-headed approach for the ensemble knowl- edge distillation. The shared component keeps the model computationally and memory efficient while diversity is captured through the heads matching the individual ensemble members. Secondly, show through experimental evaluation that Hydra outperforms existing distillation methods for both classification and regression tasks with respect to predictive test performance. Lastly, we in- vestigate Hydra's 's behaviour in terms of in-domain and distr--distributionn data and demonstrate that Hydra comes closest to the ensemble behaviour in comparison to existing distillation methods.
Novelty and significance. Ensembles of models have successfully improved predictive performance and yielded robust measures of uncertainty. However, existing distillation methods do not retain the diversity of the ensemble (beyond its average predictive behavior), or need make strong paramet- ric assumptions that are not applicable in regression settings. To the best of our knowledge, our approach is the first employ multi-headed architecture the context of ensemble distillation. It is simple to implement, does not make the strong parametric assumptions, requires few modifi- cations to the distilled ensemble model and works well in practice, thereby making it attractive to apply to wide range of ensemble models and tasks.
HYDRA: A MULTI-HEADED APPROACH
With focus offline distillation, goal is to train student network to match the predic- tive distribution of the teacher models, which an ensemble of (deep) neural networks. For- mally, given a dataset = ()) y( J we consider an ensemble M models' parameters Oens = (0 = and prediction outputs (x()())(x()) (000) ) For simplicity, a single data instance pair will be referred as (x. y).
In Hinton et al., 2015 Balan et al., (2155) distilling an ensemble of models into a single neural network is achieved by minimizing the Keibbec (KL) divergence between the student's predictive distribution {p(y]x, [x; (asaan)) and the expected predictive distribution of an ensemble:
M
C(9ems, Busau)) E. KL
Σ
(p(y)x; mm p(y]x; Quisan))
M
2

---

Under review as a conference paper at ICLR 2020
Hydra builds upon the approach of knowledge distillation and extends it to multi-headed student model. Hydra is defined as a (deep) neural network with a single body and M heads. For distillation, Hydra has many heads as there ensemble members. The distillation model is parametrized by LM Ohydra = (Oboyy, (0head, m==1] ). i.e. the body, Oodyy is shared among all heads (Ohead.m))n)))
In terms of number of parameters, we assume the heads to be much lighter than the shared part, so that the distillation is still meaningful. In practice, we use N k, 0 Σ < N, layers of the original ensemble member architecture and the original final layer(s) as head. The objective is to minimize the average KL divergence between each head m and corresponding ensemble member m. We differentiate between two tasks, classification and regression.
Classification. For classification tasks, the ensemble of models has access to D during training, with each x belonging to one of C classes, i.e., E [1,2.... ,,...,,,,, C}. Assuming Zm f (Sewa (x) correspond to the logits, also termed unnormalised log probabílities, of head m. the categorical distribution over class labels for a sample over class c is computed as:
exp(Zm,e./
p[c|x; Omydramm) = p(clzm) ES=1 exp(zmm.j
(2)
where 7' is temperature re-scaling the logits. As discussed in (Hinton et al. 2015 Malinin et al, 2019) the distribution the teacher network is often "sharp", which can limit the common support between the output distribution of the model and the target empirical distribution. Mimimizing KL- divergence between distributions with limited non-zers zero common support is known to be particularly difficult. To alleviate this issue, follow the common practice (Hinton et al. 2013 [ong & Chai
2018 Lan et al. 2018) to use temperature to "heat up" both distributions and increase common sup- port during training. At evaluation, T is set 1. The soft probability distributions temperature of T are used to match the teacher ensemble of models by minimizing the average KL divergence between each head m and ensemble model m:
M Σ
[(x, y; Ohydras Onns)
(L(p(y|x; Oens,m) || p(y) x; Obody. Ohead.m )) m==
M
Compared to the objective of knowledge distillation I we can observe that the average the ensemble members is pulled out of the KL. Ignoring the constant entropy terms, this objective is reduced to standard cross entropy loss:
M Σ ply|x; Oens,m) log x; Obody: Opead.m).
L(x, y; Orend, 0ens)
M
We scale our objective by 2 as the gradient magnitudes produced by the soft targets are scaled by 1/T2, By multiplying the loss term by factor of T22 we ensure that the relative contributions to additional regularization losses remain roughly unchanged (Song & Chai 2018 Lan et al. 2018).
Regression. We focus on heteroscedastic regression tasks where each ensemble member m out- puts a mean un and of (æ) given an input x!_ The output modeled as p(yla, f,.) (u (I), o (x)) for given head and the ensemble of models are trained by minimizing the negative log-likelihood. Traditional knowledge distillation matches single Gaussian ("student") outputting (ri) and oisan (1) to mixture of Gaussians (a "teacher" ensemble):
C(x, y; Oens, Oaisin) = KI.
1 M
Σ
N (x), 2.000 I (v), ddssiin (2:)
With Hydra, each head m outputs a mean Jhydramm (z) and variance of nydm, (x) and optimizes the KL divergence between each head output and corresponding ensemble member output:
In our concrete implementation, neural network outputs the mean (m(()) (x) and log standard deviation log (x)) which thereafter exponentiate.
3

---

Under review as a conference paper at ICLR 2020
1 M
[(x, y; Bens, Ohydra) =
Σ
KL
(1) o2(00 N (Jhydraam (z), hyutra m (x))
M
m=1
1 M
(omitting constants)
N(y (mm o2) (x)) log ((y; Maistill oissnn dy
X
M
m==
o (ar) + (Jmm (2) Maistill (1:))
1 log
+
M
20 disun (4)
where the final line uses the fact that each KL term has analytical solution.
Training with multi-head growth. Hydra is trained in two phases. In the first phase, Hydra mimics know ledge distillation in that is trained until convergence with single head the "Hinton head"-to match the average predictions of the ensemble. Hydra is then extended by M - heads, all of which initialized with the parameter values of the "Hinton head". The resulting M heads are finally further trained to match the individual predictions of the M ensemble members (according to objective 3 In practice, we sometimes experienced difficulties for Hydra to converge in absence of this initialization scheme, and for the cases where different initialization worked, this two-phase training scheme typically led to overall quicker convergence.
3 EVALUATION
In this section, we demonstrate that Hydra not only best matches the behavior of the teacher ensem- ble in terms of uncertainty quantification but also improves the predictive performance compared to existing distillation approaches, over both classification and regression tasks.
Datasets. For visualizing and explaining model uncertainty in Subsection 3.1 we used spiral toy dataset. For classification, we used two datasets: MNIST and CIFAR-IO. 10. For evaluating MNIST, we use test set as well increasingly shifted data (increasingly rotated horizontally translated images) and Fashion-MNIST. For CIFAR-10, 10, we report performance of the test set, the cyclic trans- lated test set and 80 different corrupted test sets as well as SVHN. The pre-processing for MNIST and CIFAR-10, ct al. 2019). For regression, we conducted experiments on the standard regression datasets from the UCI repository (Asuncion & Newman 2007), following the protocol of Bui et al. 2016)).
well as the generation schemes of the corrupted images, are taken from Ovadia
Ensemble setup. For the toy dataset, we trained 10 ensemble models, each of which consists of a multi-layer perceptron (MLP) with two hidden layers of 100 units each. For MNIST and CIFAR, we each trained ensembles of 50 models. For ensemble training we used different architectures for CIFAR-10 and MNIST. We used an MLP architecture for MNIST and the ResNet-20 V1 architecture for CIFAR-IO. For additional details for MNIST and CIFAR-I0 models, we refer to (Ovadia et al., 2019 and their open-source code For all regression tasks, the same model is optimized: an MLP with single hidden layer of 50 units with softplus activation for each dataset, except for case of the larger protein structure dataset, prot where 100 units were used (following Bui et al. (2016)).
Distillation setup. We compare our work with two core distillation approaches, Knowledge Distil- lation (Hinton et al., 2015) and Prior Networks (Malinin et al., 2019 Malinin Gales) 2018). All baseline models have the same architecture as the ensemble for distillation. For MNIST, Hydra uses the original ensemble member architecture and adds an MLP with two hidden layers of 100 units each as head. For CIFAR-IO, the original Resnet20 VI model without the last residual block was used as body. The final distillation model reported here has one residual block per head.
Evaluation metries. We report all evaluation metrics on the test set based on the best validation loss from training. For both classification and regression, we evaluate the negative log likelihood (NLL) well as Brier score which depends the predictive uncertainty and model uncertainty (MU). NLL is a proper scoring rule and a popular metric for evaluating predictive uncertainty (Ovadia
hhtpps://g thu.ihhubbb.om
resea htaaa//ccecooooaaarrarar///////r
benchmark,
4

---

Under review as a conference paper at ICLR 2020
Moddel uncersiiity [MD
Total uncertaity (Tu)
Expecied daia und wete (HDUU)
-02
da de de 5
&
220 -as
(a) Ensemble
Model uncertalnty (M)
Total uncertainty (TU)
Expected data uncertainty (EDU)
(c) Spiral dataset
de da
5
33
(b) Hydra
Figure 2: Model uncertainty, total uncertainty and expected data uncertainty applied in-domain and -usto--distriuution data for both (a) ensemble and (b). The original dataset is visualized (c), where each color corresponds to single class.
Dataset
Knowledge distillation
Prior Networks
Hydra
(Hinton et al.
2015
Malinin al., 20199
MNIST (test)
N/A
2.60 10 23.31- 10-5
3.11 10 10
MNIST (shifted)
N/A
0.1938 1.0076
0.1552
1,00
(roST(rotated)
N/A
0.2943 1. 0100
0.0528
:: 2.0116
Fashion-MNIST
N/A
0.1636 1: 1065
0.2717 :: 0866
CIFAR-10 (test)
N/A
0.0877 1.1371
0.0983
0738
CIFAR-I0 (cyclic translation)
N/A
0.1330 +.002s
0.1646 0.00
CITAR-I0 (skewed)
N/A
0.0877 1.14
0.0938
SVHN
N/A
0.3916
1208
0.3611 #.1191
Table 1: Average absolute difference of model uncertainty and std. of several runs (N 3) between
distillation models and ensemble of models. Knowledge distillation (Hinton et al. 2015) was not reported, as it does not offer uncertainty quantification. Lower better.
Model uncertainty, introduced by Depeweg al. (2017); Malinin et (2019), is measure of the spread or disagreement of an ensemble based on mutual information. A detailed description of MU and exemplary visualization can be found in Subsection 3.I For classification we additionally measure classification accuracy and the Brier score. NLL has the disadvantage of over-emphasizing tail probabilities (Quinonero-Candela et al. 2005). In contrast, the Brier score, proper scoring rule that takes into account both calibration and accuracy, is not as strongly skewed as NLL (Gneiting & Raftery 2007). For a given input-output pair (x* yY) with y" E [1, C/j the Brier Σ (t' (pyp(y- c)=*))² where 1-- = 1 if e y", and 0 otherwise.
al. 2019)
score is defined as BS-
3.1 UNCERTAINTY QUANTIFICATION
We assess Hydra's 's ability to distill uncertainty metries from an ensemble on classification tasks with MNIST and CIFAR- 10. One way to quantify uncertainty is through model uncertainty (Depeweg et al. 2017 Malinin et al., 2019) which measures the spread or disagreement of an ensemble. Model uncertainty estimates the mutual information between the categorical label and model parameters can be expressed as the difference the total uncertainty and the expected data uncertainty. where total uncertainty is the entropy of the expected predictive distribution and expected data un- certainty is the expected entropy of individual predictive distribution:
MUly, 00'0''')'' Dj H ((Epooro)(p(lk> 0)) Ep(e HH(ply|x*, O))
(6)
Model uncertainty
Total uncertainty
Expected data uncertainty
5

---

Under review as a conference paper at ICLR 2020
095
0.2
Ersemibe
0.00
Hydia (hood-last res, Modi
050
075-
Fror Nemeodo
(evairted
035
Byyhe (hrett herk
Networts
Vel
-0.-
Hysia (hoeiinted ymu. toca)
Hydka (hodd-last
Malia throseeted teway Pror Netteroo
Ilycha lheddtled brekkt
Netterrts
-0.5
04-
DG
Wi
VMM
Taet
Insembbee
Rydta thred Yes Morks
(hend-iast Nied0
Hyrta (bredabbe Norkt
Nelworks
64
62-
Vili
1952
Intarity of 5kOW
(a) Accuracy (ACC), Brier score (BS). negative log- likelihood (NLL)
(b) Model uncertainty, total uncertainty, expected data uncertainty
Figure 3: Model evaluation metrics plotted against intensity of distributional shift for CIFAR- 10. The plots contain all distillation models as well as the original ensemble of models. We observe that Hydra consistently more accurately captures the uncertainty of the ensemble.
The total uncertainty will be high whenever the model is uncertain- both in regions of severe class overlap and out-of-domain. However, for uncertainty are poor, resulting in high model uncertainty. Figure 2 visualizes the model uncertainty and its decomposition for the spiral toy dataset, with Subfigures 2a and 2b showing results for the ensemble and Hydra respectively. The results show that Hydra successfully model uncertainty and its decomposition, though with a slight decrease in scale. We observe, expected, low model uncertainty where classes overlap due to both high total uncertainty and expected data uncertainty and high model uncertainty where the border of in-domain and
data the estimates of expected data
data.
Further, we evaluate Hydra with respect to in-domain and duitto-distribution behavior for models trained on each MNIST and CIFARIO. We report the average absolute difference of model uncer- tainty as a measure how "close" the distillation method and ensemble are matched in Table Hydra outperforms for MNIST test and shifted versions of MNIST, however, Hydra performs worse than Prior Network for several CIFAR1O 10 datasets. Looking closer at the performance of one of the worse results, CIFARIO 10 (cyclic translation), we plotted all evaluation metrics against the intensity of skew. In order to rule out, that the worse performance is related to model capacity, we added the results from Hydra trained with a larger head of three residual blocks. As visualized in Figure Hydra matches the behaviour of the ensemble the best in terms of accuracy, Brier score and NLL. With model uncertainty, Hydra with a larger head configuration seems to improve on overall performance compared to Prior Network. Inspecting the decomposition of model uncertainty, total uncertainty and expected data uncertainty, it seems like Hydra is more capable to "mimic" the behaviour the of ensemble. However, the uncertainty scales of Prior Networks is larger than the ones of both the ensemble and Hydra, leading to an overall better model uncertainty.
3.2 TEST PERFORMANCE ON COMMON CLASSIFICATION AND REGRESSION BENCHMARKS
real im- age datasets, MNIST and CIFAR-10, using the setup described in Section 3 We report all metrics for MNIST in Table 2 and for CIFAR-10 10 in Table 3 as well as model capacity and efficiency in Table For MNIST we can see that all distillation methods match the accuracy of the target ensem- ble, but Hydra outperforms both knowledge distillation and prior networks in terms of capturing the ensemble uncertainty, almost matching the ensemble predictive NLL (Hydra NLL 0.0465 matching ensemble NLL 0.0439) and Brier score (Hydra -0.9776 versus ensemble -0.9780). For the more challening CIFAR-I0 10 dataset all distillation methods approach but do not quite match the high accu- racy of the target ensemble (0.9226 ensemble accuracy). Among distillation methods Hydra has the smallest gap in terms of accuracy. All distillation methods retain gap in NLL performance com- pared to the ensemble, but Hydra again has a significantly smaller NLL (0.3179 NLL) compared to Prior networks (0.4392 NLI.). In-distribution model uncertainty (MU) is comparable for both Prior
Classification Performance on MNIST and CIFAR-10. We investigate Hydra

---

Under review as a conference paper at ICLR 2020
Model
ACC
NLL
BS +
MU
Ensemble (M 50)
0.9851
0.0439
-0.9780
9.97 10-6
Prior Networks
0.9842
0.0521
1 1.0007
-0.9285
0.1158
(Malinin 2
Knowledge distillation
0.9843 + 0013
0.0497
0000
-0.9764
+ +.0005
N/A
[Hinton et al. (2015)
Hydra
0.9857
0.0004
0.0465
0.0004
-0.9776
#4.60 10-6
2.28 10
-1.90-10m
(head- = [100, 100, 10)
Table 2: Average test performance and std. over several runs (n 3) for different baselines and Hydra for MNIST. For all models, classification accuracy (ACC), negative log-likelihood (NLL), Brier score (BS) and model uncertainty (MU) are reported. For #Params and FLOPs, we also report in brackets the percentage with respect to the ensemble. Bold numbers represent best performance with respect to the specific evaluation metric (columns) across distillation models (rows) except for MU, where we report the model uncertainty closest to the ensemble one.
Model
ACC
NLL +
BS
MU
Ensemble (M = 50)
0,9226
0.2392
-0.9033
0.1055
Prior Networks
0.8731 4.0085
0.4392
4.0020
-0.8231
#.0015
0.0280 +.0009
(Malinin et al. (019)
Knowledge distillation
0.8933 4.0020
0.3598
4.0023
-0.8373
4. 4.0004
N/A
(Hinton et al. (015
Hydra
0.8992 1.0011
0.3179
1: 1.0000
-0.8468 1: 1.0030
0.0074 1.0005
(head = last res. block)
Table 3: Average test performance and std. over several runs (n = 3) for different baselines and
Hydra for CIFAR-10. For all models, classification accuracy (ACC), negative log-likelihood (NLL), Brier score (BS) and Model Uuncertainty (MU) are reported. Bold numbers represent best per- formance with respect to the specific evaluation metric (columns) across distillation models (rows) except for MU, where we report the model uncertainty closest to the ensemble one.
Networks (0.0280) and Hydra (0.0074) but quite a bit smaller compared to target ensemble MU of 0.1055, meaning is possible improve uncertainty quantification in all distillation methods tested.
Regression Performance on UCI Regression Datasets. We trained both Knowledge Distilla- tion (Hinton al. 2015) and Hydra standard regression UCI datasets shown in Table 4 Here Prior Networks are not applicable because for the case of probabilistic regression we cannot take averages of distributions For regression Hydra outperforms knowledge distillation in terms of pre- dictive performance (NLL) because Hydra produces more flexible output in the form of a Gaussian mixture model with one Gaussian component per head, whereas Knowledge Distillation can produce only single Gaussian component.
4 RELATED WORK
We now review related work in both distillation and multi-headed neural architectures.
Distillation. In (Hinton et al., 2015). "student" network-i.ee the outcome of the distillation
process-is trained to match the average predictions of the "teacher" network(s). This methodology has been later successfully applied to the distillation of Bayesian ensembles (Balan al. 2015) where ensemble members correspond to Monte Carlo approximation to the posterior distribution parallel line of research has focused on co-distillation, also sometimes referred to as online dis- tillation to further reduce the overall training cost (Zhang et al., 2018]] Anil et al. 2018 Lan et al. 2018). In this setting, both the student and teacher networks are learned simultaneously and this form of mutual training acts as a regularization mechanism. Also, distillation has been recently the topic of theoretical analysis to better explain empirical success (Lopez-Paz al. 2015 Phuong & Lamperi 2019).
Formally, the average of two Gaussian densities is no longer Gaussian in general.
'Note that, in this paper, we focus primarily on instances of distillation where the supervision occurs at the level of the predictive probabilities, although other strategies, e.g., based on the feature representations in the networks (Ha & Caruanal 2014 Romero al. 2014), have heen proposed.
7

---

Under review as a conference paper at ICLR 2020
Dataset
Ensemble
Prior Network
Knowledge distillation (Hinton et al. (2015))
Hydra (head = [10.] 1)
Malinin et al.
2019
bost 2.3780
NA
2.3893 #.0089
2.3805 4.0084
concr
3.0585
N/A
3.1231
#.0054
3.0982 2.0065
ener
1.2756
N/A
I.5236
#.0277
1.4402 #
kin8
-1.2977
N/A
-1.2343 1.0014
-1.2555 1.0034
nava
-7.5983
N/A
-6.4340
1.2710
-7.1987 11.
powe 2.8861
N/A
2.8940
#.0007
2.8921 +.0007
prot 2.8272
N/A
2.8970
+.0007
2.8829 +.0064
wine 0.9111
N/A
0.9112
1:0012
0.91 13 4:0010
yach -6.1640
N/A
0.3837
#.0174
0.3489 #0.0136
Table 4: UCT regression benchmark (Dua & Taniskidou 2017). Average test negative log-likelihood (NLL) and std. over several runs 3) of the 9 different datasets considered. As Prior Net-
works (Malinin et al. 2019) cannot be applied to regression tasks, we denote this with "N/A".
MNIST
CIFAR-10
Model
#Params J
FLOPs
#Params J
FLOPs J
Ensemble (M = 50)
9,960,500
9.0 10
13,722,100
4.09 10th
Prior Networks
199,210 [2%)
1.8 10 [2%]
274,442 [2%]
8.18 10a [2%]
(Malinin et al. 2019)
Knowledge distillation
199,210 [2%)
1.8 10" [2%]
274,442 [2%)
8.18 10 [2%]
(Hinton et al. (2015)
Hydra
1,757,700 700 [17.6%]
2.5 10 [27.7%]
(head 1100, 100, 10)
Hydra
3,950,324 [28.7%]
5.45 10" [3.3%%
(head=last res.block)
Table 5: Model capacity and efficiency: We report number of model parameters (#Params) and number of floating point operations (FLOPs) for the ensemble and all distillation models. We also report in brackets the percentage with respect to the ensemble.
Closest to our approach is the work Lan al. (2018). Their method consists in training multiple student models whose combined predictions induce an ensemble teacher model. While we share conceptual similarities with their work, we depart from their formulations in several ways. First, we focus on the offline ensemble setting (Hinton al. 2015) where we start from a pre-defined ensemble whose training may be difficult to replicate inside a co-distillation process. Second, our approach follows different goal: we consider multiple branches, or heads in terminology, to individually match the behavior of each teacher model. We do so to preserve the diversity of the teacher ensemble which is, for instance, essential in out-of-distribution tasks (Lakshminarayanan S et al., 2017 Ovadia et al.) (019). Third, our methodology has a conceivably simpler design, as reflected by our single-component objective function average KL divergence between each head and corresponding teacher model-and and the absence of a learned gating mechanism to linearly combine the logits of student models (Lan al., 2018).
Multi-headed architectures. This type of architecture can be motivated by several benefits, e.g., reduction memory footprint due the sharing of parameters, speed-up of the training process, regularization effect due to the introduction of auxiliary training objectives, or the transfer of information across different tasks. As discussed in detail by [Song & Chai (2018), there exist multiple strategies to define body of shared parameters together with heads, c.g. hicrarchical sharing pattern. In this work, we concentrate on simple multi-headed architectures, where the heads will correspond to either the last-layer of the original network small extensions thereof. While multi-headed architectures were used for online distillation in (Lan et al.) 2018), we reiterate that our goal is different from this previous work in that exploit the multiple heads match and mirror the individual members of a pre-defined teacher ensemble.
CONCLUSION
We presented Hydra, a simple and effective approach to distillation for ensemble models. Hydra preserves diversity in ensemble member predictions and have demonstrated on standard models that capturing this information translates into improved performance and better uncertainty quan- tification. While Hydra improves previous approaches believe that further improve
8

---

Under review as a conference paper at ICLR 2020
distillation performance by leveraging techniques from fields studying sets of related learning such as meta-learning and domain adaptation.
REFERENCES
Rohan Anil, Gabriel Pereyra, Alexandre Passos, Robert Ormandi, George Dahl, and Geoffrey E Hinton. Large scale distributed neural network training through online distillation. preprint arXiv: 1804.0333, 2018.
Arthur Asuncion and David Newman. Uci machine learning repository. 2007.
Jimmy Ba and Rich Caruana. Do deep nets really need to be deep? In Advances in neural informa- tion processing systems, 2014.
Anoop Korattikara Balan, Vivek Rathod. Kevin P Murphy, and Max Welling. Bayesian dark knowl- edge. In Advances in Neural Information Processing Systems, pp. 3438-3446, 2015.
Cristian Bucilu, Rich Caruana, and Alexandru Niculescu-Mizze Model compression. In Proceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining, pp- 535-541. ACM, 2006.
Thang Bui, Daniel Hernández-Lobato Jose Hernandez-Lobato, Yingzhen Li, and Richard Turner. Deep gaussian processes for regression using approximate expectation propagation. Interna- tional Conference on Machine Learning, pp. 1472-1481, 2016.
Stefan Depeweg, José Miguel Hernández-Lobato, Finale Doshi- Velez, and Steffen Udluft. De- composition of uncertainty for active learning and reliable reinforcement learning in stochastic systems, stat, 1050:11, 2017.
Thomas G Dietterich. Ensemble methods in machine learning, In International workshop on multi- ple classifier systems, pp. 1-15. Springer, 2000.
Dheeru Dua and E Karra Taniskidou. UCI machine learning repository. University of Califormia, Irvine, School of Information and Computer Sciences, 2017.
Tilmann Gneiting and Adrian E Raftery. Strictly proper scoring rules, prediction, and estimation. Journal of the American Statistical Association, 102(477):399-37,, 2007.
2007.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. preprint arXiv:1503.025:: 2015.
Balaji Lakshminarayanan. Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Advances in Neural Information Processing Systems, pp. 1402-6413,20 2017.
Xu Lan, Xiatian Zhu, and Shaogang Gong. Knowledge distillation by on-the-fly native ensemble. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pp. 7528-7538, 2018.
David Lopez-Paz, Léon Bottou, Bernhard Schölkopf, and Vladimir Vapnik. Unifying distillation and privileged information. preprint arXiv:1511.0360 2015.
Andrey Malinin and Mark Gales. Predictive uncertainty estimation via prior networks. In Advances in Neural Information Processing Systems, pp. 7047-7058, 2018.
Andrey Malinin, Bruno Mlodozeniec, and Mark Gales. Ensemble distribution distillation. preprint arXiv:1905.00076, 2019.
Ian Osband, Charles Blundell, Alexander Pritzel, and Benjamin Van Roy. Deep exploration via bootstrapped dqn. In Advances in neural information processing systems, pp. 4026-4034, 2016.
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, D Sculley, Sebastian Nowozin, Joshua V Dil- lon, Balaji Lakshminarayanan, and Jasper Snoek. Can you trust your model's uncertainty? eval- uating predictive uncertainty under dataset shift. preprint arXiv: arXiv: 2019.

---

Under review as a conference paper at ICLR 2020
Mary Phuong and Christoph Lampert. Towards understanding knowledge distillation. In Interna- tional Conference on Machine Learning, pp. 5142-5151, 2019.
Joaquin Quinonero-Candela, Carl Edward Rasmussen, Fabian Sinz, Olivier Bousquet, and Bern- hard Schölkopf. Evaluating predictive uncertainty challenge. In Machine Learning Challenges Workshop, pp. 1-27. Springer, 2005.
Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio. Fitmets: Hints for thin deep nets. arXiv preprint arXiv:1412.6550, 2014.
Tom Sercu, Christian Puhrsch, Brian Kingsbury, and Yann LeCun. Very deep multilingual convo- lutional neural networks for Ivesr. In 2016 IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pp.
IEEE, 2016.
Guocong Song and Wei Chai. Collaborative learning for deep neural networks. In Advances in Neural Information Processing Systems, pp. 1832-1841, 2018.
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Du- mitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1-9, 2015.
Ying Zhang, Tao Xiang, Timothy M Hospedales, and Huchuan Lu. Deep mutual learning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 4320- 4328, 2018
10


Reviews:
[{"review_id": "ryexo70hKr", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Overview:\nThis work introduces a new method for ensemble distillation. The problem of making better ensemble distillation methods seems relevant as ensembles are still one of the best ways to estimate uncertainty in practice (although see concerns below). The method itself is a simple extension of earlier “prior networks”: the original method suggested to fit a single network to mimick a distribution produce by given ensemble, and here authors suggest to use multi-head (one head per individual ensemble member) in order to better capture the ensemble diversity. \nAuthors report results on multiple relatively standard benchmarks (MNIST, CIFAR, etc), and seem to outperform the baseline by a small margin. The choice of baselines is reasonable.\n\nWriting:\nThe paper is well-written, illustrations are good.\n\nDecision:\nThe method itself is very easy to implement, and does seem to outperform the baseline (prior networks). However, I am a bit concerned that the method itself seems like a trivial extension of the prior work, and does not really provide much addition insight. In addition, the results are reported on a set of small-scale benchmarks and seem incremental: it can be OK, but it would be really great to see a somewhat more realistic application.\nThus, I am on the fence with this one, but generally positive about this work, thus “weak accept” rating.\n\nQuestions / concerns:\n* I honestly do not see the point on having an additional column in tables if all the values are N/A. \n* The names in Table 4 are mixed up.\n* Arguably, a lot of applications that would actually rely on uncertainty estimation might require online training of some sort. This means that in those scenarios one does not actually have access to a pre-trained ensemble. I understand that this might not be the main focus of this work, but it seems like a major limitation of “distillation” approaches in general, which should / could be addressed in some way?\n\n<update>\nThanks for a detailed answer. I am not very convinced by the argument about online-vs-offline training, but I do not see a reason to decrease my rating. I can see this work useful in practice.\n</update>\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "HJgwwEGEFH", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary & Pros\n- This paper proposes a simple yet effective distillation scheme from an ensemble of independent models to a multi-head architecture for preserving the diversity of the ensemble.\n- The proposed scheme provides the same advantages of the ensemble in terms of uncertainty estimation and predictive performance, but it is computationally efficient compared to the ensemble.\n- This paper experimentally shows that multi-head architecture performs well on MNIST and CIFAR-10 in terms of accuracy and uncertainty.\n\nConcerns #1: Novelty of the proposed method\n- The multi-head architectures have been widely used in various settings, especially multi-task learning. As the authors mentioned, it also used for online distillation [1]. Although its goal is different from this paper, just applying such multi-head architectures seems to be incremental.\n\nConcerns #2: Insufficient experiments\n- To evaluate OOD detection quality, ID/OOD datasets should be stated and various metrics (e.g., AUROC) should be measured like other literature, e.g., Table 2 in [2]. Such OOD detection quality is important to evaluate the quality of uncertainty estimation.\n- This paper provides experiments on only small-sized 10-class datasets, MNIST and CIFAR10. To verify the effectiveness of the proposed distillation method, other large-sized datasets should be tested, e.g., CIFAR-100, ImageNet.\n- There is no ablation study on the effect of the number of size of heads in Hydra. To achieve similar performance to the ensemble, how many heads are required?\n\nConcerns #3: Week efficiency\n- As reported in Table 5, in the case of CIFAR10, Hydra has 14x more parameters and 6x more FLOPs. Despite such a large number of parameters, the performance gain seems to be incremental.\n- A comparison with an ensemble with M=14 models should be tested because this ensemble has the same number of parameters compared to Hydra with M=50 heads. I think it might achieve good performance on the evaluation metrics.\n\nConcerns #4: Incremental improvements\n- Accuracy gain is too marginal even Hydra uses 14x more parameters.\n- ONE [1] might be a stronger baseline because ONE achieves 94% accuracy on CIFAR-10 using ResNet32 with only 2~3 heads while Hydra achieves only 90% even it uses 50 heads. Moreover, since ONE has multiple heads, uncertainty estimation is also available. So it should be compared with the proposed method.\n- In OOD detection tasks, Hydra underperforms Prior Networks on 5 of 8 datasets (note that PN (2.60) is better than Hydra (3.11) in the case of MNIST (test)). To overcome this gap, the proposed method requires more parameters.\n\n[1] Zhu, Xiatian, and Shaogang Gong. \"Knowledge Distillation by On-the-Fly Native Ensemble.\" Advances in Neural Information Processing Systems. 2018.\n[2] Andrey Malinin and Mark Gales. \"Predictive Uncertainty Estimation via Prior Networks.\" Advances in Neural Information Processing Systems. 2018.\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "ByxD8uBpFr", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes to distill the predictions of an ensemble with a multi-headed network, with as many heads as members in the original ensemble. Distillation proceeds by minimizing the KL divergence between the predictions of each ensemble member with the corresponding head in the student network. Experiments illustrate that the multi-headed architecture approximates the ensemble marginally better than approaches that use a network with a single head.\n\nThe paper presents a straightforward idea and fairly unsurprising results —  a multi-headed architecture with each head matching an ensemble member more faithfully represents the original ensemble. This improved fidelity, however, comes at the cost of increased computation and storage requirements (which scale linearly with the size of the ensemble). It is unclear that the marginal improvements demonstrated justify the increased cost and how this approach would scale to larger ensembles. The paper would have been more interesting if the authors had managed to demonstrated significant improvements over competitors on not toy (MNIST / CIFAR) problems. Unfortunately, this is not the case and the fact that similar ideas (Lan et al.) have been proposed in the past (which the authors, to their credit, cite) leads me to recommend a rejection.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work introduces a simple and effective method for ensemble distillation. The method is a simple extension of earlier “prior networks”: it differs in which, instead of fitting a single network to mimic a distribution produced by the ensemble, this work suggests to use multi-head (one head per individual ensemble member) in order to better capture the ensemble diversity. This paper experimentally shows that multi-head architecture performs well on MNIST and CIFAR-10 (they added CIFAR-100 in the revised version) in terms of accuracy and uncertainty.

While the method is effective and the experiments on CIFAR-100 (a harder task) improved the paper, the reviewers (myself included) pointed out in the discussion phase that the limited novelty remains a major weakness. The proposed method seems like a trivial extension of the prior work, and does not provide much additional insight. To remedy this shortcoming, I suggest the authors provide extensive experimental supports including various datasets and ablation studies.  

Another concern mentioned in the discussion is the fact that these small improvements are in spite of the fact that the proposed method ends up using many more parameters than the baselines. Including and comparing different model sizes in a full fledged experimental evaluation would better convey the trade-offs of the proposed approach.

Decision: Reject