Layer Adaptive Node Selection in Bayesian Neural Networks: Statistical Guarantees and Implementation Details

Abstract

Sparse deep neural networks have proven to be efficient for predictive model building in large-scale studies. Although several works have studied theoretical and numerical properties of sparse neural architectures, they have primarily focused on the edge selection. Sparsity through edge selection might be intuitively appealing; however, it does not necessarily reduce the structural complexity of a network. Instead pruning excessive nodes in each layer leads to a structurally sparse network which would have lower computational complexity and memory footprint. We propose a Bayesian sparse solution using spike-and-slab Gaussian priors to allow for node selection during training. The use of spike-and-slab prior alleviates the need of an ad-hoc thresholding rule for pruning redundant nodes from a network. In addition, we adopt a variational Bayes approach to circumvent the computational challenges of traditional Markov Chain Monte Carlo (MCMC) implementation. In the context of node selection, we establish the fundamental result of variational posterior consistency together with the characterization of prior parameters. In contrast to the previous works, our theoretical development relaxes the assumptions of the equal number of nodes and uniform bounds on all network weights, thereby accommodating sparse networks with layer-dependent node structures or coefficient bounds. With a layer-wise characterization of prior inclusion probabilities, we also discuss optimal contraction rates of the variational posterior. Finally, we provide empirical evidence to substantiate that our theoretical work facilitates layer-wise optimal node recovery together with competitive predictive performance.

Sequential Bayesian Neural Subnetwork Ensembles

“…Baselines. Our baselines include the frequentist model of a deterministic deep neural network (trained with SGD), BNN [39], spike-and-slab BNN for node sparsity [30], single forward pass ensemble models including rank-1 BNN Gaussian ensemble [10], MIMO [11], and EDST ensemble [12], multiple forward pass ensemble methods: DST ensemble [12] and Dense ensemble of deterministic neural networks. For fair comparison, we keep the training hardware, environment, data augmentation, and training schedules of all the models same.…”

“…Dynamic sparsity learning for our sequential ensemble of sparse BNNs is achieved via spike-and-slab prior: a Dirac spike (δ 0 ) at 0 and a uniform slab distribution elsewhere [40]. We adopt the sparse BNN model of [30] to achieve the structural sparsity in Bayesian neural networks. Specifically a common indicator variable z is used for all the weights incident on a node which helps to prune away the given node while training.…”

“…The Evidence Lower Bound (ELBO) is represented as L = −E q(θ) [log p(D|θ)] + d KL (q(θ), p(θ)), where L in the case of SSBNN [30], which are base learners in our SeBayS ensembles, consists of KL of discrete variables leading to a non-differentiable loss and thus creates a challenge in practical implementation. Instead, [30] apply continuous relaxation, i.e., to approximate discrete random variables by a continuous distribution. Specifically, the continuous relaxation approximation is achieved through Gumbel-softmax (GS) distribution [49], [50], that is q(z lj ) ∼ Bernoulli(γ lj ) is approximated by q(z lj ) ∼ GS(γ lj , τ ), where…”

“…where τ is the temperature. We fix the τ = 0.5 for all the experiments, similar to [30]. zlj is used in the backward pass for easier gradient calculation, z lj is used for the exact sparsity in the forward pass.…”

“…First, the automatic data-driven sparsity learning in Bayesian neural networks is achieved using sparsity-inducing priors [27]. Second, the use of group sparsity priors [28][29][30] provides structural sparsity in Bayesian neural networks leading to significant computational gains. We leverage the automated structural sparsity learning using spike-and-slab priors similar to [30] in our approach to sequentially generate multiple Bayesian neural subnetworks with varying sparse connectivities which when combined yields highly diverse ensemble.…”