Edge-preserving Bayesian neural networks via hierarchical sparsity-promoting priors

Abstract

Recovering piecewise-smooth functions with jump discontinuities from noisy, indirect, and incomplete data is a central task in computational inverse problems. Bayesian neural networks (BNNs) with heavy-tailed weight priors have emerged as effective edge-preserving priors for such problems. However, posterior inference for these models remains computationally challenging, and existing approaches settle for point estimates obtained by directly optimizing heavy-tailed log-densities. We propose GSBL-BNNs, which model the neural network parameters hierarchically using sparsity-promoting priors from generalized sparse Bayesian learning (GSBL). The hierarchical formulation recovers several popular marginal priors, including Laplace and Student's $t$ (including Cauchy). A key modeling ingredient, novel among edge-preserving BNNs and native to the hierarchical formulation, couples each neuron's bias to its weights through a shared hyperparameter, so that the induced jump discontinuities are free to occur anywhere in the domain while the conditionally Gaussian structure underlying efficient inference is preserved. Furthermore, we develop a block-coordinate descent (BCD) algorithm for maximum a posteriori (MAP) estimation. This BCD approach alternates between standard network training with iteratively reweighted $\ell^2$-regularization and exact, closed-form hyperparameter updates for gamma and inverse gamma hyperpriors. Building on the BCD algorithm, we perform approximate Bayesian inference and uncertainty quantification (UQ) using deep ensembles. Computational experiments---including signal denoising, deblurring, and inverting a two-dimensional hyperbolic conservation law--- show that GSBL-BNNs deliver accurate, edge-preserving reconstructions with built-in UQ. Compared with Gaussian- and Cauchy-BNNs, they yield more accurate predictive means and avoid overconfident uncertainty estimates.