Bayesian Sparse learning with preconditioned stochastic gradient MCMC and its applications

TL;DR

This paper introduces a Bayesian sparse deep learning algorithm using preconditioned stochastic gradient MCMC, which improves convergence and efficiency in non-convex models by adapting to local geometry and optimizing hyperparameters.

Contribution

The paper proposes a novel Bayesian sparse learning method with preconditioned SGLD that addresses slow mixing issues and ensures asymptotic convergence in complex neural network models.

Findings

Demonstrates improved accuracy on synthetic regression tasks.

Shows efficient learning of solutions to elliptic PDEs.

Provides theoretical convergence guarantees.

Abstract

In this work, we propose a Bayesian type sparse deep learning algorithm. The algorithm utilizes a set of spike-and-slab priors for the parameters in the deep neural network. The hierarchical Bayesian mixture will be trained using an adaptive empirical method. That is, one will alternatively sample from the posterior using preconditioned stochastic gradient Langevin Dynamics (PSGLD), and optimize the latent variables via stochastic approximation. The sparsity of the network is achieved while optimizing the hyperparameters with adaptive searching and penalizing. A popular SG-MCMC approach is Stochastic gradient Langevin dynamics (SGLD). However, considering the complex geometry in the model parameter space in non-convex learning, updating parameters using a universal step size in each component as in SGLD may cause slow mixing. To address this issue, we apply a computationally manageable preconditioner in the updating rule, which provides a step-size parameter to adapt to local geometric properties. Moreover, by smoothly optimizing the hyperparameter in the preconditioning matrix, our proposed algorithm ensures a decreasing bias, which is introduced by ignoring the correction term in preconditioned SGLD. According to the existing theoretical framework, we show that the proposed algorithm can asymptotically converge to the correct distribution with a controllable bias under mild conditions. Numerical tests are performed on both synthetic regression problems and learning the solutions of elliptic PDE, which demonstrate the accuracy and efficiency of present work.