Hydra: Preserving Ensemble Diversity for Model Distillation

1 INTRODUCTION . Deep neural networks have achieved impressive performance , however , they tend to make overconfident predictions and poorly quantify uncertainty ( Lakshminarayanan et al. , 2017 ) . It has been demonstrated that ensembles of models improve predictive performance and offer higher quality uncertainty quantification ( Dietterich , 2000 ; Lakshminarayanan et al. , 2017 ; Ovadia et al. , 2019 ) . A fundamental limitation of ensembles is the cost of computation and memory at evaluation time . 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 . This idea goes back to the foundational work of Hinton et al . ( 2015 ) , itself inspired by earlier ideas developed by Bucilu 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 et al. , 2017 ; Ovadia et al. , 2019 ) . 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 ) . Proper uncertainty quantification is especially crucial for safety-related tasks and applications . To overcome this limitation , Malinin et al . ( 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 ) . However , this imposes a strong parametric assumption on the distillation process . 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 . 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 . ( 2015 ) . Our distillation model is comprised of ( 1 ) a single body and ( 2 ) as many heads as there are members in the original ensemble . Each head is assigned to an ensemble member distillation to ensemble models Hinton et al . ( 2015 ) train a network to imitate the average ensemble prediction . Hydra instead learns to 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 . and tries to mimic the individual predictions of this ensemble member , as illustrated in Figure 1 . The heads share the same body network whose role is to provide a 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 1 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 N −k , 0 ≤ k N , layers of the original ensemble members for the body and duplicating the layer for the heads . Summary of contributions . Firstly , we present a multi-headed approach for the ensemble knowledge distillation . The shared component keeps the model computationally and memory efficient while diversity is captured through the heads matching the individual ensemble members . Secondly , we show through experimental evaluation that Hydra outperforms existing distillation methods for both classification and regression tasks with respect to predictive test performance . Lastly , we investigate Hydra ’ s behaviour in terms of in-domain and out-of-distribution 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 to make strong parametric assumptions that are not applicable in regression settings . To the best of our knowledge , our approach is the first to employ a multi-headed architecture in the context of ensemble distillation . It is simple to implement , does not make the strong parametric assumptions , requires few modifications 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 . 2 HYDRA : A MULTI-HEADED APPROACH . With a focus on offline 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 . Formally , given a dataset D = { ( x ( i ) , y ( i ) ) } Ni=1 , we consider an ensemble of M models ’ parameters θens = { θens , m } Mm=1 and prediction outputs { p ( y ( i ) |x ( i ) ; θens , m ) } Mm=1 . For simplicity , a single data instance pair will be referred to as ( x , y ) . 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 : L ( θens , θdistill ) = Ex [ KL ( 1 M M∑ m=1 ( p ( y|x ; θens , m ) ) ‖ p ( y|x ; θdistill ) ) ] ( 1 ) Hydra builds upon the approach of knowledge distillation and extends it to a multi-headed student model . Hydra is defined as a ( deep ) neural network with a single body andM heads . For distillation , Hydra has as many heads as there are ensemble members . The distillation model is parametrized by θhydra = { θbody , { θhead , m } Mm=1 } , i.e . the body , θbody , is shared among all heads { θhead , m } Mm=1 . 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 ≤ k 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 ofC classes , i.e. , y ∈ { 1 , 2 , . . . , C } . Assuming zm = 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 : p ( c|x ; θhydra , m ) = p ( c|zm ) = exp ( zm , c/T ) ∑C j=1 exp ( zm , j/T ) , ( 2 ) where T is a temperature re-scaling the logits . As discussed in ( Hinton et al. , 2015 ; Malinin et al. , 2019 ) 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 . Mimimizing KLdivergence between distributions with limited non-zero common support is known to be particularly difficult . 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 . At evaluation , T is set to 1 . 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 : L ( x , y ; θhydra , θens ) = 1 M M∑ m=1 KL ( p ( y|x ; θens , m ) ‖ p ( y|x ; θbody , θhead , m ) ) . ( 3 ) Compared to the objective of knowledge distillation ( 1 ) , we can observe that the average over the ensemble members is pulled out of the KL . Ignoring the constant entropy terms , this objective is reduced to standard cross entropy loss : L ( x , y ; θhead , θens ) = − T 2 M M∑ m=1 p ( y|x ; θens , m ) log p ( y|x ; θbody , θhead , m ) . ( 4 ) We scale our objective by T 2 as the gradient magnitudes produced by the soft targets are scaled by 1/T 2 . 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 ) . Regression . We focus on heteroscedastic regression tasks where each ensemble member m outputs a mean µm ( x ) and σ2m ( x ) given an input x . 1 The output is modeled as p ( y|x , θm ) = N ( µm ( x ) , σ2m ( x ) ) for a given head m and the ensemble of models are trained by minimizing the negative log-likelihood . Traditional knowledge distillation matches a single Gaussian ( “ student ” ) outputting µdistill ( x ) and σ2distill ( x ) to a mixture of Gaussians ( a “ teacher ” ensemble ) : L ( x , y ; θens , θdistill ) = KL ( 1 M M∑ m=1 N ( µm ( x ) , σ2 , ( i ) m ( x ) ) ‖ N ( µdistill ( x ) , σ2distill ( x ) ) ) ( 5 ) With Hydra , each head m outputs a mean µhydra , m ( x ) and variance σ2hydra , m ( x ) and optimizes the KL divergence between each head output and corresponding ensemble member output : 1In our concrete implementation , our neural network outputs the mean µm ( x ) and log standard deviation log ( σ2m ( x ) ) which we thereafter exponentiate . L ( x , y ; θens , θhydra ) = 1 M M∑ m=1 KL ( N ( µm ( x ) , σ2 , ( i ) m ( x ) ) ‖ N ( µhydra , m ( x ) , σ2hydra , m ( x ) ) ) , ( omitting constants ) ∝ − 1 M M∑ m=1 ∫ N ( y ; µm ( x ) , σ2 , ( i ) m ( x ) ) logN ( y ; µdistill , σ2distill ) dy = − 1 M M∑ m=1 σ2m ( x ) + ( µm ( x ) − µdistill ( x ) ) 2 2σ2distill ( x ) + 1 2 log ( 2πσ2distill ( x ) ) where the final line uses the fact that each KL term has an analytical solution . Training with multi-head growth . Hydra is trained in two phases . In the first 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 . Hydra is then extended by M − 1 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 theM 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 .

Summary: This 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.