Regression Prior Networks

1 INTRODUCTION . Neural Networks have become the standard approach to addressing a wide range of machine learning tasks ( Girshick , 2015 ; Simonyan & Zisserman , 2015 ; Villegas et al. , 2017 ; Mikolov et al. , 2013b ; a ; 2010 ; Hinton et al. , 2012 ; Hannun et al. , 2014 ; Caruana et al. , 2015 ; Alipanahi et al. , 2015 ) . However , in order to improve the safety of AI systems ( Amodei et al. , 2016 ) and avoid costly mistakes in high-risk applications , such as self-driving cars , it is desirable for models to yield estimates of uncertainty in their predictions . Ensemble methods are known to yield both improved predictive performance and robust uncertainty estimates ( Gal & Ghahramani , 2016 ; Lakshminarayanan et al. , 2017 ; Maddox et al. , 2019 ) . Importantly , ensemble approaches allow interpretable measures of uncertainty to be derived via a mathematically consistent probabilistic framework . Specifically , the overall total uncertainty can be decomposed into data uncertainty , or uncertainty due to inherent noise in the data , and knowledge uncertainty , which is due to the model having limited uncertainty of the test data ( Malinin , 2019 ) . Uncertainty estimates derived from ensembles have been applied to the detection of misclassifications , out-of-domain inputs and adversarial attack detection ( Carlini & Wagner , 2017 ; Smith & Gal , 2018 ) , and active learning ( Kirsch et al. , 2019 ) . Unfortunately , ensemble methods may be computationally expensive to train and are always expensive during inference . A class of models called Prior Networks ( Malinin & Gales , 2018 ; 2019 ; Malinin , 2019 ; Sensoy et al. , 2018 ) was proposed as an approach to modelling uncertainty in classification tasks by emulating an ensemble using a single model . Prior Networks parameterize a higher order conditional distribution over output distributions , such as the Dirichlet distribution . This enables Prior Networks to efficiently yield the same interpretable measures of total , data and knowledge uncertainty as an ensemble . Unlike ensembles , the behaviour of Prior Networks ’ higher-order distribution is specified via a loss function , such as reverse KL-divergence ( Malinin & Gales , 2019 ) , and training data . However , such Prior Networks yield predictive performance consistent with that of a single model trained via Maximum Likelihood , which is typically worse than that of an ensemble . This can be overcome via Ensemble Distribution Distillation ( EnD2 ) ( Malinin et al. , 2020 ) , which is an approach that allows distilling an ensemble into Prior Network such that measures of ensemble diversity are preserved . This enables to retain both the predictive performance and uncertainty estimates of an ensemble at low computational and memory cost . Finally , it is important to point out that a related class of evidential methods has concurrently appeared ( Sensoy et al. , 2018 ; Amini et al. , 2020 ) . Structurally they yield models similar to Prior Networks , but are trained in a different fashion . While Prior Networks have many attractive properties , they have only been applied to classification tasks . In this work we develop Prior Networks for regression tasks by considering the Normal-Wishart distribution - a higher-order distribution over the parameters of multivariate normal distributions . Specifically , we extend theoretical work from ( Malinin , 2019 ) , where such models are considered , but never evaluated . We derive all measures of uncertainty , the reverse KL-divergence training objective , and the Ensemble Distribution Distillation objective in closed form . Regression Prior Networks are then evaluated on synthetic data , selected UCI datasets and the NYUv2 and KITTI monocular depth estimation tasks , where they are shown to yield comparable or better performance to state-of-the-art single-model and ensemble approaches . Crucially , they enable , via EnD2 , to retain the predictive performance and uncertainty estimates of an ensemble within a single model . 2 REGRESSION PRIOR NETWORKS . In this section we develop Prior Network models for regression tasks . While typical regression models yield point-estimate predictions , we consider probabilistic regression models which parameterizes a distribution p ( y|x , θ ) over the target y ∈ RK . Typically , this is a normal distribution : p ( y|x , θ ) = N ( y|µ , Λ ) , { µ , Λ } = f ( x ; θ ) ( 1 ) where µ is the mean , and Λ the precision matrix , a positive-definite symmetric matrix . While normal distributions are usually defined in terms of the covariance matrix Σ = Λ−1 , parameterization using the precision is more numerically stable during optimization ( Bishop , 2006 ; Goodfellow et al. , 2016 ) . While a range of distributions over continuous random variables can be considered , we will consider the normal as it makes the least assumptions about the nature of y and is mathematically simple . As in the case for classification , we can consider an ensemble of networks which parameterize multivariate normal distributions ( MVN ) { p ( y|x , θ ( m ) ) } Mm=1 . This ensemble can be interpreted as a set of draws from a higher-order implicit distribution over normal distributions . A Prior Network for regression would , therefore , emulate this ensemble by explicitly parameterizing a higher-order distribution over the parameters µ and Λ of a normal distribution . One sensible choice is the formidable Normal-Wishart distribution ( Murphy , 2012 ; Bishop , 2006 ) , which is a conjugate prior to the MVN . This parallels how the Dirichlet distribution , the conjugate prior to the categorical , was used in classification Prior Networks . The Normal-Wishart distribution is defined as follows : NW ( µ , Λ|m , L , κ , ν ) = N ( µ|m , κΛ ) W ( Λ|L , ν ) ( 2 ) wherem and L are the prior mean and inverse of the positive-definite prior scatter matrix , while κ and ν are the strengths of belief in each prior , respectively . The parameters κ and ν are conceptually similar to precision of the Dirichlet distribution α0 . The Normal-Wishart is a compound distribution which decomposes into a product of a conditional normal distribution over the mean and a Wishart distribution over the precision . Thus , a Regression Prior Network ( RPN ) parameterizes the NormalWishart distribution over the mean and precision of normal output distributions as follows : p ( µ , Λ|x , θ ) = NW ( µ , Λ|m , L , κ , ν ) , { m , L , κ , ν } = Ω = f ( x ; θ ) ( 3 ) where Ω = { m , L , κ , ν } is the set of parameters of the Normal-Wishart predicted by neural network . The posterior predictive of this model is the multivariate Student ’ s T distribution ( Murphy , 2012 ) , which is the heavy-tailed generalization of the multivariate normal distribution : p ( y|x , θ ) = Ep ( µ , Λ|x , θ ) [ p ( y|µ , Λ ) ] = T ( y|m , κ+ 1 κ ( ν −K + 1 ) L−1 , ν −K + 1 ) ( 4 ) In the limit , as ν →∞ , the T distribution converges to a normal distribution . The predictive posterior of the Prior Network given in equation 4 only has a defined mean and variance when ν > K + 1 . Figure 1 depicts the desired behaviour of an ensemble of normal distributions sampled from a Normal-Wishart distribution . Specifically , the ensemble should be consistent for in-domain inputs in regions of low/high data uncertainty , as in figures 1a-b , and highly diverse both in the location of the mean and in the structure of the covariance for out-of-distribution inputs , as in figure 1c . Samples of continuous output distributions from a regression Prior Network should yield the same behaviour . Measures of Uncertainty Given an RPN which displays these behaviours , we can compute closedform expression for all uncertainty measures previously discussed for ensembles and Dirichlet Prior Networks ( Malinin , 2019 ) . We can obtain measures of knowledge , total and data uncertainty by considering the mutual information between y and the parameters of the output distribution { µ , Λ } : I [ y , { µ , Λ } ] ︸ ︷︷ ︸ Knowledge Uncertainty = H [ Ep ( µ , Λ|x , θ ) [ p ( y|µ , Λ ) ] ] ︸ ︷︷ ︸ Total Uncertainty −Ep ( µ , Λ|x , θ ) [ H [ p ( y|µ , Λ ) ] ] ︸ ︷︷ ︸ Expected Data Uncertainty ( 5 ) This expression consists of the difference between the differential entropy of the posterior predictive and the expected differential entropy of draws from the Normal-Wishart prior . We can also consider the expected pairwise KL-divergence ( EPKL ) between draws from the Normal-Wishart prior : K [ y , { µ , Λ } ] = − Ep ( y|x , θ ) [ Ep ( µ , Λ|x , θ ) [ ln p ( y|µ , Λ ) ] ] − Ep ( µ , Λ|x , θ ) [ H [ p ( y|µ , Λ ) ] ] ( 6 ) This is an upper bound on mutual information ( Malinin , 2019 ) . Notably , estimates of data uncertainty are unchanged . One practical use of EPKL is comparison with ensembles , as it is not possible to obtain a tractable expression for the mutual information of a regression ensemble ( Malinin , 2019 ) . Alternatively , we can consider measures of uncertainty derived via the law of total variance : Vp ( µ , Λ|x , θ ) [ µ ] ︸ ︷︷ ︸ Knowledge Uncertainty = Vp ( y|x , θ ) [ y ] ︸ ︷︷ ︸ Total Uncertainty − Ep ( µ , Λ|x , θ ) [ Λ−1 ] ︸ ︷︷ ︸ Expected Data Uncertainty ( 7 ) This yields a similar decomposition to mutual information , but only first and second moments are considered . We provide closed-form expressions of ( 5 ) - ( 7 ) in appendix A . Note , however , that these variance-based measures are not scale-invariant , and are therefore sensitive to the scale of predictions which the model makes . This is relevant to applications such as depth-estimation , where we are likely to encounter images with a wide range of possible depths , and therefore varying scale of predictions . We omit the closed-form expressions for all terms here and instead provide them in appendix A. RKL training objective Having discussed how to construct Prior Networks for regression , we now discuss how they can be trained . Prior Networks are trained using a multi-task loss , where an in-domain loss Lin and an out-of-distribution ( OOD ) loss Lout are jointly minimized : L ( θ , Dtr , Dout ) = Eptr ( y , x ) [ Lin ( y , x , θ ) ] + γ · Epout ( x ) [ Lout ( x , θ ) ] ( 8 ) The OOD loss is necessary to teach the model the limit of its knowledge ( Malinin , 2019 ) and define a decision boundary between the in-domain and out-of-domain regions . In some applications the choice of OOD data can be particularly difficult . Normal distributions sampled from a Prior Network should be consistent and reflect the correct level of data uncertainty in-domain , and diverse in both mean and precision out-of-domain . Achieving the former is challenging , as the training data only consists of samples of inputs x and targets y , there is no access to the underlying distribution , and associated data uncertainty , represented by the precision Λ . Effectively , we are attempting to train a Normal-Wishart distribution from targets sampled from normal distribution that are sampled from the Normal-Wishart , rather than on the normal distribution samples themselves , which is challenging . However , it was shown that for Dirichlet Prior Networks minimizing the reverse KL-divergence between the model and an appropriate target Dirichlet induces in expectation the correct estimate of data uncertainty ( Malinin & Gales , 2019 ) . As the Normal-Wishart , like the Dirichlet , is a conjugate prior and exponential family member , the precision can be induced in expectation by considering the reverse KL-divergence between the model p ( µ , Λ|x , θ ) and a target Normal-Wishart p ( µ , Λ|Ω̂ ( i ) ) corresponding to each x ( i ) . The appropriate Normal-Wishart is specified via Bayes ’ s rule : p ( µ , Λ|Ω̂ ( i ) ) ∝ p ( y ( i ) |µ , Λ ) β̂p ( µ , Λ|Ω0 ) ( 9 ) where p ( y ( i ) |µ , Λ ) is a normal distribution and Ω0 = { m0 , L0 , κ0 , ν0 } are the parameters of the prior p ( µ , Λ|Ω0 ) defined as follows : m0 = 1 N N∑ i=1 y ( i ) , L−10 = ν0 N N∑ i=1 ( y ( i ) −m0 ) ( y ( i ) −m0 ) T , κ0 = , ν0 = K + 1 + ( 10 ) In other words , we consider a semi-informative prior which corresponds to the mean and scatter matrix of marginal distribution p ( y ) , and we see each sample of the training data β̂ times . The hyper-parameter β̂ allows us to weigh the effect of the prior and the data . is a small value , like 10−2 , so that κ0 and ν0 yield a maximally un-informative , but proper predictive posterior . The reason to use a semi-informative prior is that in regression tasks , unlike classification tasks , uninformative priors are improper and lead to infinite differential entropy . Furthermore , we do know something about the data purely based on the marginal distribution , and it is sensible to use that as the prior . The reverse KL-divergence loss can then be expressed as : L ( y , x , θ ; β̂ , Ω0 ) = KL [ p ( µ , Λ|x , θ ) ‖p ( µ , Λ|Ω̂ ( i ) ) ] = β̂ · Ep ( µ , Λ|x , θ ) [ − ln p ( y|µ , Λ ) ] + KL [ p ( µ , Λ|x , θ ) ‖p ( µ , Λ|Ω0 ) ] + Z ( 11 ) where Z is a normalization constant independent of parameters θ . For in-domain data , β̂ can be set to a large value , and for out-of-domain training data β̂ = 0 , so that the model regresses to the prior . In-domain , the prior will add a degree of smoothing , which may prevent over-fitting and improve performance on small datasets . A large value of β̂ means the model will yield a larger κ , ν for in-domain data . This will result in the first term of ( 11 ) being very close to the expected negative log-likelihood of the predictive posterior p ( y|x , θ ) , thereby avoiding degradation of predictive performance . The derivation and closed-form expression for this loss is provided in appendix A . Ensemble Distribution Distillation An exciting task which Prior Networks can solve is Ensemble Distribution Distillation ( EnD2 ) ( Malinin et al. , 2020 ) , where the distribution of an ensemble ’ s predictions is distilled into a single model . EnD2 enables retaining an ensemble ’ s improved predictive performance and uncertainty estimates within a single model at low cost . In contrast , standard Ensemble Distillation ( EnD ) minimizes the KL-divergence between a model and the ensemble : LEnD ( φ , Dtrn ) = 1 NM N∑ i=1 M∑ m=1 KL [ p ( y|x ( i ) , θ ( m ) ) ||p ( y|x ( i ) , φ ) ] ] ( 12 ) This loses information about ensemble diversity . A complication that occurs in probabilistic regression models is that the student model will try to fit a single normal distribution on top of a mixture distribution , spreading itself across each model . This may result in both poor predictive performance and poor estimates of uncertainty . An approach to both overcome this and retain information about diversity was considered in ( Tran et al. , 2020 ; Wu et al. , 2020 ) , where the student model parameterizes a mixture distribution , where each component of the mixture models a particular model in the ensemble . We refer to this as Mixture-Density Ensemble Distillation ( MD-EnD ) : p ( y|x , φ ) = 1 M M∑ m=1 N ( y|µ ( m ) , Λ ( m ) ) , { µ ( m ) , Λ ( m ) } Mm=1 = f ( x ; φ ) LMD-EnD ( φ , Dtrn ) = 1 NM N∑ i=1 M∑ m=1 KL [ p ( y|x ( i ) , θ ( m ) ) ||N ( y|µ ( m ) , Λ ( m ) ) ) ] ] ( 13 ) This clearly overcomes the issue of distributional mismatch , resolving any issues with poor predictive performance . It also allows the model to , theoretically , retain information about ensemble diversity . However , it may be challenging to fully replicate the behaviour of each ensemble member in detail by having only multiple output heads - thus , splitting the model at an earlier point in the network may be necessary , which increases computational cost . EnD2 avoids the problem by directly modelling the bulk behaviour of the ensemble , rather than the behavior of each individual model . EnD2 can be implemented for regression via RPNs as follows . Consider an ensemble { p ( y|x , θ ( m ) ) } Mm=1 , where each model yields the mean and precision of a normal distribution . We can define an empirical distribution over the mean and precision as follows : p̂ ( µ , Λ , x ) = { { µ ( mi ) , Λ ( mi ) } Mm=1 , x ( i ) } N i=1 = Dtrn ( 14 ) EnD2 can then be accomplished by minimizing the negative log-likelihood of the ensemble ’ s mean and precision under the Normal-Wishart prior : LEnD2 ( φ , Dtrn ) = Ep̂ ( µ , Λ , x ) [ − ln p ( µ , Λ|x ; φ ) ] = Ep̂ ( x ) [ KL [ p̂ ( µ , Λ|x ) ||p ( µ , Λ|x ; φ ) ] ] + Z ( 15 ) This is equivalent to minimizing the KL-divergence between the model and the empirical distribution of the ensemble . Note that here , unlike in the previous section , the parameters of a normal distribution are available for every inputx , making forward KL-divergence the appropriate loss function . However , while this is a theoretically sound approach , the optimization might be numerically challenging . Similarly to ( Malinin et al. , 2020 ) we propose a temperature-annealing trick to make the optimization process easier . First , the ensemble is reduced to it ’ s mean : µ ( mi ) T = 2 T + 1 µ ( mi ) + T − 1 T + 1 µ̄ ( i ) , µ̄ ( i ) = 1 M M∑ m=1 µ ( mi ) Λ −1 ( mi ) T = 2 T + 1 Λ−1 ( mi ) + T − 1 T + 1 Λ̄−1 ( i ) , Λ̄−1 ( i ) = 1 M M∑ m=1 Λ̄−1 ( mi ) ( 16 ) We use inverses of the precision matrix Λ because we are interpolating the covariance matrices Σ. Secondly , the predicted κ̃ = Tκ and ν̃ = Tν are multiplied by T in order to make the Normal-Wishart sharp around the mean . The loss is divided by T to avoid scaling the gradients by T , yielding : pT ( µ , Λ|x , φ ) = NW ( µ , Λ|m , L , Tκ , Tν ) , { m , L , κ , ν } = f ( x ; φ ) LEnD2 ( φ , Dtrn ; T ) = 1 T Ep̂T ( µ , Λ , x ) [ − ln pT ( µ , Λ|x ; T , φ ) ] ( 17 ) This splits learning into two phases . First , when the temperature is high , the model learns to match the ensemble ’ s mean ( first moment ) . Second , as the temperature is annealed down to 1 , the model will gradually focus on learning higher moments of the ensemble ’ s distribution . This trick may be necessary , as the ensemble may have a highly non-Normal-Wishart distribution , which may be challenging to learn . Note that for EnD2 it may be better to parameterize the Normal-inverseWishart distribution over the mean and covariance due to numerical stability concerns . However , for consistency , we describe EnD2 in terms of the Normal-Wishart . Finally , we emphasize that EnD2 does not require OOD training data , unlike the RKL objective above . This eliminates the non-trivial challenge of finding appropriate OOD data . Related Approaches It is necessary to mention evidential approaches , specifically Deep Evidential Regression ( Amini et al. , 2020 ) . Structurally , it considers models which are similar to RPNs , and is thus compared to in section 5 . However , they do not use OOD training data nor attempt to emulate and generalize an ensemble ’ s behaviour to enforce high-uncertainty behaviour in OOD regions . Rather , they are trained by maximizing the likelihood of the T distribution ( equation ( 4 ) ) together with an evidence regularizer which forces the model to yield high ν , κ on regions with low absolute error , and low evidence in regions of high L1 error . However , this seems susceptible to pathologies , such as making the model overconfident if it over-fits the training data and achieves zero MAE , or introducing an inverse correlation between ν , κ and the scale of predictions in datasets where the scale of predictions range widely . Thus , while they yield encouraging results , their principle of action remains unclear . Another related recent research direction is in developing computationally efficient ensembles ( Wen et al. , 2020 ) . Here , rather than emulating an ensemble using a single model , the goal is to generate high diversity in predictions while re-using large portions of a neural network model , thereby making for compact ensembles . While such approaches are more computationally efficient than using multiple independent models and use only a little more memory than a single model on disk , they still require using M times as much GPU memory at run time , which can be a significant limitation in resource constrained applications , like mobile devices or autonomous vehicles .

Summary: Prior Networks (Malinin & Gales, 2018) use Dirichlet prior over categorical predictive distributions to distill ensembles for classification tasks. This paper extends Prior Networks to the regression setting by using a Normal-Wishart prior in order to attempt to match the predictive diversity. The authors define the model and loss terms including analytical derivation and evaluate their proposed approach with synthetic data, UCI datasets and monocular depth estimation.