Practical Neuron-level Pruning Framework for Bayesian Neural Networks
V. Kuboň (TU Delft - Electrical Engineering, Mathematics and Computer Science)
L. Laurenti – Graduation committee member (TU Delft - Team Luca Laurenti)
Steven Adams – Mentor (TU Delft - Team Luca Laurenti)
Avishek Anand – Graduation committee member (TU Delft - Web Information Systems)
More Info
expand_more
Other than for strictly personal use, it is not permitted to download, forward or distribute the text or part of it, without the consent of the author(s) and/or copyright holder(s), unless the work is under an open content license such as Creative Commons.
Abstract
Bayesian Neural Networks (BNNs) offer uncertainty quantification but are computationally expensive, limiting their practical deployment. This paper introduces a neuron-level pruning framework that reduces BNN complexity while preserving predictive performance. Unlike existing weight-level pruning techniques, our approach removes entire neurons, enabling significant memory savings and inference speedups without requiring specialized hardware. We propose a pruning loss based on the Wasserstein distance, balancing model sparsity and predictive accuracy. Our method is fully automatic, eliminating the need for manual hyperparameter tuning. Experimental results on UCI regression and Fashion MNIST datasets demonstrate that our framework can prune over 80% of neurons while maintaining predictive distribution integrity. Additionally, we validate the Lottery Ticket Hypothesis in the Bayesian setting, showing that pruned subnetworks retain performance and learn faster when retrained. This work represents a step toward making BNNs more scalable for real-world applications.