Original: H YDRA : P RESERVING E NSEMBLE D IVERSITY FOR M ODEL D ISTILLATION
Revised: H YDRA : P RESERVING E NSEMBLE D IVERSITY FOR M ODEL D ISTILLATION

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Ensembles of models have been empirically shown to improve predictive performance and to yield robust measures of uncertainty.
Revised: Ensembles of models have been empirically shown to improve predictive performance and to yield robust measures of uncertainty.

Original: However, they are expensive in computation and memory.
Revised: However, they are expensive in computation and memory.

Original: 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.
Revised: 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.

Original: Most existing distillation formulations summarize the ensemble by capturing its average predictions.
Revised: Most existing distillation formulations summarize the ensemble by capturing its average predictions.

Original: As a result, the diversity of the ensemble predictions, stemming from each individual member, is lost.
Revised: As a result, the diversity of the ensemble predictions, stemming from each individual member, is lost.

Original: Thus, the distilled model cannot provide a measure of uncertainty comparable to that of the original ensemble.
Revised: Thus, the distilled model cannot provide a measure of uncertainty comparable to that of the original ensemble.

Original: 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 .
Revised: 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 .

Original: The shared body network learns a joint feature representation that enables each head to capture the predictive behavior of each ensemble member.
Revised: The shared body network learns a joint feature representation that enables each head to capture the predictive behavior of each ensemble member.

Original: We demonstrate that with a slight increase in parameter count, Hydra improves distillation performance on classiﬁcation and regression settings while capturing the uncertainty behaviour of the original ensemble over both in-domain and out-of-distribution tasks.
Revised: We demonstrate that with a slight increase in parameter count, Hydra improves distillation performance on classiﬁcation and regression settings while capturing the uncertainty behaviour of the original ensemble over both in-domain and out-of-distribution tasks.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Deep neural networks have achieved impressive performance, however, they tend to make overconﬁdent predictions and poorly quantify uncertainty (Lakshminarayanan et al., 2017).
Revised: Deep neural networks have achieved impressive performance, however, they tend to make overconﬁdent predictions and poorly quantify uncertainty (Lakshminarayanan et al., 2017).

Original: It has been demonstrated that ensembles of models improve predictive performance and offer higher quality uncertainty quantiﬁcation (Dietterich, 2000; Lakshminarayanan et al., 2017; Ovadia et al., 2019).
Revised: It has been demonstrated that ensembles of models improve predictive performance and offer higher quality uncertainty quantiﬁcation (Dietterich, 2000; Lakshminarayanan et al., 2017; Ovadia et al., 2019).

Original: A fundamental limitation of ensembles is the cost of computation and memory at evaluation time.
Revised: A fundamental limitation of ensembles is the cost of computation and memory at evaluation time.

Original: A popular solution is to distill an ensemble of models into a single compact network by attempting to match the average predictions of the original ensemble.
Revised: A popular solution is to distill an ensemble of models into a single compact network by attempting to match the average predictions of the original ensemble.

Original: This idea goes back to the foundational work of Hinton et al.
Revised: This idea goes back to the foundational work of Hinton et al.

Original: , itself inspired by earlier ideas developed by Bucilu et al.
Revised: , itself inspired by earlier ideas developed by Bucilu et al.

Original: 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.
Revised: 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.

Original: 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 et al., 2017; Ovadia et al., 2019).
Revised: 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 et al., 2017; Ovadia et al., 2019).

Original: Similarly, by losing the diversity of the ensemble, this simple form of distillation makes it impossible to estimate measures of uncertainty such as model uncertainty (Depeweg et al., 2017; Malinin et al., 2019).
Revised: Similarly, by losing the diversity of the ensemble, this simple form of distillation makes it impossible to estimate measures of uncertainty such as model uncertainty (Depeweg et al., 2017; Malinin et al., 2019).

Original: Proper uncertainty quantiﬁcation is especially crucial for safety-related tasks and applications.
Revised: Proper uncertainty quantiﬁcation is especially crucial for safety-related tasks and applications.

Original: To overcome this limitation, Malinin et al.
Revised: To overcome this limitation, Malinin et al.

Original: (2019) proposed to model the entire distribution of an ensemble using a Dirichlet distribution parametrized by a neural network, referred to as a prior network (Malinin & Gales, 2018).
Revised: (2019) proposed to model the entire distribution of an ensemble using a Dirichlet distribution parametrized by a neural network, referred to as a prior network (Malinin & Gales, 2018).

Original: However, this imposes a strong parametric assumption on the distillation process.
Revised: However, this imposes a strong parametric assumption on the distillation process.

Original: Inspired by multi-headed architectures already widely applied in various applications (Szegedy et al., 2015; Sercu et al., 2016; Osband et al., 2016; Song & Chai, 2018), we propose a multiheaded model to distill ensembles.
Revised: Inspired by multi-headed architectures already widely applied in various applications (Szegedy et al., 2015; Sercu et al., 2016; Osband et al., 2016; Song & Chai, 2018), we propose a multiheaded model to distill ensembles.

Original: Our multi-headed approach—which we name Hydra —can be seen as an interpolation between the full ensemble of models and the knowledge distillation proposed by Hinton et al.
Revised: Our multi-headed approach—which we name Hydra —can be seen as an interpolation between the full ensemble of models and the knowledge distillation proposed by Hinton et al.

Original: Our distillation model is comprised of (1) a single body and (2) as many heads as there are members in the original ensemble.
Revised: Our distillation model is comprised of (1) a single body and (2) as many heads as there are members in the original ensemble.

Original: Each head is assigned to an ensemble member
Revised: Each head is assigned to an ensemble member

Original: [Figure]
Revised: [Figure]

Original: ˆ y pred yyyffff small ˆ y pred
Revised: ˆ y pred yyyffff small ˆ y pred

Original: D KL (b) Distillation of ensembles yyyfff
Revised: D KL (b) Distillation of ensembles yyyfff

Original: D (1) KL
Revised: D (1) KL

Original: ˆ yˆ yˆ yf body hhh
Revised: ˆ yˆ yˆ yf body hhh

Original: D (2) KL
Revised: D (2) KL

Original: D (3) KL
Revised: D (3) KL

Original: [Figure]
Revised: [Figure]

Original: otherwise and tries to mimic the individual predictions of this ensemble member, as illustrated in Figure 1.
Revised: otherwise and tries to mimic the individual predictions of this ensemble member, as illustrated in Figure 1.

Original: The heads share the same body network whose role is to provide a common feature representation.
Revised: The heads share the same body network whose role is to provide a common feature representation.

Original: The design of the body and the heads makes it possible to trade off the computational and memory efﬁciency against the ﬁdelity with which the diversity of the ensemble is retained.
Revised: The design of the body and the heads makes it possible to trade off the computational and memory efﬁciency against the ﬁdelity with which the diversity of the ensemble is retained.

Original: An illustration of the common knowledge distillation and ensemble distillation as well as Hydra is shown in Figureand a detailed methodology description is found in Section 2.
Revised: An illustration of the common knowledge distillation and ensemble distillation as well as Hydra is shown in Figureand a detailed methodology description is found in Section 2.

Original: 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 N − k , 0 ≤ k (cid:28) N , layers of the original ensemble members for the body and duplicating the layer for the heads.
Revised: 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 N − k , 0 ≤ k (cid:28) N , layers of the original ensemble members for the body and duplicating the layer for the heads.

Original: Summary of contributions.
Revised: Summary of contributions.

Original: Firstly, we present a multi-headed approach for the ensemble knowledge distillation.
Revised: Firstly, we present a multi-headed approach for the ensemble knowledge distillation.

Original: The shared component keeps the model computationally and memory efﬁcient while diversity is captured through the heads matching the individual ensemble members.
Revised: The shared component keeps the model computationally and memory efﬁcient while diversity is captured through the heads matching the individual ensemble members.

Original: Secondly, we show through experimental evaluation that Hydra outperforms existing distillation methods for both classiﬁcation and regression tasks with respect to predictive test performance.
Revised: Secondly, we show through experimental evaluation that Hydra outperforms existing distillation methods for both classiﬁcation and regression tasks with respect to predictive test performance.

Original: Lastly, we investigate Hydra’s behaviour in terms of in-domain and out-of-distribution data and demonstrate that
Revised: Lastly, we investigate Hydra’s behaviour in terms of in-domain and out-of-distribution data and demonstrate that

Original: Hydra comes closest to the ensemble behaviour in comparison to existing distillation methods.
Revised: Hydra comes closest to the ensemble behaviour in comparison to existing distillation methods.

Original: Novelty and signiﬁcance.
Revised: Novelty and signiﬁcance.

Original: Ensembles of models have successfully improved predictive performance and yielded robust measures of uncertainty.
Revised: Ensembles of models have successfully improved predictive performance and yielded robust measures of uncertainty.

Original: However, existing distillation methods do not retain the diversity of the ensemble (beyond its average predictive behavior), or need to make strong parametric assumptions that are not applicable in regression settings.
Revised: However, existing distillation methods do not retain the diversity of the ensemble (beyond its average predictive behavior), or need to make strong parametric assumptions that are not applicable in regression settings.

Original: To the best of our knowledge, our approach is the ﬁrst to employ a multi-headed architecture in the context of ensemble distillation.
Revised: To the best of our knowledge, our approach is the ﬁrst to employ a multi-headed architecture in the context of ensemble distillation.

Original: It is simple to implement, does not make the strong parametric assumptions, requires few modiﬁcations to the distilled ensemble model and works well in practice, thereby making it attractive to apply to a wide range of ensemble models and tasks.
Revised: It is simple to implement, does not make the strong parametric assumptions, requires few modiﬁcations to the distilled ensemble model and works well in practice, thereby making it attractive to apply to a wide range of ensemble models and tasks.

Original:  H YDRA : A M ULTI -H EADED A PPROACH
Revised:  H YDRA : A M ULTI -H EADED A PPROACH

Original: With a focus on ofﬂine distillation, our goal is to train a student network to match the predictive distribution of the teacher models, which is an ensemble of (deep) neural networks.
Revised: With a focus on ofﬂine distillation, our goal is to train a student network to match the predictive distribution of the teacher models, which is an ensemble of (deep) neural networks.

Original: Formally, given a dataset D = { ( x ( i ) , y ( i ) ) } Ni =1 , we consider an ensemble of M models’ parameters θ
Revised: Formally, given a dataset D = { ( x ( i ) , y ( i ) ) } Ni =1 , we consider an ensemble of M models’ parameters θ

Original: ens = { θ ens ,m } Mm =1 and prediction outputs { p ( y ( i )
Revised: ens = { θ ens ,m } Mm =1 and prediction outputs { p ( y ( i )

Original: | x ( i ) ; θ ens ,m ) } Mm =1 .
Revised: | x ( i ) ; θ ens ,m ) } Mm =1 .

Original: For simplicity, a single data instance pair will be referred to as ( x , y ) .
Revised: For simplicity, a single data instance pair will be referred to as ( x , y ) .

Original: In (Hinton et al., 2015; Balan et al., 2015) distilling an ensemble of models into a single neural network is achieved by minimizing the Kullback-Leibler (KL) divergence between the student’s predictive distribution { p ( y | x ; θ distill ) } and the expected predictive distribution of an ensemble:
Revised: In (Hinton et al., 2015; Balan et al., 2015) distilling an ensemble of models into a single neural network is achieved by minimizing the Kullback-Leibler (KL) divergence between the student’s predictive distribution { p ( y | x ; θ distill ) } and the expected predictive distribution of an ensemble:

Original: [Equation]
Revised: [Equation]

Original: x
Revised: x

Original: [Equation]
Revised: [Equation]

Original: Hydra builds upon the approach of knowledge distillation and extends it to a multi-headed student model.
Revised: Hydra builds upon the approach of knowledge distillation and extends it to a multi-headed student model.

Original: Hydra is deﬁned as a (deep) neural network with a single body and M heads.
Revised: Hydra is deﬁned as a (deep) neural network with a single body and M heads.

Original: For distillation,
Revised: For distillation,

Original: Hydra has as many heads as there are ensemble members.
Revised: Hydra has as many heads as there are ensemble members.

Original: The distillation model is parametrized by θ hydra = (cid:8) θ body , { θ head ,m } Mm =1 (cid:9) , i.e. the body, θ body , is shared among all heads { θ head ,m } Mm =1 .
Revised: The distillation model is parametrized by θ hydra = (cid:8) θ body , { θ head ,m } Mm =1 (cid:9) , i.e. the body, θ body , is shared among all heads { θ head ,m } Mm =1 .

Original: 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.
Revised: 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.

Original: In practice, we use N − k, 0 ≤ k (cid:28) N , layers of the original ensemble member architecture and the original ﬁnal layer(s) as head.
Revised: In practice, we use N − k, 0 ≤ k (cid:28) N , layers of the original ensemble member architecture and the original ﬁnal layer(s) as head.

Original: The objective is to minimize the average KL divergence between each head m and corresponding ensemble member m .
Revised: The objective is to minimize the average KL divergence between each head m and corresponding ensemble member m .

Original: We differentiate between two tasks, classiﬁcation and regression.
Revised: We differentiate between two tasks, classiﬁcation and regression.

Original: Classiﬁcation.
Revised: Classiﬁcation.

Original: For classiﬁcation tasks, the ensemble of models has access to D during training, with each x belonging to one of C classes, i.e., y ∈ { 1 , 2 , . . .
Revised: For classiﬁcation tasks, the ensemble of models has access to D during training, with each x belonging to one of C classes, i.e., y ∈ { 1 , 2 , . . .

Original: , C } .
Revised: , C } .

Original: Assuming z m =
Revised: Assuming z m =

Original: f θ head ,m ( f θ body ( x )) correspond to the logits, also termed unnormalised log probabilities, of head m , the categorical distribution over class labels for a sample x over a class c is computed as:
Revised: f θ head ,m ( f θ body ( x )) correspond to the logits, also termed unnormalised log probabilities, of head m , the categorical distribution over class labels for a sample x over a class c is computed as:

Original: [Equation]
Revised: [Equation]

Original: where T is a temperature re-scaling the logits.
Revised: where T is a temperature re-scaling the logits.

Original: As discussed in (Hinton et al., 2015;
Revised: As discussed in (Hinton et al., 2015;

Original: Malinin et al., 2019)
Revised: Malinin et al., 2019)

Original: the distribution of the teacher network is often “sharp”, which can limit the common support between the output distribution of the model and the target empirical distribution.
Revised: the distribution of the teacher network is often “sharp”, which can limit the common support between the output distribution of the model and the target empirical distribution.

Original: Mimimizing KLdivergence between distributions with limited non-zero common support is known to be particularly difﬁcult.
Revised: Mimimizing KLdivergence between distributions with limited non-zero common support is known to be particularly difﬁcult.

Original: To alleviate this issue, we follow the common practice (Hinton et al., 2015; Song & Chai, 2018; Lan et al., 2018) to use temperature to “heat up” both distributions and increase common support during training.
Revised: To alleviate this issue, we follow the common practice (Hinton et al., 2015; Song & Chai, 2018; Lan et al., 2018) to use temperature to “heat up” both distributions and increase common support during training.

Original: At evaluation, T is set to 1 .
Revised: At evaluation, T is set to 1 .

Original: The soft probability distributions at a 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 :
Revised: The soft probability distributions at a 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 :

Original: [Equation]
Revised: [Equation]

Original: Compared to the objective of knowledge distillation (1), we can observe that the average over the ensemble members is pulled out of the KL.
Revised: Compared to the objective of knowledge distillation (1), we can observe that the average over the ensemble members is pulled out of the KL.

Original: Ignoring the constant entropy terms, this objective is reduced to standard cross entropy loss:
Revised: Ignoring the constant entropy terms, this objective is reduced to standard cross entropy loss:

Original: [Equation]
Revised: [Equation]

Original: We scale our objective by T 2 as the gradient magnitudes produced by the soft targets are scaled by 1 /T 2 .
Revised: We scale our objective by T 2 as the gradient magnitudes produced by the soft targets are scaled by 1 /T 2 .

Original: By multiplying the loss term by a factor of T 2 we ensure that the relative contributions to additional regularization losses remain roughly unchanged (Song & Chai, 2018; Lan et al., 2018).
Revised: By multiplying the loss term by a factor of T 2 we ensure that the relative contributions to additional regularization losses remain roughly unchanged (Song & Chai, 2018; Lan et al., 2018).

Original: Regression.
Revised: Regression.

Original: We focus on heteroscedastic regression tasks where each ensemble member m outputs a mean µ m ( x ) and σ 2 m ( x ) given an input x.
Revised: We focus on heteroscedastic regression tasks where each ensemble member m outputs a mean µ m ( x ) and σ 2 m ( x ) given an input x.

Original: The output is modeled as p ( y | x, θ m ) =
Revised: The output is modeled as p ( y | x, θ m ) =

Original: N ( µ m ( x ) , σ 2 m ( x )) for a given head m and the ensemble of models are trained by minimizing the negative log-likelihood.
Revised: N ( µ m ( x ) , σ 2 m ( x )) for a given head m and the ensemble of models are trained by minimizing the negative log-likelihood.

Original: Traditional knowledge distillation matches a single Gaussian (“student”) outputting µ distill ( x ) and σ 2 distill ( x ) to a mixture of Gaussians (a “teacher” ensemble):
Revised: Traditional knowledge distillation matches a single Gaussian (“student”) outputting µ distill ( x ) and σ 2 distill ( x ) to a mixture of Gaussians (a “teacher” ensemble):

Original: [Equation]
Revised: [Equation]

Original: With Hydra, each head m outputs a mean µ hydra ,m ( x ) and variance σ 2 hydra ,m ( x ) and optimizes the
Revised: With Hydra, each head m outputs a mean µ hydra ,m ( x ) and variance σ 2 hydra ,m ( x ) and optimizes the

Original: KL divergence between each head output and corresponding ensemble member output:
Revised: KL divergence between each head output and corresponding ensemble member output:

Original: [Equation]
Revised: [Equation]

Original: where the ﬁnal line uses the fact that each KL term has an analytical solution.
Revised: where the ﬁnal line uses the fact that each KL term has an analytical solution.

Original: Training with multi-head growth.
Revised: Training with multi-head growth.

Original: Hydra is trained in two phases.
Revised: Hydra is trained in two phases.

Original: In the ﬁrst phase, Hydra mimics knowledge distillation in that it is trained until convergence with a single head—the “Hinton head”—to match the average predictions of the ensemble.
Revised: In the ﬁrst phase, Hydra mimics knowledge distillation in that it is trained until convergence with a single head—the “Hinton head”—to match the average predictions of the ensemble.

Original: Hydra is then extended by M − 1 heads, all of which initialized with the parameter values of the “Hinton head”.
Revised: Hydra is then extended by M − 1 heads, all of which initialized with the parameter values of the “Hinton head”.

Original: The resulting M heads are ﬁnally further trained to match the individual predictions of the M ensemble members (according to objective (3)).
Revised: The resulting M heads are ﬁnally further trained to match the individual predictions of the M ensemble members (according to objective (3)).

Original: In practice, we sometimes experienced difﬁculties 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.
Revised: In practice, we sometimes experienced difﬁculties 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.

Original:  E VALUATION
Revised:  E VALUATION

Original: In this section, we demonstrate that Hydra not only best matches the behavior of the teacher ensemble in terms of uncertainty quantiﬁcation but also improves the predictive performance compared to existing distillation approaches, over both classiﬁcation and regression tasks.
Revised: In this section, we demonstrate that Hydra not only best matches the behavior of the teacher ensemble in terms of uncertainty quantiﬁcation but also improves the predictive performance compared to existing distillation approaches, over both classiﬁcation and regression tasks.

Original: Datasets.
Revised: Datasets.

Original: For visualizing and explaining model uncertainty in Subsection 3.1 we used a spiral toy dataset.
Revised: For visualizing and explaining model uncertainty in Subsection 3.1 we used a spiral toy dataset.

Original: For classiﬁcation, we used two datasets: MNIST and CIFAR-10.
Revised: For classiﬁcation, we used two datasets: MNIST and CIFAR-10.

Original: For evaluating MNIST, we use its test set as well as increasingly shifted data (increasingly rotated or horizontally translated images) and Fashion-MNIST.
Revised: For evaluating MNIST, we use its test set as well as increasingly shifted data (increasingly rotated or horizontally translated images) and Fashion-MNIST.

Original: For CIFAR-10, we report performance of the test set, the cyclic translated test set and 80 different corrupted test sets as well as SVHN.
Revised: For CIFAR-10, we report performance of the test set, the cyclic translated test set and 80 different corrupted test sets as well as SVHN.

Original: The pre-processing for MNIST and CIFAR-10, as well as the generation schemes of the corrupted images, are taken from Ovadia et al.
Revised: The pre-processing for MNIST and CIFAR-10, as well as the generation schemes of the corrupted images, are taken from Ovadia et al.

Original: For regression, we conducted experiments on the standard regression datasets from the
Revised: For regression, we conducted experiments on the standard regression datasets from the

Original: UCI repository (Asuncion & Newman, 2007), following the protocol of Bui et al. (2016).
Revised: UCI repository (Asuncion & Newman, 2007), following the protocol of Bui et al. (2016).

Original: Ensemble setup.
Revised: Ensemble setup.

Original: 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.
Revised: 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.

Original: For MNIST and CIFAR, we each trained ensembles of 50 models.
Revised: For MNIST and CIFAR, we each trained ensembles of 50 models.

Original: For ensemble training we used different architectures for
Revised: For ensemble training we used different architectures for

Original: CIFAR-10 and MNIST.
Revised: CIFAR-10 and MNIST.

Original: We used an MLP architecture for MNIST and the ResNet-20 V1 architecture for CIFAR-10.
Revised: We used an MLP architecture for MNIST and the ResNet-20 V1 architecture for CIFAR-10.

Original: For additional details for MNIST and CIFAR-10 models, we refer to (Ovadia et al., 2019) and their open-source code 2 .
Revised: For additional details for MNIST and CIFAR-10 models, we refer to (Ovadia et al., 2019) and their open-source code 2 .

Original: For all regression tasks, the same model is optimized: an MLP with a 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.
Revised: For all regression tasks, the same model is optimized: an MLP with a 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.

Original: Distillation setup.
Revised: Distillation setup.

Original: We compare our work with two core distillation approaches, Knowledge Distillation (Hinton et al., 2015) and Prior Networks (Malinin et al., 2019; Malinin & Gales, 2018).
Revised: We compare our work with two core distillation approaches, Knowledge Distillation (Hinton et al., 2015) and Prior Networks (Malinin et al., 2019; Malinin & Gales, 2018).

Original: All baseline models have the same architecture as the ensemble for distillation.
Revised: All baseline models have the same architecture as the ensemble for distillation.

Original: For MNIST, Hydra uses the original ensemble member architecture and adds an MLP with two hidden layers of 100 units each as head.
Revised: For MNIST, Hydra uses the original ensemble member architecture and adds an MLP with two hidden layers of 100 units each as head.

Original: For CIFAR-10, the original Resnet20 V1 model without the last residual block was used as body.
Revised: For CIFAR-10, the original Resnet20 V1 model without the last residual block was used as body.

Original: The ﬁnal distillation model reported here has one residual block per head.
Revised: The ﬁnal distillation model reported here has one residual block per head.

Original: Evaluation metrics.
Revised: Evaluation metrics.

Original: We report all evaluation metrics on the test set based on the best validation loss from training.
Revised: We report all evaluation metrics on the test set based on the best validation loss from training.

Original: For both classiﬁcation and regression, we evaluate the negative log likelihood (NLL) as well as Brier score which depends on the predictive uncertainty and model uncertainty (MU).
Revised: For both classiﬁcation and regression, we evaluate the negative log likelihood (NLL) as well as Brier score which depends on the predictive uncertainty and model uncertainty (MU).

Original: NLL is a proper scoring rule and a popular metric for evaluating predictive uncertainty (Ovadia
Revised: NLL is a proper scoring rule and a popular metric for evaluating predictive uncertainty (Ovadia

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: et al., 2019).
Revised: et al., 2019).

Original: Model uncertainty, introduced by Depeweg et al.
Revised: Model uncertainty, introduced by Depeweg et al.

Original: ; Malinin et al.
Revised: ; Malinin et al.

Original: (2019), is a measure of the spread or disagreement of an ensemble based on mutual information.
Revised: (2019), is a measure of the spread or disagreement of an ensemble based on mutual information.

Original: A detailed description of MU and exemplary visualization can be found in Subsection 3.1.
Revised: A detailed description of MU and exemplary visualization can be found in Subsection 3.1.

Original: For classiﬁcation we additionally measure classiﬁcation accuracy and the Brier score.
Revised: For classiﬁcation we additionally measure classiﬁcation accuracy and the Brier score.

Original: NLL has the disadvantage of over-emphasizing tail probabilities (Quinonero-Candela et al., 2005).
Revised: NLL has the disadvantage of over-emphasizing tail probabilities (Quinonero-Candela et al., 2005).

Original: In contrast, the Brier score, a proper scoring rule that takes into account both calibration and accuracy, is not as strongly skewed as NLL (Gneiting & Raftery, 2007).
Revised: In contrast, the Brier score, a proper scoring rule that takes into account both calibration and accuracy, is not as strongly skewed as NLL (Gneiting & Raftery, 2007).

Original: For a given input-output pair ( x ∗ , y ∗ ) with y ∗ ∈ [1 , C ] the Brier score is deﬁned as BS = 1 K (cid:80)
Revised: For a given input-output pair ( x ∗ , y ∗ ) with y ∗ ∈ [1 , C ] the Brier score is deﬁned as BS = 1 K (cid:80)

Original: Cc =1 (cid:0) t ∗ − p
Revised: Cc =1 (cid:0) t ∗ − p

Original: ( y = c | x ∗ ) (cid:1) 2 where t ∗ c = 

Revised: ( y = c | x ∗ ) (cid:1) 2 where t ∗ c =

Original: if c = y ∗ , and 0 otherwise.
Revised: if c = y ∗ , and 0 otherwise.

Original: 3. U NCERTAINTY Q UANTIFICATION
Revised: 3. U NCERTAINTY Q UANTIFICATION

Original: We assess Hydra’s ability to distill uncertainty metrics from an ensemble on classiﬁcation tasks with
Revised: We assess Hydra’s ability to distill uncertainty metrics from an ensemble on classiﬁcation tasks with

Original: MNIST and CIFAR-10.
Revised: MNIST and CIFAR-10.

Original: One way to quantify uncertainty is through model uncertainty (Depeweg et al., 2017;
Revised: One way to quantify uncertainty is through model uncertainty (Depeweg et al., 2017;

Original: Malinin et al., 2019) which measures the spread or disagreement of an ensemble.
Revised: Malinin et al., 2019) which measures the spread or disagreement of an ensemble.

Original: Model uncertainty estimates the mutual information between the categorical label y and model parameters θ .
Revised: Model uncertainty estimates the mutual information between the categorical label y and model parameters θ .

Original: It can be expressed as the difference of the total uncertainty and the expected data uncertainty, where total uncertainty is the entropy of the expected predictive distribution and expected data uncertainty is the expected entropy of individual predictive distribution:
Revised: It can be expressed as the difference of the total uncertainty and the expected data uncertainty, where total uncertainty is the entropy of the expected predictive distribution and expected data uncertainty is the expected entropy of individual predictive distribution:

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: The total uncertainty will be high whenever the model is uncertain - both in regions of severe class overlap and out-of-domain.
Revised: The total uncertainty will be high whenever the model is uncertain - both in regions of severe class overlap and out-of-domain.

Original: However, for out-of-distribution data themodels’ understanding of data is low and thus, the estimates of expected data uncertainty are poor resulting in high model uncertainty.
Revised: However, for out-of-distribution data the estimates of expected data uncertainty are poor, resulting in high model uncertainty.

Original: Figure 2 visualizes the model uncertainty and its decomposition for the spiral toy dataset, with Subﬁgures 2a and 2b showing results for  ensemble and Hydra respectively.
Revised: Figure 2 visualizes the model uncertainty and its decomposition for the spiral toy dataset, with Subﬁgures 2a and 2b showing results for the ensemble and Hydra respectively.

Original: The results show that Hydra successfully model uncertainty and its decomposition, though with a slight decrease in scale.
Revised: The results show that Hydra successfully model uncertainty and its decomposition, though with a slight decrease in scale.

Original: We observe, as expected, a low model uncertainty where classes overlap due to both high total uncertainty and expected data uncertainty and high model uncertainty where at the border of in-domain and out-of-distribution data.
Revised: We observe, as expected, a low model uncertainty where classes overlap due to both high total uncertainty and expected data uncertainty and high model uncertainty where at the border of in-domain and out-of-distribution data.

Original: Further, we evaluate Hydra with respect to in-domain and out-of-distribution behavior for models trained on each MNIST and CIFAR10.
Revised: Further, we evaluate Hydra with respect to in-domain and out-of-distribution behavior for models trained on each MNIST and CIFAR10.

Original: We report the average absolute difference of model uncertainty as a measure how “close” the distillation method and ensemble are matched in Table 1.
Revised: We report the average absolute difference of model uncertainty as a measure how “close” the distillation method and ensemble are matched in Table 1.

Original: Hydra outperforms for MNIST test and shifted versions of MNIST, however, Hydra performs worse than
Revised: Hydra outperforms for MNIST test and shifted versions of MNIST, however, Hydra performs worse than

Original: Prior Network for several CIFAR10 datasets.
Revised: Prior Network for several CIFAR10 datasets.

Original: Looking closer at the performance of one of the worse results, CIFAR10 (cyclic translation), we plotted all evaluation metrics against the intensity of skew.
Revised: Looking closer at the performance of one of the worse results, CIFAR10 (cyclic translation), we plotted all evaluation metrics against the intensity of skew.

Original: 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.
Revised: 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.

Original: As visualized in Figure 3, Hydra matches the behaviour of the ensemble the best in terms of accuracy, Brier score and NLL.
Revised: As visualized in Figure 3, Hydra matches the behaviour of the ensemble the best in terms of accuracy, Brier score and NLL.

Original: With model uncertainty, Hydra with a larger head conﬁguration seems to improve on overall performance compared to Prior Network.
Revised: With model uncertainty, Hydra with a larger head conﬁguration seems to improve on overall performance compared to Prior Network.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 3. T EST P ERFORMANCE ON C OMMON C LASSIFICATION AND R EGRESSION B ENCHMARKS
Revised: 3. T EST P ERFORMANCE ON C OMMON C LASSIFICATION AND R EGRESSION B ENCHMARKS

Original: Classiﬁcation Performance on MNIST and CIFAR-10.
Revised: Classiﬁcation Performance on MNIST and CIFAR-10.

Original: We investigate Hydra on two real image datasets, MNIST and CIFAR-10, using the setup described in Section 3.
Revised: We investigate Hydra on two real image datasets, MNIST and CIFAR-10, using the setup described in Section 3.

Original: We report all metrics for MNIST in Table 2 and for CIFAR-10 in Table 3 as well as model capacity and efﬁciency in
Revised: We report all metrics for MNIST in Table 2 and for CIFAR-10 in Table 3 as well as model capacity and efﬁciency in

Original: Table 5.
Revised: Table 5.

Original: For MNIST we can see that all distillation methods match the accuracy of the target ensemble, 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).
Revised: For MNIST we can see that all distillation methods match the accuracy of the target ensemble, 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).

Original: For the more challening CIFAR-10 dataset all distillation methods approach but do not quite match the high accuracy of the target ensemble (0.9226 ensemble accuracy).
Revised: For the more challening CIFAR-10 dataset all distillation methods approach but do not quite match the high accuracy of the target ensemble (0.9226 ensemble accuracy).

Original: Among distillation methods Hydra has the smallest gap in terms of accuracy.
Revised: Among distillation methods Hydra has the smallest gap in terms of accuracy.

Original: All distillation methods retain a gap in NLL performance compared to the ensemble, but Hydra again has a signiﬁcantly smaller NLL (0.3179 NLL) compared to
Revised: All distillation methods retain a gap in NLL performance compared to the ensemble, but Hydra again has a signiﬁcantly smaller NLL (0.3179 NLL) compared to

Revised: Prior networks (0.4392 NLL).

Revised: In-distribution model uncertainty (MU) is comparable for both Prior

Original: [Table] [Table]
Revised: [Table] [Table]

Revised: possible

Original: Regression Performance on UCI Regression Datasets.
Revised: Regression Performance on UCI Regression Datasets.

Original: We trained both Knowledge Distillation (Hinton et al., 2015) and Hydra on standard regression UCI datasets shown in Table 4.
Revised: We trained both Knowledge Distillation (Hinton et al., 2015) and Hydra on standard regression UCI datasets shown in Table 4.

Original: Here
Revised: Here

Original: Prior Networks are not applicable because for the case of probabilistic regression we cannot take averages of distributions.
Revised: Prior Networks are not applicable because for the case of probabilistic regression we cannot take averages of distributions.

Original: For regression Hydra outperforms knowledge distillation in terms of predictive performance (NLL) because Hydra produces a more ﬂexible output in the form of a Gaussian mixture model with one Gaussian component per head, whereas Knowledge Distillation can produce only a single Gaussian component.
Revised: For regression Hydra outperforms knowledge distillation in terms of predictive performance (NLL) because Hydra produces a more ﬂexible output in the form of a Gaussian mixture model with one Gaussian component per head, whereas Knowledge Distillation can produce only a single Gaussian component.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: We now review related work in both distillation and multi-headed neural architectures.
Revised: We now review related work in both distillation and multi-headed neural architectures.

Original: Distillation.
Revised: Distillation.

Original: In (Hinton et al., 2015), a “student” network—i e. the outcome of the distillation process—is trained to match the average predictions of the “teacher” network(s).
Revised: In (Hinton et al., 2015), a “student” network—i.e. the outcome of the distillation process—is trained to match the average predictions of the “teacher” network(s).

Original: This methodology has been later successfully applied to the distillation of Bayesian ensembles (Balan et al., 2015) where ensemble members correspond to a Monte Carlo approximation to the posterior distribution.
Revised: This methodology has been later successfully applied to the distillation of Bayesian ensembles (Balan et al., 2015) where ensemble members correspond to a Monte Carlo approximation to the posterior distribution.

Original: A parallel line of research has focused on co-distillation , also sometimes referred to as online distillation to further reduce the overall training cost (Zhang et al., 2018;
Revised: A parallel line of research has focused on co-distillation , also sometimes referred to as online distillation to further reduce the overall training cost (Zhang et al., 2018;

Original: Anil et al., 2018; Lan et al., 2018).
Revised: Anil et al., 2018; Lan et al., 2018).

Original: In this setting, both the student and teacher networks are learned simultaneously and this form of mutual training acts as a regularization mechanism.
Revised: In this setting, both the student and teacher networks are learned simultaneously and this form of mutual training acts as a regularization mechanism.

Original: Also, distillation has been recently the topic of theoretical analysis to better explain its empirical success (Lopez-Paz et al., 2015; Phuong & Lampert, 2019).
Revised: Also, distillation has been recently the topic of theoretical analysis to better explain its empirical success (Lopez-Paz et al., 2015; Phuong & Lampert, 2019).

Original: [Table] [Table]
Revised: [Table] [Table]

Original: Closest to our approach is the work of Lan et al.
Revised: Closest to our approach is the work of Lan et al.

Original: Their method consists in training multiple student models whose combined predictions induce an ensemble teacher model.
Revised: Their method consists in training multiple student models whose combined predictions induce an ensemble teacher model.

Original: While we share conceptual similarities with their work, we depart from their formulations in several ways.
Revised: While we share conceptual similarities with their work, we depart from their formulations in several ways.

Original: First, we focus on the ofﬂine ensemble setting (Hinton et al., 2015) where we start from a pre-deﬁned ensemble whose training may be difﬁcult to replicate inside a co-distillation process.
Revised: First, we focus on the ofﬂine ensemble setting (Hinton et al., 2015) where we start from a pre-deﬁned ensemble whose training may be difﬁcult to replicate inside a co-distillation process.

Original: Second, our approach follows a different goal: we consider multiple branches, or heads in our terminology, to individually match the behavior of each teacher model .
Revised: Second, our approach follows a different goal: we consider multiple branches [Table] or heads in our terminology [Table]

Revised: to individually match the behavior of each teacher model .

Original: We do so to preserve the diversity of the teacher ensemble which is, for instance, essential in out-of-distribution tasks (Lakshminarayanan et al., 2017; Ovadia et al., 2019).
Revised: We do so to preserve the diversity of the teacher ensemble which is, for instance, essential in out-of-distribution tasks (Lakshminarayanan et al., 2017; Ovadia et al., 2019).

Original: Third, our methodology has a conceivably simpler design, as reﬂected by our single-component objective function—the average KL divergence between each head and corresponding teacher model—and the absence of a learned gating mechanism to linearly combine the logits of the student models (Lan et al., 2018).
Revised: Third, our methodology has a conceivably simpler design, as reﬂected by our single-component objective function—the average KL divergence between each head and corresponding teacher model—and the absence of a learned gating mechanism to linearly combine the logits of the student models (Lan et al., 2018).

Original: Multi-headed architectures.
Revised: Multi-headed architectures.

Original: This type of architecture can be motivated by several beneﬁts, e.g., a reduction in memory footprint due to the sharing of parameters, a speed-up of the training process, a regularization effect due to the introduction of auxiliary training objectives, or the transfer of information across different tasks.
Revised: This type of architecture can be motivated by several beneﬁts, e.g., a reduction in memory footprint due to the sharing of parameters, a speed-up of the training process, a regularization effect due to the introduction of auxiliary training objectives, or the transfer of information across different tasks.

Original: As discussed in detail by Song & Chai (2018), there exist multiple strategies to deﬁne a body of shared parameters together with heads , e.g. hierarchical sharing pattern.
Revised: As discussed in detail by Song & Chai (2018), there exist multiple strategies to deﬁne a body of shared parameters together with heads , e.g. hierarchical sharing pattern.

Original: In this work, we concentrate on simple multi-headed architectures, where the heads will correspond to either the last-layer of the original network or small extensions thereof.
Revised: In this work, we concentrate on simple multi-headed architectures, where the heads will correspond to either the last-layer of the original network or small extensions thereof.

Original: 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 we exploit the multiple heads to match and mirror the individual members of a pre-deﬁned teacher ensemble.
Revised: 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 we exploit the multiple heads to match and mirror the individual members of a pre-deﬁned teacher ensemble.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We presented Hydra , a simple and effective approach to distillation for ensemble models.
Revised: We presented Hydra , a simple and effective approach to distillation for ensemble models.

Original: Hydra preserves diversity in ensemble member predictions and we have demonstrated on standard models that capturing this information translates into improved performance and better uncertainty quantiﬁcation.
Revised: Hydra preserves diversity in ensemble member predictions and we have demonstrated on standard models that capturing this information translates into improved performance and better uncertainty quantiﬁcation.

Original: While Hydra improves on previous approaches we believe that we can further improvedistillation performance by leveraging techniques from ﬁelds studying sets of related learning such as meta-learning and domain adaptation.
Revised: While Hydra improves on previous approaches we believe that we can further improve