Source-linked AI summary
Sparse Bayesian Neural Networks via Sticky Piecewise Deterministic Markov Processes
August Arnstad, Geir Olve Storvik, Filippo Pagani, Leiv Rønneberg
TL;DR
BNN posterior sampling is difficult, and dense weights can remain uninformative or retain prior variance. This paper uses sticky PDMPs to sample spike-and-slab BNN posteriors with exact zeros when numerical rate bounds hold. The authors report high sparsity, modest predictive-performance costs, and greater robustness to Gaussian input noise than dense samplers.
Problem
BNN posterior sampling is difficult, and weights uninformative to the data can retain prior variance while dense networks are overparameterized.
Source: §1 · p. 1Method
Sticky and non-sticky PDMPs sample BNN posteriors, with sticky versions incorporating spike-and-slab priors to set weights exactly to zero during sampling.
Source: §1 · p. 2Results
Sticky PDMPs set many weights exactly to zero at modest predictive-performance cost, and sparse BNNs can be more robust to Gaussian input noise than dense ones.
Source: §6 · pp. 9–10Takeaways & Limitations
Exact-zero sampling yields sparse BNNs without thresholding or post-processing, while retaining predictive performance close to dense samplers.
Source: §5.4 · p. 9, §6 · p. 9Exactness depends on rate bounds holding throughout each window; violations between proposed event times may go undetected.
Source: §4.2 · p. 5
Abstract
from arXiv · showhide
There is growing evidence that the performance of a neural network depends largely on a small subset of its weights, which are typically identified by training a dense network and pruning it afterwards. Within a Bayesian framework, pruning can be directly incorporated through a spike-and-slab prior, providing a proper treatment of the uncertainties involved, both with respect to pruning and to parameter uncertainty. In Bayesian neural networks (BNNs), sampling by MCMC is notoriously hard. We demonstrate that modern Monte Carlo methods based on piecewise deterministic Markov processes (PDMPs) can be used to sample posteriors of BNNs asymptotically exact whenever numerically constructed rate bounds hold. Spike-and-slab priors can be incorporated through sticky versions of the PDMPs, setting weights exactly to zero during sampling. We examine the predictive performance and convergence of the samplers, and show that sparse BNNs sampled this way can be more robust to input noise than dense ones.
1 Introduction
Bayesian neural networks are difficult to sample because exact methods remain costly, while dense, overparameterized weights can be weakly identified and poorly suited to spike-and-slab priors. This paper adapts sticky PDMPs to sample BNN posteriors with exact zeros, subject to valid numerical rate bounds.
Motivation
- Asymptotically exact BNN samplers target the posterior without bias but remain too costly for practical use, while minibatch stochastic-gradient samplers introduce bias.Variational inference instead replaces the posterior with a tractable distribution.Source: §1 · p. 1
- Weights uninformative to the data can retain prior variance, while spike-and-slab priors assign positive probability to exact zeros but are difficult for gradient-based samplers.Pruning can substantially reduce BNN weights with little predictive-performance loss, motivating direct sparse Bayesian inference.Source: §1 · p. 1
Background
- Sticky PDMPs target spike-and-slab posteriors directly, producing sparsity during sampling without post-processing or thresholding.Figure 1 illustrates a sticky ZigZag path and posterior inclusion probabilities for BNN weights and MNIST input pixels.Source: §1 · p. 2, Figure 1 · p. 2
Contributions
- Sticky and non-sticky PDMPs are adapted to BNNs with cheap numerical bounds, and full-batch sampling is asymptotically exact when those bounds hold.The paper also provides a GPU-accelerated PyTorch implementation and evaluates the samplers against NUTS and a sparse variational method.Source: §1 · p. 2
- The paper reports that exact zeros can yield high sparsity and improve robustness to input noise.It examines predictive performance and convergence of the samplers.Source: §1 · p. 2
2 Related work
Related work covers Bayesian neural-network priors and sampling methods, including the difficulty of rate bounds for PDMPs and prior sticky-PDMP applications. The paper positions its work among approaches to sparsification and continuous-time sampling.
Bayesian neural-network inference
- Hamiltonian Monte Carlo remains a gold standard for asymptotically exact inference, while variational inference has made faster approximate techniques attractive for modern neural networks.The cited work describes the scale of modern neural networks as motivating this shift.Source: §2 · p. 2
Bayesian neural-network priors
- Spike-and-slab priors assign positive probability to exact zeros, whereas global-local priors such as the horseshoe require a pruning threshold after sampling.A posterior-concentration result for deep networks under spike-and-slab priors is reported, but the cited work proposes no sampler for that posterior.Source: §2 · p. 2
PDMP sampling
- The bouncy particle sampler, ZigZag, and Boomerang are established PDMP samplers, with the latter two following the BPS in the cited history.The cited passages identify their origins and references.Source: §2 · pp. 2–3
- PDMP event-time simulation usually requires an upper rate bound; prior approaches use analytic decompositions, numerical windows, adaptive grids, or discretization with Metropolis correction.The cited BNN implementation uses a different bounding strategy.Source: §2 · p. 3
Sticky PDMPs
- Sticky PDMPs allow parameters to remain exactly zero and make non-reversible moves between sub-models; prior continuous-time applications required tractable rate bounds.Stochastic-gradient sticky ZigZag instead discretizes time, incurring bias of the order of the step size.Source: §2 · p. 3
3 Bayesian neural networks
The paper formulates Bayesian neural-network prediction through a posterior over weights given training data. Because the predictive integral is unavailable in closed form, posterior samples support Monte Carlo prediction; PDMPs provide a continuous-time alternative to discrete-time MCMC.
Bayesian neural networks
- For a network f(x; β) with weights and biases β, the target posterior is proportional to the data likelihood times the parameter prior.The training data are denoted D, and the likelihood is p(D|β).Source: §3 · p. 3
- Because the predictive integral is unavailable in closed form for neural networks, posterior samples can be used for a Monte Carlo estimate of the predictive distribution.The integral averages the data predictive distribution over the posterior of the parameters.Source: §3 · p. 3
- The paper considers continuous-time PDMP samplers as an alternative to discrete-time MCMC, noting that they extend naturally to spike-and-slab priors.Source: §3 · p. 3
4 Piecewise deterministic Markov processes
PDMP samplers target a BNN posterior through deterministic motion interrupted by random events, while sticky variants let weights remain exactly zero. The section describes event simulation, practical rate bounds, and the condition under which sampling remains exact.
4 Piecewise deterministic Markov processes
- PDMPs target the posterior as the position marginal of an extended distribution, evolving deterministically between events governed by a rate and transition kernel.ZigZag events select a coordinate in proportion to its rate and reverse that coordinate’s velocity.Source: §4 · pp. 3–4
- Because event rates vary along trajectories, implementations use Poisson thinning with a tractable upper bound and accept proposals according to the rate-to-bound ratio.For sticky samplers, the simulation window ends at the earliest of the bound horizon, a freeze time, or a thaw clock.Source: §4 · p. 4, §4.1 · p. 4, §4.2 · p. 5
4.1 Sticky PDMPs
- Sticky PDMPs handle spike-and-slab targets by freezing a coordinate at zero for an exponential time, then resuming motion with its previous velocity.The fraction of time a coordinate stays frozen estimates its posterior probability of being zero.Source: §4.1 · p. 4
4.2 Piecewise constant bounds
- A single-segment tangent bound uses endpoint rate and derivative evaluations; compared with a multi-segment scheme, it reduces rate-evaluation costs for BNNs.Building a window costs two evaluations, or one when the previous window ended without an event; one UCI model saw a 1.6–2.5-fold reduction in cost per event.Source: §4.2 · p. 5
- Sampling is asymptotically exact only if the numerical rate bound holds throughout each window, although violations are detected only at proposed event times.A rate may switch between concave and convex within a window, and a violation elsewhere can go unnoticed.Source: §4.2 · p. 5
- The implementation accepts a PyTorch module and parameter prior, supports GPU acceleration, and reports detected bound violations for user inspection.The user need not derive event-rate bounds for a new model.Source: §4.2 · pp. 5–6
5 Experiments
Experiments compare dense and sticky PDMPs with NUTS and a sparse variational method across synthetic regression, UCI benchmarks, convergence tests, and MNIST. Sticky sampling produces exact zeros, with predictive and convergence behavior varying by setting.
5 Experiments
- The experiments compare dense PDMPs and NUTS targeting the usual posterior with sticky samplers targeting a spike-and-slab posterior, alongside LBBNN; regression is evaluated with RMSE, NLL, and CRPS.The study spans one-dimensional regression, Boston and other UCI datasets, convergence analysis, and MNIST image classification.Source: §5 · p. 6
5.1 A one-dimensional non-linear example
- On the sharp dataset, sticky samplers reach test RMSE 0.109 versus 0.115–0.116 for dense samplers, while recovering flat, saturating regions more closely.All samplers capture the bump, with its peak slightly underestimated.Source: §5.1 · p. 6, Table 1 · p. 7, Figure 2 · p. 7
5.2 Benchmark datasets
- On Boston, sticky samplers set about two-thirds of weights to zero at w = 0.3 and 85–89% at w = 0.1, with small predictive differences from dense samplers on the small network.The LBBNN trails all samplers in RMSE and CRPS, even at matched sparsity.Source: §5.2 · p. 7, Table 2 · p. 8
5.3 Convergence analysis
- On the small network, all samplers converge; medium-network ZigZag chains agree but mix slowly, while Boomerang chains do not converge.The authors attribute Boomerang’s behavior to a narrow Gaussian reference around overfitting MAP estimates; NUTS converges on both networks and is more effective per gradient evaluation.Source: §5.3 · p. 8, Figure 3 · p. 9
5.4 Image data
- On MNIST, all samplers stay within 0.8 percentage points of SGD accuracy, while sticky ZigZag sets 92–98% of weights exactly to zero.Under unclipped Gaussian noise, sticky ZigZag is more robust than dense ZigZag but less robust than SGD and both MAPs.Source: §5.4 · p. 9, Figure 4 · p. 9
6 Discussion
The discussion reports that sticky PDMPs produce exactly sparse BNNs with modest predictive-performance cost, and can improve robustness to Gaussian input noise. It also emphasizes constraints on convergence and exactness, alongside directions for further work and disclosures about AI use.
6 Discussion
- Sticky PDMPs set many weights exactly to zero without post-processing, with modest predictive-performance cost and better predictions than a sparse variational approximation under the same prior.The paper also reports improved robustness to Gaussian input noise relative to a dense BNN.Source: §6 · p. 9
- A cheap tangent-segment rate bound is valid when the rate does not switch between concave and convex within the adaptive window; MNIST results show greater robustness to Gaussian input noise than a dense BNN.UCI experiments examine convergence across chains initialized at different MAPs.Source: §6 · p. 10
- On the first Boston split, dense ZigZag mixes about 20 times more slowly than NUTS on the small network and 50 times more slowly on the medium network per gradient evaluation, measured by effective sample size of log σ.The comparison is approximate; the authors also report that exactness depends on unverifiable bounds for a BNN, MNIST runs are single minibatches not expected to converge, and chains explore locally from a MAP.Source: §6 · p. 10
- Proposed directions include adapting ZigZag velocities, valid subsampled-gradient bounds, an adaptive Boomerang reference, sparser priors or architectures, and scaling to larger networks.The authors also call for Python software comparable to NumPyro's HMC tooling.Source: §6 · p. 10
AI use statement
- The authors report using generative AI for feedback, implementation, analysis, literature work, figures, code, and drafting, and say they reviewed AI-assisted work and verified methods, code, sources, and text.They state that AI did not generate synthetic datasets, design the methodology, develop theoretical models, or provide critical ingredients for proofs.Source: AI use statement · p. 10
A Obtaining samples from the continuous trajectory
The section estimates trajectory expectations after discarding an initial burn-in, using uniformly sampled times to approximate averages. Reservoir sampling lets the same sampling scheme work when large trajectories are processed in segments and their total duration is not yet known.
A Obtaining samples from the continuous trajectory
- The first 20% of trajectory time is discarded as burn-in, and T0 marks the start of the retained trajectory used to estimate the expectation.The estimator averages f(β_t) over [T0, T].Source: A Obtaining samples from the… · p. 14
- Uniformly sampled times on [T0, T] yield an unbiased sample-average estimate of the trajectory estimator, unlike a grid-based Riemann-sum approximation.In the experiments, times are drawn independently and uniformly, then sorted into time order.Source: A Obtaining samples from the… · p. 14
- A uniform trajectory time can equivalently be generated by choosing an interval with probability proportional to its duration, then sampling uniformly within that interval.The interval weight is Tj/(T − T0).Source: A Obtaining samples from the… · p. 14
- When the full trajectory duration is unknown, reservoir sampling across discarded time-segments retains uniformly distributed sample times while keeping only one segment of the skeleton in memory.After processing, samples before T0 = 0.2 T are discarded.Source: A Obtaining samples from the… · p. 14
B Implementation details
The implementation represents network parameters as layerwise weights and biases, with task-specific likelihoods and priors. It initializes chains from MAP estimates, including a pruned MAP for sticky samplers on image tasks.
B Implementation details
- Networks use a recursive forward pass with elementwise tanh activations, except the MNIST feedforward network uses ReLU; the output parameterizes the task likelihood.Weights and biases are collected across layers as β.Source: B Implementation details · p. 14
B.1 Model specifications
- MNIST classification uses a categorical likelihood, while regression uses a Gaussian likelihood with learned or, for one-dimensional datasets, fixed noise scale.Non-sticky PDMPs and NUTS use fan-in-scaled Gaussian weight priors; sticky PDMPs use spike-and-slab priors, and biases have Gaussian priors but cannot freeze.Source: §B.1 · pp. 14–15
B.2 MAP and initialisation
- For each dataset and split, the authors find a same-architecture MAP estimate by minimizing the objective with Adam, using constant or cosine-annealed learning rates.The image-model MAP is pruned by zeroing small weights while maintaining accuracy tolerance on 2,000 held-out test images, then refitted.Source: §B.2 · p. 15
- All samplers start at the MAP for a fair comparison, but this initializes chains within one mode and can hinder full state-space exploration.For complex tasks, the MAP is also described as essential for a well-behaving sampler.Source: §B.2 · p. 15
C The Boomerang sampler
The Boomerang sampler replaces linear motion with an elliptical flow around a Gaussian reference, chosen to approximate the BNN posterior. The section describes how the reference and sticky dynamics are set for BNN sampling.
C The Boomerang sampler
- Boomerang uses an elliptical flow instead of linear motion, with a Gaussian reference whose mean and covariance are typically chosen to resemble the posterior.The covariance also defines the velocity distribution; the flow handles the Gaussian reference component exactly, leaving event rates to reflect departures from it.Source: C The Boomerang sampler · pp. 15–16
- Boomerang events are sampled by thinning the rate, and contour reflections are supplemented by Poisson-rate velocity refreshment to ensure ergodicity.The refreshed velocity is drawn from N(0, Σ).Source: C The Boomerang sampler · p. 16
C.1 Reference measure for BNNs
- The reference mean is set to the MAP, and its diagonal covariance combines prior precision with a scaled empirical Fisher estimate, with each diagonal entry floored at 10^-8.The overall reference precision is hand-scaled; the log-noise-scale entry for regression instead uses the curvature of its log prior and is not multiplied by that scale.Source: §C.1 · p. 16
C.2 Sticky Boomerang sampler
- At zero, an active coordinate freezes with its velocity stored and later thaws after an exponential waiting time; frozen coordinates are excluded from event rates and reflections.At refreshment, frozen coordinates keep the sign of their stored velocity and redraw its magnitude and thawing time; MAP-centering changes active motion, not the sticky dynamics.Source: §C.2 · p. 16
D.1 Data
The appendix records the experimental data and settings, defines predictive metrics, and describes the comparison baselines.
D.1 Data
- Boston Housing and UCI data use five random 90/10 train/test splits; image experiments use all 60,000 MNIST training images and 8,000 test images.For MNIST, 2,000 test images select the MAP pruning threshold; metrics use the remaining 8,000.Source: §D.1 · pp. 16–17
D.2 Hyperparameters
- Prior scales are fixed per experiment, while Boomerang reference scale is tuned per dataset and UCI inclusion weights w = 0.3 and w = 0.1 provide two sparsity levels.The ZigZag adds a constant excess switching rate to ensure ergodicity.Source: §D.2 · p. 17
D.3 Performance metrics
- Regression metrics include RMSE, NLL, CRPS, and central 90% interval coverage; image metrics include accuracy, NLL, ECE, and Brier score.UCI results average over five splits; one-dimensional and image experiments are single runs.Source: §D.3 · pp. 17–18
D.4 Baselines
- NUTS uses the same priors and likelihood as dense PDMPs, while the image experiments also compare against a separately fitted dense SGD network.NUTS uses four chains with 1,000 warmup and 1,000 retained draws; SGD is trained by maximum likelihood without the prior.Source: §D.4 · p. 18
D.4.1 LBBNN method
The appendix details the variational LBBNN baseline and supplements the sampler evaluation with bounding details, validation experiments, and further predictive comparisons.
D.4.1 LBBNN method
- LBBNN augments each weight with a binary inclusion variable, allowing weights to be exactly zero under a spike-and-slab prior.The implemented version uses the same prior inclusion probability and slab as sticky samplers; posterior indicators are sampled from their learned probabilities.Source: §D.4.1 · p. 18
D.5 The bounding strategy
- The rate bound is adapted by balancing empty windows against rejected proposals, with a 1% inflation factor for numerical slack.The adaptation uses α = 1.01 for one-dimensional and UCI experiments and α = 1.02 for image models.Source: §D.5 · p. 19
D.6 Computing infrastructure
- The one-dimensional, small UCI, medium UCI, and MNIST runs took about 6 minutes, 30 minutes, 2.5–4.5 hours, and 16 hours, respectively.Runs used an Apple M3 laptop for one-dimensional and small UCI experiments and a server with four NVIDIA A10 GPUs for image models.Source: §D.6 · p. 19
E Additional experiments
- The additional experiments cover sampler validation, sparse linear regression, further one-dimensional and UCI datasets, prior-inclusion ablations, convergence, bounding, and minibatching.The section identifies these as supplementary analyses beyond the main experiments.Source: E Additional experiments · p. 19
E.1 Banana distribution
- On the banana target, ZigZag and Boomerang draws cover the target including its tails, and their marginals closely approximate the exact distribution.Boomerang exploration depends strongly on reference scale: c = 0.1 performs comparably to ZigZag, whereas c = 10 misses the banana arms; refreshment rate has a smaller effect.Source: §E.1 · pp. 19–20, Figure 5 · p. 20, Figure 6 · p. 20
E.2 Sparse linear regression
- On sparse linear regression, sticky PDMP posterior coefficient means agree closely with Gibbs, and inclusion-probability deviations mostly fall within two Monte Carlo standard deviations.The deviation scatter is symmetric around zero with no trend in P(γi = 1 | y).Source: §E.2 · p. 21, Figure 7 · p. 21
E.3 More one-dimensional examples
- Dense samplers and NUTS give nearly identical one-dimensional fits; dense methods slightly outperform sticky samplers on Hernandez and multiscale, while methods perform alike on gap data.On gap data, sticky ZigZag has the lowest RMSE, sticky ZigZag and LBBNN tie in CRPS, and dense Boomerang and NUTS have the lowest NLL.Source: §E.3 · p. 22, Table 4 · p. 22
E.4 More UCI datasets
Additional UCI results show competitive sticky samplers alongside dataset- and network-dependent sparsity, while prior inclusion and bounding choices affect performance and computational cost.
E.4 More UCI datasets
- Sticky samplers are competitive with non-sticky samplers and NUTS on the additional UCI datasets, while LBBNN is behind in RMSE and CRPS on every dataset.On the small concrete network, sticky-sampler RMSE rises from 5.15 to 5.64 as sparsity increases.Source: §E.4 · p. 22, Table 5 · p. 23, Table 6 · p. 24
- Achieved sparsity depends on dataset and network: small networks reach 42–61% at w = 0.3 and 51–80% at w = 0.1 across the three datasets.For medium networks, sparsity remains close to the prior sparsity 1 − w.Source: §E.4 · p. 22
- The authors attribute some medium-network sampler differences to possible nonconvergence, but convergence analysis was conducted only on Boston.They note that the attribution is uncertain.Source: §E.4 · p. 23
E.5 More convergence diagnostics
- The convergence diagnostics report rank-normalised R-hat and bulk ESS; the authors treat R-hat ≤ 1.01 and bulk ESS above 400 as signs of convergence.PDMP ESS is calculated from prediction draws ordered in time, not from the continuous-time trajectory average.Source: §E.5 · p. 23, Table 8 · p. 25
E.6 Effect of prior inclusion weight
- On Boston, achieved sparsity closely follows prior sparsity 1 − w for both networks and sticky samplers, though small networks on three other datasets retain more weights than the prior implies.The prior inclusion probability w must be chosen in advance and largely determines sparsity.Source: §E.6 · pp. 23, 25
- For the small network, predictive performance improves as w increases and levels off around w = 0.5, where sticky samplers track NUTS closely; very sparse models lose accuracy.The medium-network sticky Boomerang stays above NUTS, and its flat curve is interpreted as slow mixing rather than an effect of w.Source: §E.6 · p. 25
E.7 Comparison of bounding strategies
- In the constant-rate cost analysis, the single-segment bound needs about (K + 3)/4 times fewer gradient evaluations per event than the grid strategy.For short freeze or thaw windows, the expected gain is about 3/2 per event.Source: §E.7 · pp. 25–26
- On Boston split 0, the single-segment scheme yields 1.84, 2.52, 1.66, and 1.58 times as many events per gradient evaluation for ZigZag, sticky ZigZag, Boomerang, and sticky Boomerang.Each version used exactly 10^6 gradient evaluations.Source: §E.7 · p. 26
E.8 More MNIST results
On MNIST, sticky samplers produce highly sparse FFN weights while losing at most 0.7 percentage points of accuracy against SGD; LeNet-5’s dense and sticky ZigZag are about equally robust to noise. Minibatching introduces bias, and UCI comparisons show no systematic change in predictive performance.
E.8 More MNIST results
- The pruned MAP has the lowest FFN NLL; dense ZigZag has the lowest FFN accuracy but lowest calibration error, and leads LeNet-5 in accuracy and Brier score.Sticky Boomerang is slightly behind sticky ZigZag on every predictive metric.Source: §E.8 · p. 26
- LeNet-5’s six shared 5 × 5 first-layer filters are structured for both samplers, while dense and sticky ZigZag are about equally robust to Gaussian noise.Figure 12 shows the first-layer filters and robustness across noise standard deviations σ on the [0, 1] pixel scale.Source: §E.8 · p. 26, Figure 12 · p. 29
- 98% against 91% of FFN weights are zero for sticky samplers versus β⋆ ref, while accuracy is at most 0.7 percentage points below SGD.Results use the 8000 MNIST test images not used for pruning.Source: §E.8 · p. 26, Table 9 · p. 27
E.9 Minibatching
- Minibatching does not change UCI predictive performance systematically, with some comparisons worse and others better; the authors note this does not guarantee the same behavior on MNIST.The MNIST full-batch comparison was infeasible within their budget; Table 10 compares small-network runs across five UCI splits using batches of 32 observations.Source: Table 10 · p. 27, §E.9 · p. 27
- Minibatching uses a fresh batch for each bounding horizon, simulating a rescaled minibatch posterior within that horizon; gradient noise means the process no longer has the exact posterior as its invariant distribution.The authors apply this approximate scheme to their MNIST models.Source: §E.9 · p. 27