Scaling Ensemble Distribution Distillation to Many Classes with Proxy Targets

1 Introduction . As machine learning models are being deployed ever more widely into production , especially in safetycritical applications such as self-driving cars or medical diagnostics , obtaining reliable estimates of uncertainty in their predictions becomes increasingly important to guaranteeing their safety and reliability [ 1 ] . Ensemble methods are known to yield both improved predictive performance and reliable uncertainty estimates [ 2 , 3 , 4 , 5 , 6 ] . Furthermore , ensembles allow separating out total uncertainty into data and knowledge uncertainty2 . The former is the intrinsic uncertainty due to class overlap and noise inherent in the data , while the latter is the model ’ s uncertainty due to lack of understanding of the test data [ 7 ] . Estimates of knowledge uncertainty are often used to detect anomalous and unfamiliar inputs [ 8 , 9 , 10 , 7 ] as well as in active learning [ 11 ] . Despite their numerous advantages , using ensembles can be computationally prohibitive . Obtaining predictions in real time is often expensive even for a single model , and the hardware requirements for serving an ensemble scale linearly with size . As a result , over the past several years the area of ensemble distillation has gained increasing attention . Broadly speaking , distillation methods aim to train a single model which can approximate the behavior of an ensemble , but at low computational cost . In the simplest and most frequently used form of distillation [ 12 ] , the student model is trained to capture the average prediction of the ensemble : for example , in case of classification this reduces ⇤Equal contribution , order decided by coin toss . 2Data and Knowledge Uncertainty are also known as Aleatoric and Epistemic uncertainty . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . to minimizing the KL-divergence between the model and the ensemble mean . While this method allows the student to obtain predictive performance comparable to that of the original ensemble , the information about its distributional properties ( in other words , its diversity ) is lost in the process . As ensemble-based uncertainty estimation utilizes the information about disagreement between its members , such distillation approaches fail to preserve valuable information from the ensemble . Recently , several works have proposed distillation procedures that capture information about both the mean and distribution of ensemble predictions within a single model [ 13 , 14 , 15 , 16 ] . These can be broadly split into methods where the student has multiple output heads - one per ensemble member [ 14 , 15 ] , and those which yield a higher-order distribution over the ensemble ’ s outputs [ 13 , 16 ] . The former are conceptually simpler , but more computationally expensive and may model the ensemble in excessive detail ; the latter approaches are more light-weight , but also complex and , as we show in this work , suffer from optimization issues . This latter class of approaches , which we refer to as Ensemble Distribution Distillation ( EnD2 ) , is the focus of this work . Ensemble Distribution Distillation views outputs of each ensemble member as samples from a higher-order Dirichlet or Normal-Wishart distribution , and the student model learns the parameters of that distribution . Typically , EnD2 is done by maximizing the likelihood of the ensemble ’ s output distributions under the Dirichlet or Normal-Wishart prior . While this criterion is theoretically sound , we show that for large-scale classification tasks with many classes , gradient-based optimization of it is highly problematic , which limits its usefulness in real-life production scenarios . In this work , we investigate the poor convergence of models trained with Ensemble Distribution Distillation at scale . We analyze the the Dirichlet log-likelihood criterion and show that it leads to high gradient norm values that affect optimization . Specifically , if the ensemble members ’ output distributions have most probability mass allocated to a few classes , with the remainder spread among a long tail of exponentially less likely classes , then the gradients associated with the tail-classes will be significantly larger than those associated with high-probability classes . As a result , the optimization procedure focuses on modelling this distribution of tail-class probabilities , preventing convergence . To solve this , we propose a two-step solution . First , we transform the empirical distribution of ensemble member predictions into a Proxy-target Dirichlet distribution with the same statistics and to use it as the target during distribution distillation . Second , we show that it is crucial to minimize the reverse KL-divergence between the model and the Proxy-Dirichlet , as minimization of the forward KL-divergence exacerbates the optimization issues . The proposed training procedure allows the model to converge , mitigating the issue of gradient explosion . We demonstrate this by distribution-distilling ensembles of models trained on the ImageNet classification , LibriSpeech speech recognition and WMT17 English-German language translation datasets , where there are 1000 , 5000 and classes , respectively . To accomplish the latter two , we extend EnD2 to autoregressive models and refer to this approach as Sequence Ensemble Distribution Distillation ( SenD2 ) . On both datasets the distribution-distilled models outperform the baselines and yield uncertainty estimates competitive this those of the original ensemble . Thus , our contributions are as follows : • We draw attention to and analyze the convergence issues of Dirichlet likelihood optimization when applied to tasks with large numbers of classes ; • We propose an alternative two-step training approach for Ensemble Distribution Distillation which overcomes the aforementioned optimization issues ; • We adapt Ensemble Distribution Distillation to auto-regressive models and propose Sequence Ensemble-Distribution Distillation ( SEnD2 ) . 2 Preliminaries : Ensembles and Distillation . We view ensembles within a Bayesian framework , where the model parameters ✓ are random variables over which a prior p ( ✓ ) is placed . The posterior distribution p ( ✓|D ) is obtained via Bayes ’ rule : p ( ✓|D ) = p ( D|✓ ) p ( ✓ ) p ( D ) / p ( D|✓ ) p ( ✓ ) ( 1 ) Consider an ensemble of models { P ( y|x⇤ , ✓ ( m ) ) } Mm=1 sampled from the posterior : P ( y|x , ✓ ( m ) ) M m=1 ! P ( y|⇡ ( m ) ) M m=1 , ⇡ ( m ) = f ( x ; ✓ ( m ) ) , ✓ ( m ) ⇠ p ( ✓|D ) ( 2 ) where ⇡ are the parameters of a categorical distribution [ P ( y = ! 1 ) , · · · , P ( y = ! K ) ] T. The predictive distribution , or predictive posterior , for a test input x⇤ is obtained by marginalization : P ( y|x⇤ , D ) = Ep ( ✓|D ) ⇥ P ( y|x⇤ , ✓ ) ⇤ ⇡ 1 M MX m=1 P ( y|x⇤ , ✓ ( m ) ) ( 3 ) In practice this is intractable and we approximate via Monte-Carlo sampling . Given the ensemble , the entropy of the predictive posterior is a measure of total uncertainty . Knowledge uncertainty can be assessed via measures of the spread , or ‘ disagreement ’ , of the ensemble such as mutual information : I [ y , ✓|x⇤ , D ] | { z } Knowledge Uncertainty = H ⇥ Ep ( ✓|D ) [ P ( y|x⇤ , ✓ ) ] ⇤ | { z } Total Uncertainty Ep ( ✓|D ) ⇥ H [ P ( y|x⇤ , ✓ ) ] ⇤ | { z } Expected Data Uncertainty ( 4 ) An alternative measure of spread , which is particularly useful for ensemble of autoregressive sequence models , is the reverse-mutual information ( RMI ) [ 17 ] : M ⇥ y , ✓|x , D ⇤ = Ep ( ✓|D ) h EP ( y|x , D ) h ln P ( y|x , D ) P ( y|x , ✓ ) ii ( 5 ) While ensembles yield improved predictive performance and meaningful uncertainty estimates , they are computationally expensive . To reduce inference-time cost , an ensemble can be distilled into a single model P ( y|x ; ) by minimizing the KL-divergence to the predictive posterior of the ensemble : L EnD ( , Dens ) =Ep̂ ( x ) h KL ⇥ P ( y|x , D ) || P ( y|x ; ) ⇤i ( 6 ) This approach has been applied to tasks such as image classification and machine translation [ 12 , 18 ] . While distillation allows a single model to capture the ensemble ’ s predictive quality and estimates of total uncertainty at low cost , information about ensemble diversity is lost . Consequently , we can not obtain estimates of knowledge uncertainty , which are particularly useful for anomaly detection [ 7 , 13 ] . Recently , a distillation technique called Ensemble Distribution Distillation ( EnD2 ) was proposed [ 13 ] . Here , both the mean and diversity of an ensemble are captured within a single model via distillation into a Prior Network [ 19 ] , which is a model that parameterizes the Dirichlet distribution : p ( ⇡|x ; ) =Dir ( ⇡|↵ ) , ↵ = ez , z = f ( x ; ) , ↵c > 0 , ↵0 = KX c=1 ↵c ( 7 ) Distribution distillation is done in two steps . Firstly , a transfer dataset Dens = { x ( i ) , ⇡ ( i,1 : M ) } Ni=1 ⇠ p̂ ( x , ⇡ ) is composed of the inputs x ( i ) from the training set D = { x ( i ) , y ( i ) } Ni=1 and the ensemble ’ s output distributions { ⇡ ( i,1 : M ) } Ni=1 for each input . Secondly , given this transfer set , the model p ( ⇡|x ; ) is trained by minimizing the negative log-likelihood of each categorical distribution ⇡ ( im ) : L EnD2 ( , Dens ) = Ep̂ ( x ) ⇥ Ep̂ ( ⇡|x ) [ ln p ( ⇡|x ; ) ] ⇤ ( 8 ) Given a distribution-distilled Prior Network , the predictive distribution is given by the expected categorical distribution ⇡̂ under the Dirichlet prior : P ( y = ! c|x ⇤ ; ) = Ep ( ⇡|x⇤ ; ) [ P ( y = ! c|⇡ ) ] = ⇡c = ↵cPK k=1 ↵k = ezc PK k=1 e zk ( 9 ) Measures of total and knowledge uncertainty are obtained by considering the mutual information between the prediction y and the parameters of ⇡ of the categorical : I [ y , ⇡|x⇤ ; ] | { z } Knowledge Uncertainty = H ⇥ Ep ( ⇡|x⇤ ; ) [ P ( y|⇡ ) ] ⇤ | { z } Total Uncertainty Ep ( ⇡|x⇤ ; ) ⇥ H [ P ( y|⇡ ) ] ⇤ | { z } Expected Data Uncertainty ( 10 ) Alternative measures , such as RMI , can also be obtained , as detailed in appendix C. An alternative approach to capture ensemble diversity is to distill an ensemble into a mixture model which yields a separate output head for each ensemble member [ 14 , 15 ] . However , the principal downside of this approach is that it attempts to model the ensemble in excessive detail , as it models the behaviour of each ensemble member , rather than bulk distributional properties . This requires more flexible and powerful models . As a result , for good performance , it is necessary to split the model into multiple heads at an early stage , which significantly increases the computational and memory complexity . In contrast , EnD2 has a fixed computational and memory cost of one model , regardless of the size of the original ensemble , as it models the ensemble ’ s bulk distributional properties .

Summary: This paper considers uncertainty estimation for neural networks, specifically predictive uncertainty and knowledge uncertainty. For these problems, deep ensembles are considered state-of-the-art, but are expensive at training time and, more importantly, inference time. To address the latter problem, this paper explores Ensemble Distribution Distillation (EnD^2), an approach previously proposed by Malinin et al (2020) which consists of distilling the ensemble into a single model that captures both the mean and diversity of the original ensemble. However, this paper notes that EnD^2 fails to converge for problems where there is a large number of classes, which is attributed to the use of the forward KL divergence as the distillation criterion. Instead, the reverse-KL divergence is proposed using a proxy Dirichlet distribution to represent the ensemble. The precision of the proxy distribution is estimated using Stirling's approximation, and some constraints are imposed on the model parameters to ensure stability. Experiments on ImageNet (several variants) and NMT show that the proposed approach successfully retrains some of the accuracy and calibration benefits of the corresponding deep ensemble.