@Article{JML-2-1,
author = {Susanna, Lange and Wei, Deng and Ye, Qiang and Lin, Guang},
title = {Batch Normalization Preconditioning for Stochastic Gradient Langevin Dynamics},
journal = {Journal of Machine Learning},
year = {2023},
volume = {2},
number = {1},
pages = {65--82},
abstract = {<p style="text-align: justify;">Stochastic gradient Langevin dynamics (SGLD) is a standard sampling technique for uncertainty
estimation in Bayesian neural networks. Past methods have shown improved convergence by including a preconditioning of SGLD based on RMSprop. This preconditioning serves to adapt to the local geometry of the
parameter space and improve the performance of deep neural networks. In this paper, we develop another
preconditioning technique to accelerate training and improve convergence by incorporating a recently developed batch normalization preconditioning (BNP), into our methods. BNP uses mini-batch statistics to improve
the conditioning of the Hessian of the loss function in traditional neural networks and thus improve convergence. We will show that applying BNP to SGLD will improve the conditioning of the Fisher information
matrix, which improves the convergence. We present the results of this method on three experiments including a simulation example, a contextual bandit example, and a residual network which show the improved
initial convergence provided by BNP, in addition to an improved condition number from this method.</p>},
issn = {2790-2048},
doi = {https://doi.org/10.4208/jml.220726a},
url = {https://global-sci.com/article/87596/batch-normalization-preconditioning-for-stochastic-gradient-langevin-dynamics}
}