Function-Space Variational Inference for Deep Bayesian Classification

1 INTRODUCTION . Deep learning ( Goodfellow et al. , 2016 ) has enabled powerful classification models capable of working with complex data modalities and scaling to large data sets ( Krizhevsky et al. , 2012 ) , outperforming non-parametric models such as Gaussian processes ( Rasmussen & Williams , 2005 ) in these domains . The aim of Bayesian neural networks ( BNN ) is to provide these complex models with function priors for regularization , generalization and uncertainty quantification ( UQ ) useful in prediction tasks ( Gal , 2016 ; Wilson & Izmailov , 2020 ; Fortuin , 2021 ; Abdar et al. , 2021 ) . Predictive uncertainty is crucial for the deployment of machine learning systems in real-world settings , as it provides a degree of safety ( McAllister et al. , 2017 ) , trust ( Lim et al. , 2019 ) , sample efficiency ( Deisenroth & Rasmussen , 2011 ; Gal et al. , 2017 ) and human-in-the-loop cooperation ( Filos et al. , 2019 ) . Unfortunately , the highly redundant and nonlinear parameterization of modern deep learning systems has been shown to be challenging for Bayesian methods , sometimes resulting in inadequate uncertainty quantification ( Foong et al. , 2020 ; 2019 ; Osband et al. , 2018 ; Ovadia et al. , 2019 ; Wenzel et al. , 2020 ; Yao et al. , 2019 ) even on toy tasks ( e.g . Figure 3 ) . In this work , we consider function-space variational inference1 ( Sun et al. , 2019 ) ( fVI ) for classification . This inference method reasons about priors in function space , specifying the prior with respect to their predictive distribution rather than their weight distribution . By approximately representing the posterior in function space using a variational implicit process ( Ma et al. , 2019 ) , our approach can be used with several popular BNN representations such as deep ensembles ( Lakshminarayanan et al. , 2017 ) and Monte Carlo dropout ( Gal & Ghahramani , 2016 ) . By positing that the implicit predictive of a stochastic network classifier is Dirichlet distributed , rather than solely categorical , we gain a richer predictive uncertainty representation while retaining the same mean predictions as categorical classifiers . Furthermore , the approach provides a unifying perspective of prior work which use the Dirichlet distribution and function-space regularization ( Malinin & Gales , 2018 ; Malinin et al. , 2020 ; Joo et al. , 2020 ; Sensoy et al. , 2018 ; 2021 ) , through a principled objective that does not require a specific model or data transformations . We demonstrate that this method improves uncertainty quantification and adversarial robustness across a range of popular models and datasets for both small- and large-scale inference . 1While prior work uses the term functional VI , we use function-space VI for clarity . 1.1 CONTRIBUTIONS . We incorporate Dirichlet predictive priors into deep Bayesian networks using function-space inference . Concretely , our contributions are : 1 . We demonstrate an extension of variational implicit processes into the classification setting . 2 . We propose a function-space variational inference method for classification , using an effec- tive KL divergence estimator from softmax samples via Newton-based Dirichlet estimation . 2 RELATED WORK . We build on prior work that considers classification in the Bayesian setting and function-space inference of Bayesian neural networks . 2.1 BAYESIAN CLASSIFICATION . Compared to regression , classification is non-trivial for Bayesian methods due to the nonlinear link function required to predict the class labels . As a result , closed-form Bayesian models , such as Gaussian processes ( GP ) , require approximate inference methods such as the Laplace approximation ( Rasmussen & Williams , 2005 ) , variational inference ( Gibbs & Mackay , 2000 ; Hensman et al. , 2015 ; Salimbeni et al. , 2018 ; Izmailov et al. , 2018 ) , and expectation propagation ( Hernandez-Lobato & Hernandez-Lobato , 2016 ) . The Pólya-Gamma data augmentation trick ( Polson et al. , 2013 ) has enabled scalable closed-form variational training of sparse Gaussian process classifiers ( Wenzel et al. , 2019 ) . Gaussian processes have also been used with a Dirichlet predictive using a log-normal approximation ( Milios et al. , 2018 ) . Classification with Bayesian neural networks is possible through a wide range of approximate inference methods , including Markov chain Monte Carlo ( Neal , 1995 ; Zhang et al. , 2020 ) , ( mean-field ) variational inference ( MFVI ) ( Graves , 2011 ; Blundell et al. , 2015 ; Zhang et al. , 2018 ) , Laplace approximations ( MacKay , 1992 ; Denker & LeCun , 1991 ; Ritter et al. , 2018 ; Khan et al. , 2019 ; Immer et al. , 2021 ; Daxberger et al. , 2021 ) , ensembles ( Lakshminarayanan et al. , 2017 ; Osband et al. , 2018 ; Barber & Bishop , 1998 ; Pearce et al. , 2020 ) , expectation propagation ( Hernández-Lobato & Adams , 2015 ) and Monte Carlo dropout ( Gal & Ghahramani , 2016 ; Kingma et al. , 2015 ) . Radial BNNs ( Farquhar et al. , 2020 ) are motivated as a practical alternative to MFVI BNNs that uses Gaussian weight priors and posteriors . By sampling weights in a radial fashion , they avoid the pathologies encountered when sampling high-dimensional Gaussian distributions . Rank 1 BNNs ( Dusenberry et al. , 2020 ) combine ensembles and weight priors . Using the shared BatchEnsemble structure ( Wen et al. , 2020 ) and rank 1 covariance parameterizations , Rank 1 BNNs have a scalable memory requirement . Alternatively , the Laplace bridge ( Hobbhahn et al. , 2021 ) approximately maps a Dirichlet predictive density backwards through the softmax into a latent Gaussian predictive . A Gaussian-predictive BNN can then be trained using this latent approximation . Alternative methods avoid propagating uncertainties by predicting Dirichlet concentrations directly with neural networks . Prior networks ( Malinin & Gales , 2018 ; 2019 ) require categorical labels to be converted to Dirichlets , and mimic fVI as the objective consists of two KL divergences , for in- and outside the data distribution respectively . Another training method used prior networks to distill a trained ensemble into a single model ( Malinin et al. , 2020 ) . Similarly , belief matching ( Joo et al. , 2020 ) converts the labels to Dirichlets using Bayes rule . This method also echoes fVI , as the training objective is an fELBO where the index set is the training data . Another method converts the training labels to categorical probabilities and uses a Bayes risk objective with KL regularization against a function-space prior ( Sensoy et al. , 2018 ; 2021 ) . Compared to these methods , we use fVI as a principled approach to inference that allows us to use standard BNN parameterizations with the categorical likelihood , avoiding the need to design networks and data representations that facilitate this training approach . For a longer discussion on this prior work , see Appendix A . 2.2 FUNCTION-SPACE VARIATIONAL INFERENCE . Function-space inference is best understood through Gaussian processes , which are closed-form distributions over functions . Using this view , sparse GPs may be viewed as approximate inference over functions ( de G. Matthews et al. , 2016 ) , minimizing the fKL from its exact posterior via inducing points . The function-space variational BNN ( fBNN ) by Sun et al . ( 2019 ) uses the fKL to use explicit or implicit stochastic processes as function-space priors . They use a GP trained on the data as a prior , which can be viewed as a form of empirical Bayes . Similar approaches take a mirror descent view for batch training ( Shi et al. , 2019 ) . However , function-space variational inference can yield ill-defined objectives in many cases Burt et al . ( 2021 ) . Neural linear models have also been used with fVI , as they have an attractive closed-form Gaussian predictive ( Watson et al. , 2021a ; b ) . Concurrent work by Rudner et al . ( 2021 ) has also adopted fVI for classification , by linearizing the network about a Gaussian weight distribution to estimate the fKL . This model works with a Gaussian ( latent ) predictive prior and posterior which loses the intuitive aspect of function-space priors . Moreover , the linearization requires computation of the weight Jacobian during training , for which the memory requirement scales with the number of model parameters and outputs . Variational implicit processes ( Ma et al. , 2019 ) take a stochastic process perspective of stochastic parametric regressors and , for the regression setting with Gaussian likelihoods , proposes a wake-sleep procedure for inference . Wang et al . ( 2019 ) propose particle optimization methods , using finite representations of functions to learn a particle representation of the function-space posterior through the gradient flow of the log posterior . Function-space inference is also an attractive approach to continual learning , e.g . Pan et al . ( 2020 ) . 3 DIRICHLET PREDICTIVE PRIORS FOR BAYESIAN CLASSIFICATION . In this section , we discuss how function-space variational inference can be applied to classification using variational implicit processes ( VIP ) ( Ma et al. , 2019 ) . 3.1 VARIATIONAL IMPLICIT PROCESSES FOR CLASSIFICATION . LetD ⊂ ( X ×Y ) = { ( xn , yn ) } Nn=1 be the training data consisting of N observed pairs of input data xn and corresponding K-dimensional , one-hot class label vectors yn . Deep learning approaches to classification typically assume that a class label y follows a categorical distribution whose parameters are predicted by a neural network . To this end , a network φ with weights w defines a deterministic function f which maps an input x ∈ X to an element fx ∈ ∆K−1 , where ∆K−1 denotes the K−1 simplex and elements from ∆K−1 are K-dimensional probability vectors . More precisely , we write fx = g ( x , w ) = σ ( φ ( x ; w ) ) and y ∼ Cat ( ·|fx ) . The softmax function σ ensures that fx is a valid probability vector , i.e . that fxk ≥ 0 and ∑K k=1 fxk = 1 . In conventional ML training , the weightsw are optimized by maximizing log p ( D|f ) = log ∏ D Cat ( y|fx ) = ∑ D ∑K k=1 yk log fxk , where ∏ D and ∑ D denote ∏ ( x , y ) ∈D and ∑ ( x , y ) ∈D respectively , and φ , σ , andw are implicit in f . In the Bayesian setting , this categorical predictive is expected to express epistemic uncertainty in its prediction . For many BNN models , this amounts to sampling from a weight distribution p ( w ) and averaging the predictions to produce a posterior categorical predictive . We argue that this process throws away the epistemic uncertainty of the classifier , which is captured in M posterior predictive samples f ( 1 ) x , . . . , f ( M ) x . To capture this uncertainty , we study these implicit predictions and instead posit they are drawn from a posterior predictive Dirichlet distribution p ( fx ) . The Dirichlet density allows us to capture the variance in the categorical samples , as seen in Figure 2 . In general , the model becomes a generative mapping from X to the space of K-dimensional Dirichlet distributions . Sampling the weights and evaluating g ( x , · ) is equivalent to sampling the Dirichlet predictive , while varying x with fixed weight w corresponds to evaluating a classifier g ( · , w ) . Together , the joint variation ofw and x defines an implicit stochastic process which we assume describes an implicit K-dimensional Dirichlet distribution for each x ∈ X . When parameterized by a neural network , we can treat this stochastic process as a variational implicit process ( Example 3 , Ma et al . ( 2019 ) ) . Definition 3.1 ( Implicit stochastic processes , from Ma et al . ( 2019 ) ) . An implicit stochastic process ( IP ) is a set of random variables f , such that any finite subset fx1 : L = { fx1 , fx2 , ... , fxL } with L ∈ N has a joint distribution which is implicitly defined by the generative process , w ∼ p ( w ) , fxl = g ( xl , w ) , ∀xl ∈ X , 1 ≤ l ≤ L. Applying Definition 3.1 to stochastic network classifiers requires two key considerations . Firstly , the weight distribution p ( w ) takes a parameterized form , qθ ( w ) , with parameters θ we wish to optimize . The function g ( · , · ) represents a specific network architecture , such as a ResNet . Further details regarding the variations of qθ ( w ) that we consider can be found in Section D in the Appendix . Secondly , in the classification setting we are now reasoning over a joint distribution of Dirichlet distributions on a fixed K−1 simplex . Note this is distinct from the Dirichlet process ( Teh , 2010 ) , which is a trans-dimensional prior over discrete distributions , such that K is inferred . Informally , in our setting , a finite collection of L random variables would be an element from the K−1 simplex raised to the power of L. The stochastic process can be defined as a random variable from the K−1 simplex to the power of the data domain X . This is not a distribution we can reason about easily , in contrast to the Gaussian process where the joint and marginal distributions are all Gaussian ( Rasmussen & Williams , 2005 ) . However , as we only interact with the joint distribution implicitly , we avoid the complexity of representing it directly . Using this VIP perspective , we can use the stochastic networks as an implicit variational posterior over functions q ( f |θ ) , allowing us to adopt high-performance BNN representations while avoiding distributions over infinite parameter spaces , as in e.g . Bayesian nonparametrics ( Hjort et al. , 2010 ) . Further details regarding the formal stochastic process definition can be found in Section C in the Appendix . While modeling the predictive distribution as a Dirichlet captures the variation in the predictive samples , the inference procedure is not affected if we use weight-space priors . To use this rich predictive density to inform our inference strategy , we can design a Dirichlet predictive prior in function space and apply variational function-space inference to optimize qθ ( w ) .

Summary: The authors propose Function space variational inference for deep Bayesian classification, an inference framework which extends ideas around evidential learning for classification by applying function space Variational Inference and marginalizing over latent functions. Concretely, this leads to a method where the estimated Dirichlet is not directly regressed from a forward pass of a function, but is estimated as a parameter that would match the empirical distribution over sampled logits given multiple samples from the function posterior (using a nested optimization step), and hence can be plugged into the fVI framework. The authors show this is a flexible way to model uncertainty for classification and can be combined with diverse models (ensembles, MCDropout and various Bayesian models). The method produces good results.