DEUP: Direct Epistemic Uncertainty Prediction

1 INTRODUCTION . A remaining great challenge for machine learning research is purposeful knowledge-seeking by learning agents , which can benefit from estimation of epistemic uncertainty , i.e. , a measure of lack of knowledge , which could potentially be eliminated with enough data if the learner converges to a Bayes-optimal predictor . Epistemic uncertainty estimation is already a key ingredient in active learning and Bayesian optimization ( Aggarwal et al. , 2014 ; Frazier , 2018 ) as well as exploration in Reinforcement Learning ( Kocsis & Szepesvári , 2006 ; Tang et al. , 2017 ) . Epistemic uncertainty ( EU ) thus tells us how much could be gained from learning around a particular area of state-space or input data space . But how should we even quantify it ? Much previous work has focused on model variance ( Gal & Ghahramani , 2016 ; Lakshminarayanan et al. , 2017 ) as a proxy thereof , i.e. , how different are the functions compatible with the learner ’ s preferences and the data . However , this approach does not take into account model bias ( or misspecification ) , due to the learner ’ s finite capacity or preferences , self-imposed to avoid overfitting or reduce computational costs . Recent work ( Masegosa , 2020 ) has also shown that Bayesian methods are suboptimal for learning predictive models when the model is misspecified . Fig . 1 illustrates with Gaussian Processes ( GP ) that a model ’ s posterior variance is not necessarily a good predictor of the reducible generalization error . 0.5 0.0 0.5 1.0 1.5 2.0 2.5 0 2 Ground truth GP mean 0.5 0.0 0.5 1.0 1.5 2.0 2.5 0 2 4 Epistemic uncertainty GP variance 0.5 0.0 0.5 1.0 1.5 2.0 2.5 0 2 4 Epistemic uncertainty DEUP 's prediction learner ) , and it is also the error which may be achieved by a Bayes-optimal predictor . Consider the case of a learning system with universal approximation properties , like modern deep learning , whose effective capacity can be increased with more data , either via model selection , early stopping or other means . In such cases , the part of the error which is not aleatoric is reducible with more data . For this reason , we propose to define epistemic uncertainty as the difference between the expected generalization error and aleatoric uncertainty . Estimating it is useful to an active or interactive learner because it could be reduced with enough data , especially in regions of the input space we care about , i.e. , where EU is large . In order to estimate EU , we propose DEUP , for Direct Epistemic Uncertainty Prediction , where we train a side learner ( the uncertainty predictor ) with an appropriate objective and training data in order to predict this generalization error . An estimate of EU can thus be obtained by subtracting an estimate of aleatoric uncertainty . In this paper we thus focus on two predictors : the main predictor ( for the learning task of original interest ) and an error predictor which can be used to predict epistemic uncertainty of the main predictor at given points , given some knowledge or estimate of the aleatoric uncertainty . The proposed methodology can be seen as an end-to-end approach to epistemic uncertainty estimation , i.e. , to help estimate where a learner would acquire the most information in terms of out-of-sample loss reduction , without relying on the brittle relationship between model variance and generalization error , which can get murky in the presence of misspecification ( bias ) , as we argue below . In order to select points with high EU , one would still need some kind of search ( in the space of x ) , e.g. , to propose new candidates for active learning ( Aggarwal et al. , 2014 ; Nguyen et al. , 2019 ; Bengio et al. , 2021 ) or sequential model optimization ( Kushner , 1964 ; Jones et al. , 1998 ; Snoek et al. , 2012 ) , or exploration in RL ( Kocsis & Szepesvári , 2006 ; Osband et al. , 2016 ; Janz et al. , 2019 ) . An interesting potential advantage of the proposed approach , compared with model variance estimators , is that DEUP error predictors can be explicitly trained to care about , and calibrate for estimating the generalization error for examples which may come from a distribution different from the distribution of most of the training examples , i.e. , an out-of-distribution ( OOD ) setting . Such non-stationarities arise naturally in contexts like active learning and reinforcement learning ( RL ) because the learner explores new areas of the input space . In these non-stationary settings , we typically want to retrain the main predictor as we acquire new training data , not just because more training data is generally better but also to better track the changing distribution and generalize to yet unseen but upcoming OOD inputs . This setting makes it challenging to train the main predictor but it also entails an even greater non-stationarity for the training data seen by the error predictor : a large error initially made at a point x before x is incorporated in the training set ( along with an outcome y ) will typically be greatly reduced after updating the main predictor with ( x , y ) . We propose several strategies to cope with this non-stationarity , the most important being to rely on additional features ( such as log-density estimates and model variance estimates at x ) . These features account for variations in prediction error arising due to updates of the main predictor , making DEUP ’ s training distribution more stationary . The experiments presented here explore first the scenario of purely trying to estimate epistemic uncertainty in terms of the recoverable generalization error , and second how this can be used in contexts of sequential model optimization ( or black-box optimization ) and exploration in RL ( where the issue of non-stationarity arises ) . The main contributions of this paper are thus the following : 1 . Novel end-to-end approach for directly estimating the epistemic uncertainty as the lack of knowledge about the true f∗ ( x ) = E [ Y |x ] , the Direct Epistemic Uncertainty Prediction ( DEUP ) approach . 2 . Experimental evidence that this direct estimation of epistemic uncertainty can lead to more precise estimates of the lack of knowledge than variance-based methods . 3 . A mean to mitigate the issue of non-stationarity arising in training the uncertainty estimator in interactive settings . 4 . A successful validation of the above in the contexts of ( a ) sequential model optimization and ( b ) exploration in RL . 2 ALEATORIC AND EPISTEMIC UNCERTAINTIES , MODEL VARIANCE . Consider a supervised learning algorithm ( or learner ) L mapping a dataset D to a predictive function f̂ = L ( D ) . L tries to minimize the expected value of a supervised learning loss l ( f̂ ( x ) , y ) ∈ R under unknown P ( Y |x ) . In this section , we focus on regression tasks with the squared loss l ( ŷ , y ) = ( ŷ−y ) 2 , with y ∈ R. In Appendix B , we provide the corresponding results for a general loss function . Definition 2.1 The expected loss of a predictor f̂ at x is defined as : U ( f̂ , x ) = ∫ ( f̂ ( x ) − y ) 2dP ( y|x ) . ( 1 ) Note how the expected loss is an unknown quantity , because we generally do not have access to the true data generating distribution P ( Y |x ) . The errors made by any predictor f̂ at x are due to both the inherent randomness of P ( Y |x ) ( aleatoric uncertainty ) and the lack of knowledge of the predictor that can be tackled by acquiring more information around x ( epistemic uncertainty ) . Next , we will define these concepts more formally . Because of this natural decomposition of the expected loss , we will also refer to it as total uncertainty , and we will use the two terms interchangeably . Bayes-optimal predictors f∗ satisfy the following equation at every x : ∀ỹ ∈ R ∫ ( f∗ ( x ) − y ) 2dP ( y|x ) ≤ ∫ ( ỹ − y ) 2dP ( y|x ) . They depend on the underlying data distribution only and not on learner L or trained predictor f̂ . Additionally , all Bayes-optimal predictors have the same total uncertainty for all x. Aleatoric uncertainty is the irreducible expected error , i.e. , that made by a Bayes-optimal predictor : Definition 2.2 We define aleatoric uncertainty at x as the total uncertainty of any Bayes-optimal predictor f∗ at x : A ( x ) = U ( f∗ , x ) . ( 2 ) and we note that by definition A ( x ) ≤ U ( f , x ) , ∀f ∀x . We now propose a definition of epistemic uncertainty of predictor f̂ as the gap between the error of f̂ at x and the lowest possible error at x , i.e. , the reducible part of the loss , given more knowledge . Definition 2.3 We define epistemic uncertainty E ( f̂ , x ) of a predictor f̂ at x as E ( f̂ , x ) = U ( f̂ , x ) −A ( x ) = U ( f̂ , x ) − U ( f∗ , x ) . ( 3 ) Using these definitions , we can present the main result about the regression setting . Proposition 1 In a regression task with ground truth P ( y|x ) = N ( y ; f∗ ( x ) , σ2 ( x ) ) , E ( f̂ , x ) = ( f̂ ( x ) − f∗ ( x ) ) 2 U ( f̂ , x ) = Ey∼P ( .|x ) [ ( f̂ ( x ) − y ) 2 ] = ( f̂ ( x ) − f∗ ( x ) ) 2 + σ2 ( x ) A ( x ) = Ey∼P ( .|x ) [ ( f∗ ( x ) − y ) 2 ] = σ2 ( x ) The proof is given in Appendix C. Additionally , we present in Appendix B a similar result for the negative log-likelihood loss , with discrete or continuous outputs . We will discuss in the remainder of this section the relation between this new definition of epistemic uncertainty and the more traditional way of modelling it through model variance or entropy . Consider a parametric model p ( Y |x , θ ) and a learner maintaining a distribution over parameters θ ∈ Θ , each corresponding to a predictor f in a parametric set of functions S , possibly starting from a prior p ( θ ) that would lead to a posterior distribution p ( θ|D ) . Clearly , the fact that multiple θ ’ s and corresponding values of f are compatible with the data and the prior indicates lack of knowledge . Because the lack of knowledge indicates where an interactive learner should acquire more information , this justifies the usage of dispersion measures , such as the variance or the entropy of the posterior predictive , as proxies for epistemic uncertainty . However , the limited capacity of S or the prior p ( θ ) may pull the optimal f̃ in S away from the Bayes-optimal predictor f∗ . We can refer to these self-imposed constraints as a form of bias , in the sense that the learner is usually biased towards the prior preferences , e.g. , towards smoother predictors . In particular , this issue can arise when training neural networks with limited data , where the network may not use all of its capacity due to implicit ( and not fully understood ) regularization properties of SGD ( Kale et al. , 2021 ) , explicit regularization , early stopping or a combination of these , which can induce a preference on the functions it learns . In Fig . 2 , we illustrate this gap with the bias function b ( x ) = f∗ ( x ) − f̃ ( x ) . Because of this bias , model variance can not be an accurate measure of lack of knowledge E ( f̂ , x ) = ( f̂ ( x ) − f∗ ( x ) ) 2 in the general case , as we have shown with the results in Fig . 1 . In Deep Ensembles ( Lakshminarayanan et al. , 2017 ) e.g. , if all the networks in the ensemble tend to fail in systematic ways , this aspect of prediction failure will not be captured by variance . Whereas Deep Ensembles variance provide us uncertainty regarding which of the networks we could draw is correct , this does not tell us how poor that network is even in a noise-free setting . On the other hand , with flexible models like neural networks , adding examples around x where b ( x ) 2 is large may allow to increase capacity around x and reduce b ( x ) 2 . In section 5 , we will provide more experimental evidence supporting our claim that directly estimating the epistemic uncertainty as the reducible part of the loss can lead to better measures of lack of knowledge , compared to variance-based methods , especially in interactive learning settings .

Summary: Given some supervised task, this paper redefines the uncertainty of a solution as its generalisation risk. The authors then propose to learn a secondary function to estimate the generalisation risk of the first. This is done by using held-out points as training data for the secondary model. In turn, this held out used to estimate the primary model error comes from a k-fold split.  The inputs to the secondary model are the data points being evaluated, estimates of the predictive variance of the primary model, density estimates from a generative model and whether the point has been observed by the primary model.  The authors posit that the advantage of their method is that it can capture uncertainty due to model selection bias, something omitted by existing work, which focuses on the variance of the learnt estimator. The authors provide a diverse range of experiments: OOD rejection in image classification, active learning for drug discovery, function optimisation and exploration in reinforcement learning.