Practical Neuron-level Pruning Framework for Bayesian Neural Networks
V. Kuboň (TU Delft - Electrical Engineering, Mathematics and Computer Science)
Luca Laurenti – Graduation committee member (TU Delft - Mechanical Engineering)
Steven Adams – Mentor (TU Delft - Mechanical Engineering)
Avishek Anand – Graduation committee member (TU Delft - Electrical Engineering, Mathematics and Computer Science)
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.