BayesAdapter: Being Bayesian, Inexpensively and Robustly, via Bayesian Fine-tuning

Despite their theoretical appealingness , Bayesian neural networks ( BNNs ) are falling far behind in terms of adoption in real-world applications compared with deterministic NNs , mainly due to their limited scalability in training and low fidelity in uncertainty estimates . In this work , we develop a new framework , named BayesAdapter , to address these issues and bring Bayesian deep learning to the masses . The core notion of BayesAdapter is to adapt pre-trained deterministic NNs to be BNNs via Bayesian fine-tuning . We implement Bayesian fine-tuning with a plug-and-play instantiation of stochastic variational inference , and propose exemplar reparameterization to reduce gradient variance and stabilize the finetuning . Together , they enable training BNNs as if one were training deterministic NNs with minimal added overheads . During Bayesian fine-tuning , we further propose an uncertainty regularization to supervise and calibrate the uncertainty quantification of learned BNNs at low cost . To empirically evaluate BayesAdapter , we conduct extensive experiments on a diverse set of challenging benchmarks , and observe satisfactory training efficiency , competitive predictive performance , and calibrated and faithful uncertainty estimates . 1 INTRODUCTION . Much effort has been devoted to developing flexible and efficient Bayesian deep models to make accurate , robust , and well-calibrated decisions ( MacKay , 1992 ; Neal , 1995 ; Graves , 2011 ; Blundell et al. , 2015 ) , with Bayesian neural networks ( BNNs ) as popular examples . The principled uncertainty quantification inside BNNs is critical for realistic decision-making , well evaluated in scenarios ranging from model-based reinforcement learning ( Depeweg et al. , 2016 ) and active learning ( Hernández-Lobato & Adams , 2015 ) , to healthcare ( Leibig et al. , 2017 ) and autonomous driving ( Kendall & Gal , 2017 ) . BNNs are also known to be capable of resisting over-fitting . However , there are fundamental obstacles posed in front of ML practitioners when trying to push the limit of BNNs to larger datasets and deeper architectures : ( i ) The scalability of the existing BNNs is generally restrictive owing to the essential difficulties of learning a complex , non-degenerate distribution over parameters in a high-dimensional and over-parameterized space ( Liu & Wang , 2016 ; Louizos & Welling , 2017 ; Sun et al. , 2019 ) . ( ii ) The Bayes posteriors learned from scratch are often systematically worse than their point-estimate counterparts in terms of predictive performance when “ cold posterior ” strategies are not applied ( Wenzel et al. , 2020 ) . ( iii ) It is shown that the BNNs have the possibility to assign low ( epistemic ) uncertainty for realistic out-of-distribution ( OOD ) data ( e.g. , adversarial examples ) , rendering their uncertainty estimates unreliable in safety-critical scenarios ( Grosse et al. , 2018 ) . To solve these problems , we present a scalable workflow , named BayesAdapter , to learn more reliable BNNs . In a holistic view , we unfold the learning of a BNN into two steps : deterministic pre-training of the deep neural network ( DNN ) counterpart of the BNN followed by several-round Bayesian fine-tuning . This enables us to learn a principled BNN with slightly more efforts than training a regular DNN , and provides us with the opportunities to embrace qualified off-the-shelf pre-trained DNNs ( e.g. , those on PyTorch Hub ) . The converged parameters of the deterministic model serve as a strong start point for Bayesian fine-tuning , allowing us to bypass extensive local optimum suffered by a direct learning of BNN1 . To render the fine-tuning in the style of training normal NNs , we resort to stochastic variational inference ( VI ) to update the approximate posterior . We develop optimizers with built-in weight decay for the parameters of the variational distribution to absorb the regularization effects from the prior , and develop exemplar reparametrization to reduce the gradient variance . Moreover , to make the uncertainty estimation of the learned models reliable , we propose to additionally , explicitly regularize the model to behave uncertainly on representative foreseeable OOD data during fine-tuning . This regularization takes the form of a margin loss , and is readily applicable to most of the existing BNNs . Figure 1 depicts the whole framework of BayesAdapter . Extensive empirical studies validate the efficiency and effectiveness of our workflow . In summary , our contributions are as follows : 1 . We propose BayesAdapter , to quickly and cheaply adapt a pre-trained DNN to be Bayesian without compromising performance when facing new tasks . 2 . We provide an easy-to-use instantiation of stochastic VI , which allows learning a BNN as if training a deterministic NN and frees the users from tedious details of BNN . 3 . We augment the fine-tuning with a generally applicable uncertainty regularization term to rectify the predictive uncertainty according to a collection of OOD data . 4 . Extensive studies validate that BayesAdapter is scalable ; the delivered BNN models are high-quality ; and the acquired uncertainty quantification is calibrated and transferable . 2 BAYESADAPTER . In this section , we first motivate BayesAdapter by drawing a connection between maximum a posteriori ( MAP ) and Bayesian inference . We then describe the proposed procedure Bayesian fine-tuning , and a practical and robust implementation of stochastic VI to realize it . Figure 1 illustrates the overall workflow of BayesAdapter . 2.1 FROM DNNS TO BNNS . Let D = { ( xi , yi ) } ni=1 be a given training set , where xi ∈ Rd and yi ∈ Y denote the input data and label , respectively . A DNN model can be fit via MAP as following : max w 1 n ∑ i [ log p ( yi|xi ; w ) ] + 1 n log p ( w ) . ( 1 ) We use w ∈ Rp to denote the high-dimensional model parameters , and p ( y|x ; w ) as the predictive distribution associated with the model . The prior term p ( w ) , when taking the form of an isotropic Gaussian , reduces to the common L2 weight decay regularizer in optimization . Despite the wide adoption , DNNs are known to be prone to over-fitting , generating over-confident predictions , and are unable to convey valuable information on the trustworthiness of their predictions . Naturally , Bayesian neural networks ( BNNs ) come into the picture to address these limitations . 1Here the BNN mainly refers to mean-field variational BNNs , and the results in Sec 4.1 testify this point . Typically , a BNN imposes a prior p ( w ) on model parameters , which is put together with the likelihood p ( D|w ) to infer the posterior p ( w|D ) . Among the wide spectrum of BNN algorithms ( MacKay , 1992 ; Neal , 1995 ; Graves , 2011 ; Blundell et al. , 2015 ; Liu & Wang , 2016 ; Gal & Ghahramani , 2016 ; Louizos & Welling , 2017 ) , variational BNNs are particularly promising due to their ease of training compared with other BNN variants . Formally , variational BNNs derive a θ-parameterized varitional distribution q ( w|θ ) to approximate the true posterior p ( w|D ) , by maximizing the evidence lower bound ( ELBO ) ( scaled by 1/n ) : max θ Eq ( w|θ ) [ 1 n ∑ i log p ( yi|xi ; w ) ] ︸ ︷︷ ︸ Lell − 1 n DKL ( q ( w|θ ) ‖p ( w ) ) ︸ ︷︷ ︸ Lc , ( 2 ) whereLell is the expected log-likelihood andLc is the complexity loss . By casting posterior inference into optimization , Eq . ( 2 ) makes the training of BNNs more approachable . However , most existing BNNs2 trained under such a criterion exhibit limitations in scalability and performance ( Osawa et al. , 2019a ; Wenzel et al. , 2020 ) compared with their deterministic counterparts , mainly attributed to the higher difficulty of learning high-dimensional distributions than point estimates , and challenges in finding non-degenerated optima of highly nonlinear functions characterized by NNs . Given that MAP converges to the mode of the Bayesian posterior , it might be plausible to adapt pretrained deterministic DNNs to be Bayesian economically . Following this hypothesis , we propose to repurpose the converged parameters w∗ of MAP , and use it to instantiate q ( w|θ ) as a Gaussian N ( w ; θ ) with θ = ( µ , Σ ) , where µ is initialized as w∗ and Σ ∈ Rp×p denotes the covariance . Then , we arrive at a BNN with posterior predictive : p ( y|x , D ) = EN ( w ; µ , Σ ) p ( y|x ; w ) ≈ 1 S S∑ s=1 p ( y|x ; w ( s ) ) , where w ( s ) ∼ N ( w ; µ , Σ ) , s = 1 , ... , S. ( 3 ) Eq . ( 3 ) is also called Bayes ensemble , where µ is perturbed , and the predictions from multiple likely models are assembled . Σ controls the magnitude of perturbation . A classic method to generate an informative Σ is by Laplace approximation ( Bleistein & Handelsman , 1986 ) , but it is more like a postprocessing procedure , lacking the flexibility to jointly adapt the mean and covariance of the Gaussian posterior w.r.t . data , and its naive implementation without strong assumptions may be computationally prohibitive . Instead , we suggest a more practical workflow – that fine-tunes the approximate posterior N ( w ; µ , Σ ) by maximizing the ELBO with randomly initialized Σ . 2.2 BAYESIAN FINE-TUNING IN THE STYLE OF FINE-TUNING DNNS . We develop practical learning algorithms under the stochastic VI scheme to fine-tune the imperfect variational posterior , and to cope with contemporary ML frameworks . In the following , we discuss how to deal with each term in Eq . ( 2 ) . Algorithm 1 gives an overview of BayesAdapter . Complexity loss Lc . Without losing generality , we assume an isotropic Gaussian prior p ( w ) = N ( w ; 0 , σ20I ) . Then the complexity loss is derived as : Lc = − 1 n DKL ( N ( w ; µ , Σ ) ‖N ( w ; 0 , σ20I ) ) = −µ Tµ+ tr ( Σ ) 2σ20n + log detΣ 2n + c , ( 4 ) where tr and det are matrix trace and determinant , respectively . c is a constant . The gradients of Lc w.r.t . µ and Σ can be estimated precisely as : ∇µLc = − µ σ20n , ∇ΣLc = σ20Σ −1 − I 2σ20n . ( 5 ) Eq . ( 5 ) indicates that maxµ Lc amounts to applying a weight decay regularizer with coefficient λ = 1 σ20n on µ , which can be conveniently optimized by leveraging the built-in weight decay modules in ML frameworks such as TensorFlow ( Abadi et al. , 2016 ) or PyTorch ( Paszke et al. , 2019 ) . Directly computing ∇ΣLc involves matrix inversion . Implementing the posterior as matrix-variate Gaussian is an alternative , while existing algorithms for matrix-variate Gaussian posterior typically exhibit high complexity in time or memory , limited compatibility with contemporary NN building block operations ( e.g. , convolution ) , and struggle to scale with data-parallel distributed training ( Louizos & Welling , 2016 ; Sun et al. , 2017 ; Osawa et al. , 2019b ) . To simplify the implementation and boost scalability , we assume a fully factorized Gaussian variational by devising Σ as 2We use BNNs equivalently with variational BNNs in the following text when there is no ambiguity . diag ( exp ( 2ψ ) ) , where ψ ∈ Rp is the parameter to be optimized along with µ ( i.e. , θ = ( µ , ψ ) ) . Injecting this into Eq . ( 5 ) gets a more concise gradient estimator : ∇ψLc = 1/n−λ exp ( 2ψ ) , meaning that maxψ Lc adds an exponential weight decay of ψ with coefficient λ , which can be realized by modifying only two lines of code on top of de facto DL frameworks ( see Figure 1 ) . Expected log-likelihood Lell . With the complexity loss expressed as weight decay , we now develop efficient ways for calculating the Lell at the end of forward pass , and for performing backpropagation afterwards . In particular , we derive a Monte Carlo ( MC ) estimation of Lell based on reparameterization ( Kingma & Welling , 2013 ) : we sample a p-dimensional Gaussian noise ∼ N ( 0 , I ) , then obtain the sampled parameter for the whole mini-batch B of data viaw = µ+exp ( ψ ) , given which we approximate Lell with L′ell = 1|B| ∑ ( xi , yi ) ∈B log p ( yi|xi ; w ) . The gradients of µ and ψ can be derived automatically with autodiff libraries , thus the training resembles that of normal DNNs . However , gradients derived by L′ell might exhibit high variance , caused by sharing one set of sampled parameters w across all the training instances in B . Local reparameterization is proposed to reduce the variance , but it requires at least 2x forward-backward FLOPS than vanilla reparameterization ( refer to Kingma et al . ( 2015 ) for more details ) . Flipout ( Wen et al. , 2018 ) is an alternative solution . But it is only suitable for perturbation based MC estimation and its modeling assumptions make Flipout unable to handle complex variational posterior like a FLOW ( Louizos & Welling , 2017 ) , or an implicit model ( Shi et al. , 2018b ) . Besides , it is still as slow as local reparameterization . To mitigate these issues , we propose exemplar reparametrization ( ER ) which samples a separate set of parameters for every exemplar in the minibatch . Formally , for ∀xi ∈ B , we draw w ( i ) = µ + exp ( ψ ) ( i ) where ( i ) ∼ N ( 0 , I ) , and approximate the expected log-likelihood by L∗ell = 1|B| ∑ ( xi , yi ) ∈B log p ( yi|xi ; w ( i ) ) . Obviously , ER is distribution agnostic , and is readily applicable to various variational distributions . While ER generates more parameters at training , they are mostly temporary , and the resultant computational FLOPS are provably identical to that of the vanilla reparameterization . The challenge of ER is to cope with nowadays ML frameworks and maintain computing efficiency , because off-the-shelf computation kernels in autodiff libraries typically assume a batch of instances share a common set of parameters . We present an example in Fig- ure 2 on how the standard convolution op can be converted into its exemplar version without compromising computational efficiency . The key insight here is that multiple exemplar convolutions can be expressed as a group convolution , which can be performed in parallel using a single group convolution kernel , leveraging the optimized implementations provided by various device-propriety kernel backends ( e.g . cuDNN ( Chetlur et al. , 2014 ) ) . Other common operators such as matrix multiplication are straightforward to handle ( refer to Appendix A ) . With this insight , BayesAdapter enables to obtain a BNN with only minor computational cost in addition to pre-training , and can immediately benefit from the availability of higher-performance computational kernels ( e.g. , more powerful group convolution kernel ) .

Summary: This paper introduces a fast way to get Bayesian posterior by using a pretrained deterministic model. Specifically, the authors first train a standard DNN model and then use it to initialize the variational parameters. Finally the variational parameters are optimized through standard variational inference (VI) training. To further improve uncertainty estimate, the authors propose an uncertainty regularization which maximizes the prediction inconsistency on out-of-distribution (OOD) data. Experiments including image classification and uncertainty estimates are conducted to demonstrate the proposed method.