Diversity Matters When Learning From Ensembles

1 Introduction . Deep Ensemble ( DE ) [ Lakshminarayanan et al. , 2017 ] , a simple method to ensemble the same model trained multiple times with different random initializations , is considered to be a competitive method for various tasks involving deep neural networks both in terms of prediction accuracy and uncertainty estimation . Several works have tried to reveal the secret behind DE ’ s effectiveness . As stated by Duvenaud et al . [ 2016 ] and Wilson and Izmailov [ 2020 ] , DE can be considered as an approximate Bayesian Model Average ( BMA ) procedure . Fort et al . [ 2019 ] studied the loss landscape of DEs and showed that the effectiveness comes from the diverse modes reached by ensemble members , making it well suited for approximating BMA . It is quite frustrating that most sophisticated approximate Bayesian inference algorithms , especially the ones based on variational approximations , are not as effective as DEs in terms of exploring various modes in parameter spaces . Despite being simple to train , DE incurs significant overhead in inference time and memory usage . It is therefore natural to develop a way to reduce such costs . An example of such a method is Knowledge Distillation ( KD ) [ Hinton et al. , 2015 ] , which transfers knowledge from a large teacher network to a relatively smaller student network . The student network is trained with a loss that encourages it to copy the outputs of the teacher networks evaluated at the training data . With KD , there have been several works that learn a single student network by distilling from DE teacher networks . A naïve approach would be to directly distill ensembled outputs of DE teachers to the single student network . A better way proven to be more effective is to set up a student network having multiple subnetworks ( multiple heads [ Tran et al. , 2020 ] or rank-one factors [ Mariet et al. , 2020 ] ) and distill the outputs of each ensemble member to each subnetwork in a one-to-one fashion . Nevertheless , the empirical performance of such distilled networks is still far inferior to DE teachers . ∗ Equal contribution 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . There may be many reasons why the students are not doing as well as DE teachers . We argue that a critical limitation of current distillation schemes is that they do not effectively transfer the diversity of DE teacher predictions to students . Consider an ensemble of deep neural networks trained for an image classification task . Given sufficient network capacity , each ensemble member is likely to achieve near-zero training error , meaning that the resulting outputs of ensemble members evaluated at the training set will be nearly identical , as shown in Fig . 1 . In such a situation , reusing the training data during distillation will not encourage the student to produce diverse predictions . This is quite critical , considering various results establishing that the effectiveness of DE comes from averaging out its diverse predictions . To this end , we propose a method that amplifies the diversity of students learning from DE teachers . Our idea is simple : instead of using the same training set , use a perturbed training set for which the predictions of ensemble members are likely to be diversified . To implement this idea , we employ Output Diversified Sampling ( ODS ) [ Tashiro et al. , 2020 ] , a sampling scheme that finds small input perturbations that result in significant changes in outputs . Empirically , we confirm that the inputs perturbed with ODS result in large deviations in the DE teacher outputs . We further justify our method by analyzing the role of ODS perturbation . Specifically , we show that distilling using training data perturbed with ODS can be interpreted as an approximate Jacobian matching procedure where ODS improves the sample efficiency of the approximate Jacobian matching . Intuitively , by approximately matching Jacobians of teachers and students , we are transferring the outputs of teacher networks evaluated not only on the training data points but also on nearby points . Using standard image classification benchmarks , we empirically validate that our distillation method promotes diversities in student network predictions , leading to improved performance , especially in terms of uncertainty estimation . 2 Backgrounds . 2.1 Settings and notations . The focus of this paper is on the K-way classification problem taking D-dimensional inputs . We denote a student neural network as S ( x ) : RD → [ 0 , 1 ] K and a teacher neural network as T ( x ) : RD → [ 0 , 1 ] K . S ( x ) and T ( x ) outputs class probabilities , and we denote the logits before softmax as Ŝ ( x ) and T̂ ( x ) . The kth element of an output is denoted as S ( k ) ( x ) . For DE teachers with j = 1 , . . . , M members , the jth ensemble member is denoted as Tj ( x ) . For a student network having M subnetworks , the jth subnetwork is denoted as Sj ( x ) . 2.2 Knowledge distillation . Knowledge distillation ( KD ) [ Hinton et al. , 2015 ] aims to train a student network S ( x ) by matching its outputs to the outputs of a teacher network T ( x ) . In addition to the standard cross-entropy loss between student outputs and the ground-truth labels , i.e. , LCE ( S ( x ) , y ) = − ∑K k=1 y ( k ) logS ( k ) ( x ) , KD additionally encourages the student ’ s output to be as close as possible to that of the teacher by also minimizing the KL-divergence between output class distributions of the student and the teacher : LKD ( S ( x ) , T ( x ) ; τ ) = −τ2 K∑ k=1 softmax ( k ) ( T̂ ( x ) τ ) log softmax ( k ) ( Ŝ ( x ) τ ) . ( 1 ) Here , τ > 0 is a temperature parameter which controls the smoothness of the distributions . We scale the KL divergence by τ2 since the soft targets scale as 1/τ2 . As a result , in the KD framework , the student imitates the teacher ’ s outputs for training examples by minimizing L = ( 1− α ) LCE + αLKD , where α ∈ ( 0 , 1 ) is a hyperparameter to be specified . 2.3 BatchEnsemble and one-to-one distillation . BatchEnsemble ( BE ) [ Wen et al. , 2020 ] is a lightweight ensemble method that reduces the number of parameters by weight sharing . Specifically , each layer of a BE model consists of shared weights W and rank-one factors rjs > j for j = 1 , . . . , M with M being number of subnetworks . Unlike the full ensemble that constructs full set of M weights , BE shares W across all subnetworks while the rank-one factors are kept specific to each subnetwork . The weight matrix of the jth subnetwork is computed as W ◦ rjs > j and consequently BE requires far fewer parameters than a full ensemble . During training , BE simultaneously trains all subnetworks by feeding a duplicate of the current minibatch for each rank-one factor rjs > j into the network . This procedure is easily vectorized because the computation graph for each subnetwork has the same structure , differing only in weight values . Mariet et al . [ 2020 ] proposed to further improve the performance of BE by distilling from fully trained DE teachers instead of training from scratch . During distillation , the parameters of jth subnetwork Sj ( x ) is trained by distilling from the jth DE teacher Tj ( x ) . The corresponding objective function is then written as ( c.f. , Fig . 2 without perturbation ) L = M∑ j=1 { ( 1− α ) LCE ( Sj ( x ) , y ) + αLKD ( Sj ( x ) , Tj ( x ) ; τ ) } . ( 2 ) Throughout the rest of the paper , we will use the BE as student networks and follow this one-to-one distillation framework . 2.4 Output diversified sampling ( ODS ) . ODS [ Tashiro et al. , 2020 ] is a sampling method that can maximize diversities in outputs of a given function ; it is originally introduced to replace random input perturbation to enhance the performance of the adversarial attack . Let F : RD → RK be a target function . ODS draws a perturbation vector guided from the gradient of the target function as follows . w ∼ Unif ( [ −1 , 1 ] ) K , εODS ( x , F , w ) = ∇x ( w > F ( x ) ) ‖∇x ( w > F ( x ) ) ‖2 ∈ RD . ( 3 ) The intuition behind this sampling strategy is to seek the direction in the input space that maximizes the similarity between the function output and a randomly sampled vector w. Following such a direction can make diverse output shifts due to the randomness in w . 3 Learning from Ensembles with Output Diversification . 3.1 One-to-one distillation with ODS . As discussed before , vanilla one-to-one distillation evaluates DE teachers on points where their outputs are nearly identical , so the students trained using them tend to exhibit low diversity in outputs . To avoid this undesirable behavior , we propose to instead train students on points perturbed with ODS . Let ( x , y ) be an training input pair . We first pick a random teacher Tr uniformly from { T1 , . . . , TM } and perturb x as x̃ = x+ ηεODS ( x , softmax ( T̂r ( x ) /τ ) , w ) , ( 4 ) where η > 0 is a step-size . Ideally , we would like to evaluate ODS vectors for each teacher to build M perturbed versions of x , but this would take too much time for gradient computation , especially when we have a large number of ensemble models . Instead , we just pick one of the teachers , use it to perturb the input , and share the perturbed input to train all students with the following loss function , L = M∑ j=1 { ( 1− α ) LCE ( Sj ( x ) , y ) + αLKD ( Sj ( x̃ ) , Tj ( x̃ ) ; τ ) } , ( 5 ) where we are evaluating the KL-divergence between teacher outputs and student outputs on the perturbed input . An implicit assumption here is that a diversifying direction computed from a specific random teacher can bring diversities in the outputs of all the teacher networks . This is based on the transferability assumption empirically justified in Tashiro et al . [ 2020 ] : ODS perturbations are computed on a surrogate model when we do not have access to the true model , assuming that the Jacobians computed from the surrogate model can approximate the true Jacobian to some extent . In our case , let T1 be a randomly picked teacher with which an ODS perturbation is computed , and T2 be another teacher . Assuming the transferability of Jacobian , we can let ∇xT1 ( x ) ≈ R∇xT2 ( x ) for some matrix R. Then the ODS perturbation computed from T1 ( x ) is εODS ( x , T1 , w ) ∝ w > ∇xT1 ( x ) ≈ ( R > w ) > ∇xT2 ( x ) ∝ εODS ( x , T2 , R > w ) . ( 6 ) That is , the ODS perturbation computed from T1 acts as another ODS perturbation with different random guide vector R > w . Hence , the same perturbation drives the outputs of T1 and T2 towards different directions . Note that without transferability , the Jacobians∇xT1 ( x ) and∇xT2 ( x ) would be vastly different and so are the directions w and R > w . This may be good in terms of diversification , but the resulting perturbed input act as an adversarial example not very useful for distillation , making the teachers completely disagree ( refer to Appendix A.2 for more details ) . Our experiments in Section 5.2 and Section 5.3 empirically re-confirm this transferability assumption . We can also interpret the random teacher selection as a stochastic approximation . Specifically , to find a direction maximizing diversities between teachers , we can consider an objective Ew1 , ... , wM [ M∑ j=1 w > j Tj ( x ) ] , ( 7 ) with w1 , . . . , wM i.i.d.∼ Unif ( [ −1 , 1 ] ) K . That is , we want to find a direction in the input space making the teachers { Tj } Mj=1 follow different guide vectors { wj } Mj=1 . Picking a random teacher and computing the ODS perturbation can be understood as a stochastic approximation of the gradient of this objective . Again , thanks to the transferability of the Jacobians , the variance of this stochastic gradient would be small , making it sufficient to use a single random teacher . In addition , we empirically found that scaling the ODS updates according to confidence slightly improves performance . That is , we perturb an input x as x̃ = x+ ηCmax ( x , Tr , τ ) εODS ( x , softmax ( T̂r ( x ) /τ ) , w ) , ( 8 ) where Cmax ( x , Tr , τ ) = maxk softmax ( k ) ( T̂r ( x ) /τ ) denotes the maximum class probability . This scheme takes larger steps towards ODS directions for datapoints with high confidence . We call this modification ConfODS . The training procedure with ODS perturbation is straightforward ( Algorithm 1 ) ; we can just add the ODS computation procedure and replace the loss function from the original KD training procedure . Fig . 2 depicts our overall training procedure as a diagram . Algorithm 1 Knowledge distillation from deep ensembles with ODS perturbations Require : Training data D = { ( xi , yi ) } Ni=1 . Require : Knowledge distillation weight α , temperature τ , learning rate β , ODS step-size η . Require : Deep ensemble teacher networks { Tj } Mj=1 where M denotes the size of ensembles . Let S be a student network which consists of M subnetworks { Sj } Mj=1 . Initialize parameters θ of the student network S. while not converged do Sample ( x , y ) ∼ D , w ∼ Unif ( [ −1 , 1 ] ) K , and a teacher index r uniformly from { 1 , . . . , M } . Sample diversified inputs x̃ = x+ ηεODS ( x , softmax ( T̂r ( x ) /τ ) , w ) LCE ← − ∑M j=1 ∑K k=1 y ( k ) logS ( k ) j ( x ) LKD ← −τ2 ∑M j=1 ∑K k=1 softmax ( k ) ( T̂ ( x̃ ) /τ ) log softmax ( k ) ( Ŝj ( x̃ ) /τ ) θ ← θ − β∇θ ( ( 1− α ) LCE + αLKD ) end while

Summary: This paper attempts to reduce the performance gap between an ensemble and a model that was distilled from it. The main idea is that, when distilling, the model should “absorb as much function diversity inside the ensemble as possible”. To do that, the paper proposes using ODS perturbation strategy, that samples diverse predictions for a given input form the ensemble, to use as label during distillation.