Preconditioned Crank-Nicolson Algorithms for Wide Bayesian Neural Networks

Authors: lucia  pezzetti
stefano  favaro
collegio carlo alberto
stefano  peluchetti
Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-08-29 09:44:36

Abstract: Bayesian Neural Networks represent a fascinating confluence of deep learning techniques and probabilistic reasoning, offering a compelling framework for understanding uncertainty in complex predictive models. In this paper, we consider Bayesian Neural Networks with Gaussian initialization and we investigate the use of the preconditioned Crank-Nicolson algorithm to sample from the reparametrized posterior distribution of the weights as the width of the network grows. In addition to being robust in the infinite-dimensional setting, we prove that the acceptance probability of the preconditioned Crank-Nicolson sampler approaches 1 as the width of the network goes to infinity, independently of any stepsize tuning. We then compare how the efficiency of the Langevin Monte Carlo, the preconditioned Crank-Nicolson and the preconditioned Crank-Nicolson Langevin samplers are influenced by changes in the network width in some real-world cases. In particular, we demonstrate that in wide Bayesian Neural Networks  configurations, the proposed method allows for more efficient sampling, as evidenced by a higher effective sample size and improved diagnostic results compared with the Langevin Monte Carlo algorithm.

Main paper: 
 introduction: 
 Bayesian Neural Networks (BNNs) have emerged as a powerful framework for combining deep learning with probabilistic reasoning, offering a principled approach to understanding uncertainty in complex predictive models (1; 2; 3; 4; 5; 6; 7). Despite their potential advantages, BNNs face significant challenges, particularly in sampling from high-dimensional posterior distributions of network weights. As the width of neural networks increases, standard Markov Chain Monte Carlo (MCMC) methods often struggle with efficiency and scalability. In this paper, we give theoretical and empirical guarantees to address these challenges by leveraging function-space MCMC techniques, specifically the preconditioned Crank-Nicolson (pCN) algorithm (8) , for sampling from the posterior distributions of wide BNNs. Our method exploits recent theoretical advances in the understanding of wide neural networks and offers a robust sampling procedure that remains effective as network width increases. 
background and related works: 
 BNNs extend traditional neural networks by treating weights as random variables, allowing for the quantification of uncertainty in predictions. However, BNNs have reached far less popularity than their deterministic counterparts due to the high computational requirements and limited theoretical understanding. One of the major theoretical challenges concerns the comprehension of the parameterspace behavior of BNNs, in spite of the function one. Specifically, while it is established that under Gaussian initializations, the function distributions of wide BNNs converge to the Neural Network Gaussian Process (NNGP) limit (9; 10; 11; 12; 13; 14; 15) , the dynamics of the posterior distribution remains less understood, with few exceptions (16; 17) . We contribute to this literature by exploring sampling from the posterior distribution of wide BNNs, focusing on understanding its behavior and properties from a parameter-space perspective. 
proposed method: pcn for wide bnns: 
 Our approach leverages the pCN algorithm to sample from the posterior distribution of weights in wide BNNs. The key insight is to exploit the reparametrization proposed in (17) of the BNN weights that brings the posterior distribution closer to a standard Gaussian as the network width increases. Specifically, if we consider a BNN with L hidden layers and let d l be the width of the l -th layer, the reparametrization is defined as:EQUATIONwhere θ = [θ (l) ] l=0,...,L ∼ N (0, I D ) is the collection of flattened weights of the BNN and Σ and µ are data-dependent terms defined as:EQUATIONHere, Ψ represents the scaled output of the penultimate layer, and σ 2 is the observation variance.The convergence in the KL-divergence of the reparametrized posterior distribution to a standard Gaussian N (0, I D ) as the width goes to infinity suggests potential improvements in the mixing speed of MCMC procedures, compared to sampling from the notably arduous BNN posterior. Nevertheless, standard MCMC algorithms are notoriously ill-suited for the infinite-dimensional setting and must be carefully re-tuned as the dimension increases to avoid degeneracy in the acceptance probability (18; 19) . To address these challenges, we focus on the robust pCN method, specifically designed to perform reliably in infinite-dimensional spaces.The pCN algorithm (8) for sampling from the reparametrized posterior uses the following proposal:EQUATIONwhere β ∈ [0, 1) is a stepsize parameter. The acceptance probability for this proposal is given by:EQUATION)where ℓ is the log-likelihood.Our main theoretical contribution is the following theorem:Theorem 3.1 Consider the BNN model with the reparametrization 1. The acceptance probability of the pCN algorithm to sample from the reparametrized weight posterior, for any β ∈ [0, 1), converges to 1 as the width of the network increases.This result has profound implications for sampling from wide BNNs. It guarantees that as the network becomes wider, the pCN algorithm becomes increasingly efficient, with nearly all proposals being accepted. This contrasts sharply with traditional MCMC methods, which often require careful tuning of step sizes to maintain reasonable acceptance rates in high dimensions. The proof Theorem 3.1 relies on the convergence of the empirical Neural Network Gaussian Process (NNGP) kernel to a constant independent of the weights as the network width increases and is reported in Appendix A.  
empirical results: 
 To validate our theoretical findings and demonstrate the effectiveness of the pCN sampler for wide BNNs, we conducted a series of experiments on the CIFAR-10 dataset. We replicate the setting in (17) using a fully-connected neural network with one hidden layer and varying the width from 512 to 4096 neurons. The code is available at github.com/lucia-pezzetti/Function-Space-MCMC-for-Wide-BNNs.Figure 1 shows the acceptance rates for pCN, underdamped Langevin Monte Carlo (LMC), and preconditioned Crank-Nicolson Langevin (pCNL) (8) samplers as a function of network width. The results clearly demonstrate that the pCN acceptance rate steadily increases as the network width grows, approaching 1 for very wide networks. This empirically confirms our theoretical prediction in Theorem 3.1. In contrast, the LMC sampler shows a decline in acceptance rate for wider networks, highlighting the advantage of pCN in high-dimensional settings. The empirical results also showcase that the pCNL, leveraging both the dimensional-robustness of the pCN and the gradient-informed proposal of the LMC, reaches the most desirable performances. This suggests the possibility to adapt Theorem 3.1 also to the Langevin version of the pCN sampler.Figure 2 presents the per-step Effective Sample Size (ESS) (20) for different samplers across various network widths. Both the pCN and pCNL samplers show a consistent increase in ESS as the network width grows, indicating improved sampling efficiency. This is particularly evident for larger step sizes (β = 0.2 and β = 0.1), where they outperform the LMC for wider networks. The increasing ESS demonstrates that the samplers not only maintains a high acceptance rate but also produces less correlated samples in high-dimensional spaces. The results for the smallest stepsize confirm the necessity of avoiding degeneracy in the stepsize, as this introduces autocorrelation among the collected samples and leads to a deterioration in their quality. 
discussion and conclusion: 
 In this paper, we investigated the effectiveness of the pCN sampler in sampling the posterior distribution of wide BNNs. Our method leverages recent theoretical insights into the behavior of wide neural networks and addresses the challenges of sampling in high-dimensional spaces. The poor LMC performance reflects the fact that standard MCMC procedures are ill-posed in highdimensional settings. In contrast, the pCN and pCNL samplers demonstrate constant growth in ESS as the network width increases, indicating that enhancements in acceptance rate contribute positively to efficiency and performance. Finally, the smallest stepsize, β = 0.01, heavily affects the behavior of both algorithms, introducing high autocorrelation among the samples and affecting their quality.of our work include theoretical guarantees for the convergence of the pCN acceptance probability to 1 as network width increases, and their empirical validation through extensive experiments on real-world datasets.Our findings have significant implications for Bayesian deep learning, offering a scalable and robust approach to uncertainty quantification in large neural networks and contributing to the broader goal of combining the expressiveness of deep learning with the rigorous uncertainty quantification of Bayesian methods. Future directions for this work include extending the approach to more complex network architectures, such as convolutional and recurrent neural networks, and investigating theoretical guarantees for the pCNL sampler. 
 
 Appendix: Let the assumptions of Theorem 3.1 hold. We start by analysing the general expression of the MCMC acceptance probability: a = min 1, p(ϕ * |D)q(ϕ|ϕ * ) p(ϕ|D)q(ϕ * |ϕ) .We have already shown that q(ϕ * |ϕ) = N ( 1 -β 2 ϕ, β 2 I D ). Regarding the reparamatrized weight posterior of the network, we observe that ( 17)p(ϕ|D) = p(ϕ (L+1) |ϕ (≤L) , D)p(ϕ (≤L) |D) ∝ p(ϕ (L+1) |ϕ (≤L) , D) det(Σ)exp 1 2 y T (σ 2 I n + ΨΨ T ) -1 y where p(ϕ (L+1) |ϕ (≤L) , D) ∼ N (0, I d (L)) is assured by the reparametrisation. It is then crucial to recognize the empirical NNGP kernel Kσ 2 = σ 2 I n +ΨΨ T (21) and observe that det(Σ) ∝ det( Kσ 2 ).Inserting everything in the expression of the acceptance probability we have:p(ϕ * |D)q(ϕ|ϕ * ) p(ϕ|D)q(ϕ * |ϕ) = p(ϕ * (≤L) |D)p(ϕ * (L+1) |ϕ * (≤L) , D)q(ϕ|ϕ * ) p(ϕ (≤L) |D)p(ϕ (L+1) |ϕ (≤L) , D)q(ϕ * |ϕ) = p(ϕ * (≤L) ) det(Σ * )exp 1 2 y T (σ 2 I n + Ψ * Ψ * T ) -1 y p(ϕ * (L+1) |ϕ * (≤L) , D)q(ϕ|ϕ * ) p(ϕ (≤L) ) det(Σ)exp 1 2 y T (σ 2 I n + ΨΨ T ) -1 y p(ϕ (L+1) |ϕ (≤L) , D)q(ϕ * |ϕ)Where we denote with Σ * and Ψ * the covariance matrix and scaled input matrix of the redout layer in equation 2, but for a network with weights ϕ * . Now:q(ϕ|ϕ * ) ∝ exp - 1 2β 2 ||ϕ -1 -β 2 ϕ * || 2 = exp - 1 2β 2 ||ϕ|| 2 - (1 -β 2 ) 2β 2 ||ϕ * || 2 + 1 -β 2 β 2 ϕ T ϕ * = exp - 1 2β 2 ||ϕ|| 2 - 1 2β 2 ||ϕ * || 2 + 1 2 ||ϕ * || 2 + 1 -β 2 β 2 ϕ T ϕ * From which q(ϕ|ϕ * ) q(ϕ * |ϕ) = exp 1 2 ||ϕ * || 2 - 1 2 ||ϕ|| 2and sincep(ϕ * (≤L) )p(ϕ * (L+1) |ϕ * (≤L) , D) ∝ exp - 1 2 ||ϕ * (≤L) || 2 exp - 1 2 ||ϕ * (L+1) || 2 = exp - 1 2 ||ϕ * || 2 we obtain p(ϕ * |D)q(ϕ|ϕ * ) p(ϕ|D)q(ϕ * |ϕ) = exp -1 2 ||ϕ * || 2 exp 1 2 ||ϕ * || 2 det(Σ * )exp 1 2 y T (σ 2 I n + Ψ * Ψ * T ) -1 y exp -1 2 ||ϕ|| 2 exp 1 2 ||ϕ|| 2 det(Σ)exp 1 2 y T (σ 2 I n + ΨΨ T ) -1 y = det(Σ * )exp 1 2 y T (σ 2 I n + Ψ * Ψ * T ) -1 y det(Σ)exp 1 2 y T (σ 2 I n + ΨΨ T ) -1 y ∝ det( K * σ 2 )exp 1 2 y T ( K * σ 2 ) -1 y det( Kσ 2 )exp 1 2 y T ( Kσ 2 ) -1 yTo conclude, we exploit the known convergence of the empirical NNGP kernel to a constant independent ofϕ ≤L = θ ≤L Kσ 2 → K σ 2 as d min → ∞This proves that the numerator and the denominator converge to the same quantity and, consequently, that their ratio converges to 1.Implying the thesisa = 1 ∧ p(ϕ * |D)q(ϕ|ϕ * ) p(ϕ|D)q(ϕ * |ϕ) → 1 ∧ 1 = 1 for d min → ∞An alternative approach to Theorem 3.1 involves marginalizing the weights of the network's final layer and perform the sampling procedure only on the weights of the network's inner layers. This idea is effective since it acknowledges that exact sampling can be performed from the posterior distribution of the reparametrized weights of the last layer, once the weights from all preceding layers are known, i.e.,p(ϕ|D) = p(ϕ (L+1) |ϕ (≤L) , D)p(ϕ (≤L) |D) = p(ϕ (L+1) |θ (≤L) , D)p(θ (≤L) |D)where the last equality follows directly from the reparametrisation definition and from the fact that p(ϕ (L+1) |θ (≤L) , D) ∼ N (0, I d (L) ) for any fixed value of θ (≤L) . The idea is then to simply perform pCN sampling on the posterior distribution over the inner-layers weights π(θ ≤L |D). Then, once the samples θ (W |W * ) ∝ exp - 1 2β 2 ||W -1 -β 2 W * || 2 =⇒ q(W |W * ) q(W * |W ) = exp 1 2 ||W * || 2 - 1 2 ||W || 2Putting all together and simplifying we get: a = 1 ∧ π(W * |D)q(W |W * ) π(W |D)q(W * |W ) = 1 ∧ det(Σ * )exp 1 2 y T (σ 2 I n + Ψ * Ψ * T ) -1 y det(Σ)exp 1 2 y T (σ 2 I n + ΨΨ T ) -1 y That corresponds to the same exact expression found in Appendix A and hence convergence to 1 as the layers width increases is granted by the same arguments. 



Reviews: review: title: title: Interesting theoretical and empirical findings regarding pCN MCMC algorithms for wide BNNs
main_review: main_review: This paper presents interesting theoretical findings for preconditioned Crank-Nicholson algorithms for training wide Bayesian Neural Networks. They find that the acceptance rate and efficiency of this algorithms improves with the width of the BNNs. They tested this insight on numerical experiments on CIFAR-10.
title: title: Preconditioned Crank-Nicolson Algorithms for Wide Bayesian Neural Networks
main_review: main_review: This paper describes a new sampling algorithm for Bayesian Neural Networks. The key insight is that using a new reparameterization with the Crank-Nicolson sampling algorithm can lead to better acceptance rates in high dimensions. The paper is well-written and easy to follow, although most of the complexity is in the appendix. The theoretical and experimental data presented are convincing enough to warrant further investigation. My main criticism, which also reflects the early nature of the work, is that much more experimental work is needed to consider this sampling algorithm as a viable alternative, such as:

1) Comparison with more sophisticated sampling algorithms
2) Scaling up the experimental section with deeper (more than 1 layer) networks and more challenging tasks (including real-world problems)
3) Evaluating the downstream impact of the new parameterization, sampling, and higher acceptance rate on the quality of the posterior distribution, especially in high dimensions

However, I think this paper presents early evidence that this is an interesting direction to consider.