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: The paper introduces a new method Natural Posterior Network (NatPN), an extension to the Posterior Network (Charpentier et al., 2020). NatPN aims to tackle the problem of enabling calibrated uncertainty estimates for in and out-of-distribution inputs for exponential family likelihoods parameterized by deep neural networks. In particular NatPN maps input samples into a latent space on which a normalizing flow is defined. A mapping from this latent space and its corresponding likelihood under the normalizing flow to the parameters of exponential family update rule is defined. The exponential family posterior tends towards the prior parameters for data points far from the training data and prior has little influence for data points within the training distribution.  NatPN makes two significant changes to the Posterior Network 1) it generalizes the method to the exponentially family, allowing for regression and counting tasks beyond just classification as per the Posterior Network and 2) it uses a single normalizing flow rather than class specific normalizing flows as per the Posterior Network, in theory this aids the scalability of the method to datasets with many classes. In addition a minor change is made to the update rule of the posterior parameters.  The paper presents extensive empirical results on classification, regression and counting tasks across several relatively small scale datasets. NatPN is shown to perform on par or slightly better than similar methods from the literature.