An adaptive Hessian approximated stochastic gradient MCMC method

Journal Article · Thu Feb 04 00:00:00 EST 2021 · Journal of Computational Physics

DOI:https://doi.org/10.1016/j.jcp.2021.110150· OSTI ID:1853727

Wang, Yating ^[1]; Deng, Wei ^[1]; Lin, Guang ^[2]

---

1. Purdue Univ., West Lafayette, IN (United States). Dept. of Mathematics
2. Purdue Univ., West Lafayette, IN (United States). Dept. of Mathematics. School of Mechanical Engineering. Dept. of Statistics. Dept. of Earth, Atmospheric, and Planetary Sciences

Bayesian approaches have been successfully integrated into training deep neural networks. One popular family is stochastic gradient Markov chain Monte Carlo methods (SG-MCMC), which have gained increasing interest due to their ability to handle large datasets and the potential to avoid overfitting. Although standard SG-MCMC methods have shown great performance in a variety of problems, they may be inefficient when the random variables in the target posterior densities have scale differences or are highly correlated. Here, we present an adaptive Hessian approximated stochastic gradient MCMC method to incorporate local geometric information while sampling from the posterior. The idea is to apply stochastic approximation (SA) to sequentially update a preconditioning matrix at each iteration. The preconditioner possesses second-order information and can guide the random walk of a sampler efficiently. Instead of computing and saving the full Hessian of the log posterior, we use limited memory of the samples and their stochastic gradients to approximate the inverse Hessian-vector multiplication in the updating formula. Moreover, by smoothly optimizing the preconditioning matrix via SA, our proposed algorithm can asymptotically converge to the target distribution with a controllable bias under mild conditions. To reduce the training and testing computational burden, we adopt a magnitude-based weight pruning method to enforce the sparsity of the network. Our method is user-friendly and demonstrates better learning results compared to standard SG-MCMC updating rules. The approximation of inverse Hessian alleviates storage and computational complexities for large dimensional models. Numerical experiments are performed on several problems, including sampling from 2D correlated distribution, synthetic regression problems, and learning the numerical solutions of heterogeneous elliptic PDE. The numerical results demonstrate great improvement in both the convergence rate and accuracy.

• Research Organization: Purdue Univ., West Lafayette, IN (United States)
• Sponsoring Organization: National Science Foundation (NSF); US Army Research Office (ARO); USDOE Office of Science (SC), Advanced Scientific Computing Research (ASCR)
• Grant/Contract Number: SC0021142; DMS-1555072; DMS-1736364; CMMI-1634832; CMMI-1560834; W911NF-15-1-0562
• OSTI ID: 1853727
• Alternate ID(s): OSTI ID: 1775932
• Journal Information: Journal of Computational Physics, Vol. 432, Issue C; ISSN 0021-9991
• Publisher: Elsevier
• Country of Publication: United States
• Language: English

Related Subjects

97 MATHEMATICS AND COMPUTING
adaptive Bayesian method
deep learning
Hessian approximated stochastic gradient MCMC
stochastic approximation
limited memory BFGS
highly correlated density