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 paper introduces a function-space variational inference algorithm ("fVI") that uses a Dirichlet predictive prior, and approximates the output of a (weight-space-Bayesian) neural network as a Dirichlet distribution. Inference is achieved by viewing the network as a variational impolicit process, and optimising the functional ELBO from Sun et al., 2019. This view allows them to impose priors in function-space to many existing algorithms, such as MAP, MC Dropout, ensembles, Radial BNNs, and so on. The authors show results on a toy problem, and focus on out-of-distribution performance on larger problems such as MNIST, CIFAR10 and CIFAR100, as well as performance under a fast gradient sign method adversarial attack. The fVI variants almost always perform better when the test distribution is not the same as the training distribution.