Natural Posterior Network: Deep Bayesian Predictive Uncertainty for Exponential Family Distributions

1 INTRODUCTION Accurate and rigorous uncertainty estimation is key for reliable machine learning models in safety-critical domains [ 48 ] . It quantifies the confidence of machine learning models , thus allowing them to validate knowledgeable predictions corresponding to correct/wrong predictions , flag predictions on unknown input domains corresponding to anomaly or Out-ofDistribution detection , or detect natural shifts of the data facilitating real-time model maintenance ( Filos et al. , 2019 ; Malinin et al. , 2021 ; Ovadia et al. , 2019 ) . Specifically , a reliable model can handle all these failure modes with high-quality estimates of aleatoric and epistemic uncertainty ( Gal , 2016 ) . These two levels of uncertainty allow a model to account for both irreducible data uncertainty ( e.g . a fair dice ’ s chance of 1/6 for each face ) and uncertainty due to the lack of knowledge about unseen data ( e.g . input features differing significantly from training data or a covariate shift ) respectively . Aleatoric and epistemic uncertainty levels can eventually be combined into an overall predictive uncertainty ( Gal , 2016 ) Traditional neural networks are not readily applicable in safety-critical domains as they show overconfident prediction , in particular on data that is different from training data ( Guo et al. , 2017 ; Lakshminarayanan et al. , 2017 ) . To mitigate this problem , an important model family for uncertainty estimation directly pre- dicts the parameters of a conjugate prior distribution on the predicted target probability distribution , thus accounting for the different levels of uncertainty . These models are efficient as they only require a single forward pass for target and uncertainty prediction . Most of those models focus on classification and thus predict parameters of a Dirichlet distribution ( Biloš et al. , 2019 ; Charpentier et al. , 2020 ; Gast & Roth , 2018 ; Malinin & Gales , 2018 ; 2019 ; Nandy et al. , 2020 ; Sensoy et al. , 2020 ; Shi et al. , 2020 ; Zhao et al. , 2020b ) . However , only two works ( Amini et al. , 2020 ; Malinin et al. , 2020a ) have focused on regression by learning parameters of a Normal Inverse-Gamma ( NIG ) distribution as conjugate prior . Hence , all these models are limited to a single task ( e.g . either classification or regression ) . Some approaches even require out-of-distribution ( OOD ) data at training time ( Malinin & Gales , 2018 ; 2019 ) which is an unrealistic assumption in many real-world applications where anomalies are a priori diverse , rare or unknown . Our contribution . We propose Natural Posterior Network ( NatPN ) as a new approach parametrizing conjugate prior distributions for versatile uncertainty estimation . NatPN is motivated from both the theoretical and practical perspective . ( 1 ) NatPN can estimate predictive uncertainty for any task described by the general group of exponential family distributions contrary to existing approaches from this family of models . Notably , this encompasses very common tasks such as classification , regression and count prediction which can be described with Categorical , Normal and Poisson distributions , respectively . ( 2 ) In theory , NatPN is based on a new unified exponential family framework which performs an input-dependent Bayesian update . For every input , it predicts the parameters of the posterior over the target exponential family distribution . We show that this Bayesian update is guaranteed to predict high uncertainty far from training data . ( 3 ) In practice , NatPN requires no OOD data for training , only adds a single normalizing flow density to the last predictor layer and provides fast uncertainty estimation in a single forward pass . Our extensive experiments showcase the high performances of NatPN for various criteria ( accuracy , calibration , OOD and shift detection ) and tasks ( classification , regression and count prediction ) . We illustrate the accurate aleatoric and predictive uncertainty predictions of NatPN on two toy examples for classification and regression in Fig . 1 . None of the conjugate prior related works have similar theoretical and practical properties . 2 RELATED WORK . In this section , we describe other work related to uncertainty estimation for supervised learning . We refer to Gawlikowski et al . ( 2021 ) for a detailed survey on uncertainty estimation in deep learning . Sampling-based methods . A first family of models estimates uncertainty by aggregating statistics ( e.g . mean and variance ) from different samples of an implicit predictive distribution . Examples are ensemble ( Kim & Ghahramani , 2012 ; Lakshminarayanan et al. , 2017 ; Simpson et al. , 2012 ; Wen et al. , 2020 ; Wenzel et al. , 2020 ) and dropout ( Gal & Ghahramani , 2016 ) models which provide high-quality uncertainty estimates ( Ovadia et al. , 2019 ) at the cost of an expensive sampling phase at inference time . Moreover , ensembles usually require training multiple models . Further , Bayesian neural networks ( BNN ) ( Blundell et al. , 2015 ; Maddox et al. , 2019 ; Ritter et al. , 2018 ) model the uncertainty on the weights and also require multiple samples to estimate the uncertainty on the final prediction . While recent BNNs have shown reasonably good performance ( Dusenberry et al. , 2020 ; Farquhar et al. , 2020 ; Osawa et al. , 2019 ) , modelling the distribution on the weights suffers from pathological behavior thus limiting these approaches in practice ( Foong et al. , 2020 ; Graves , 2011 ; Izmailov et al. , 2021 ) . In particular , Izmailov et al . ( 2021 ) uses an enormous computation budget by parallelizing the computation over 512 TPUv3 devices and running tens of thousands of training epochs to achieve a more exact Bayesian inference which is not suitable for practical applications . In contrast , NatPN predicts uncertainty in a single forward pass with a closed-form posterior distribution over the target variable . NatPN does not model uncertainty on the weights . Sampling-free methods . A second family of models is capable of estimating uncertainty in a single forward pass . The family of models parametrizing conjugate prior distributions is the main focus of this paper ( Amini et al. , 2020 ; Biloš et al. , 2019 ; Charpentier et al. , 2020 ; Gast & Roth , 2018 ; Malinin & Gales , 2018 ; 2019 ; Malinin et al. , 2020a ; Nandy et al. , 2020 ; Sensoy et al. , 2020 ; Shi et al. , 2020 ; Zhao et al. , 2020b ) . Beyond this family of models , we differentiate between three other families of sampling-free models for uncertainty estimation . A family of models aims at learning deep Gaussian processes by learning random features projections or by learning inducing points ( Biloš et al. , 2019 ; Lakshminarayanan et al. , 2020 ; van Amersfoort et al. , 2020 ; 2021 ) . Another family of models aims at propagating uncertainty across layers ( Gast & Roth , 2018 ; Hernandez-Lobato & Adams , 2015 ; Postels et al. , 2019 ; Shekhovtsov & Flach , 2019 ; Wang et al. , 2016 ) . They model uncertainty at the weight and/or activation levels and are generally constrained to specific transformations . In contrast , NatPN only models the uncertainty on the predicted target variable and does not enforce any constraint on the encoder architecture . Further , some of the models propagating uncertainty already used the exponential family framework ( Ranganath et al. , 2015 ; Wang et al. , 2016 ) . However , while they parametrize exponential family distributions , NatPN parametrizes the conjugate prior of the target exponential family distributions which accounts for the epistemic uncertainty . Finally , while the family of calibration models aims at calibrating predictions ( Kuleshov et al. , 2018 ; Moon et al. , 2020 ; Rahimi et al. , 2020 ; Song et al. , 2019 ; Zhao et al. , 2020a ) , NatPN aims at accurately modelling both aleatoric and epistemic uncertainty on in- and out-of-distribution data . 3 NATURAL POSTERIOR NETWORK . At the very core of NatPN stands the Bayesian update rule : Q ( θ | D ) ∝ P ( D |θ ) × Q ( θ ) where P ( D |θ ) is the target distribution of the target data D given its parameter θ , and Q ( θ ) and Q ( θ | D ) are the prior and posterior distributions , respectively , over the target distribution parameters . The target distribution P ( D |θ ) could be any likelihood describing the observed target labels . The Bayesian update has three main advantages : ( 1 ) it introduces a prior belief which represents the safe default prediction if no data is observed , ( 2 ) it updates the prior prediction based on observed target labels , and ( 3 ) it assigns a confidence for the new target prediction given the aggregated evidence count of observed target labels . While NatPN is capable to perform a Bayesian update for every possible input given the observed training data , we first recall the Bayesian background for a single exponential family distribution . 3.1 EXPONENTIAL FAMILY DISTRIBUTION . Distributions from the exponential family are very widely used and have favorable analytical properties . Indeed , ( 1 ) they cover a wide range of target variables like discrete , continuous , counts or spherical coordinates , and ( 2 ) they benefit from intuitive and generic formulae for their parameters , density functions and statistics which can often be evaluated in closed-form . Important examples of exponential family distributions are Normal , Categorical and Poisson distributions ( see Tab . 1 ) . Formally , an exponential family distribution on a target variable y ∈ R with natural parameters θ ∈ RL can be denoted as P ( y |θ ) = h ( y ) exp ( θTu ( y ) −A ( θ ) ) ( 1 ) where h : R→ R is the carrier or base measure , A : RL → R the log-normalizer and u : R→ RL the sufficient statistics ( Bishop , 2006 ; Nielsen & Nock , 2010 ) . The entropy of an exponential family distribution can always be written as H [ P ] = A ( θ ) − θT∇θA ( θ ) − E [ log h ( y ) ] ( Nielsen & Nock , 2010 ) . An exponential family distribution always admits a conjugate prior , which often also is a member of the exponential family : Q ( θ |χ , n ) = η ( χ , n ) exp ( nθTχ− nA ( θ ) ) ( 2 ) where η ( χ , n ) is a normalization coefficient , χ ∈ RL are prior parameters and n ∈ R+ is the evidence . Given a set of N target observations { y ( i ) } Ni , it is easy to compute a closed-form Bayesian update Q ( θ |χpost , npost ) ∝ P ( { y ( i ) } Ni |θ ) ×Q ( θ |χprior , nprior ) : Q ( θ |χpost , npost ) ∝ exp ( npostθTχpost − npostA ( θ ) ) ( 3 ) where χpost = npriorχprior+ ∑N j u ( y ( j ) ) nprior+N and n post = nprior + N . We see that χprior ( resp . χpost ) can be viewed as the average sufficient statistics of nprior ( resp . npost ) fictitious samples ( Bishop , 2006 ) . Further , the average sufficient statistic of fictitious samples is equal to the expected sufficient statistic of the conjugate distribution , i.e . χ = EQ ( χ , n ) [ θ ] ( Brown , 1986 ; Diaconis & Ylvisaker , 1979 ) . Thus , the parameter χpost carries the inherent aleatoric uncertainty on the target distribution with natural parameters θ , while the evidence npost aligns well with the epistemic uncertainty ( i.e . a low evidence means few prior target observations ) . We stress that the natural conjugate prior parametrization χ , n is often different from the “ well-known ” parametrization κ used by standard coding libraries . By definition , a bijective mapping m ( κ ) = ( χ , n ) from the natural parametrization to the commonly used parametrization always exists ( see examples in Tab . 1 ) . Finally , exponential family distributions always admit a closed-form posterior predictive distribution ( Gelman et al. , 2004 ) .

Summary: This paper proposes a single-pass uncertainty estimation method for neural networks by predicting the update to the natural parameters of the conjugate prior to the likelihood of the data distribution and optimizing a corresponding ‘Bayesian loss’. In that it generalizes PostNet (Charpentier et al., 2020) beyond classification. The paper reports competitive or improved performance across a range of datasets and against both standard and recent baselines.