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 Natural Posterior Networks (NatPNs), a technique to provide uncertainty-estimation to tasks where the likelihood is an exponential family, e.g. classification or regression. NatPNs leverage the properties of exponential families, i.e. conjugacy and closed-form posterior predictive for fast and elegant Bayesian Inference. This is combined with a Bayesian treatment of the loss function that allows for end-to-end training. Furthermore, it uses a normalizing flow to account for the epistemic uncertainty. The normalizing flow density is supposed to increase the evidence on the training data and thereby decrease it everywhere else. The posterior distribution is on the outputs not on the weights of the NN which makes it easier to train. The paper provides one theorem for why the posterior uncertainty is meaningful and multiple experiments showing practical improvements for multiple exponential families and datasets.    NatPNs are an extension of Posterior Networks. However, the differences are significant and well explained in the paper. From my perspective, the contributions are clearly novel enough to justify this new paper.