Push: Concurrent Probabilistic Programming for Bayesian Deep Learning
This work addresses the problem of efficient experimentation in Bayesian deep learning for researchers, though it appears incremental as it builds on existing probabilistic programming approaches.
The authors tackled the challenge of scaling Bayesian deep learning inference by introducing Push, a library that enables concurrent execution on multi-GPU hardware, achieving improved scaling behavior on vision and scientific machine learning tasks.
We introduce a library called Push that takes a probabilistic programming approach to Bayesian deep learning (BDL). This library enables concurrent execution of BDL inference algorithms on multi-GPU hardware for neural network (NN) models. To accomplish this, Push introduces an abstraction that represents an input NN as a particle. Push enables easy creation of particles so that an input NN can be replicated and particles can communicate asynchronously so that a variety of parameter updates can be expressed, including common BDL algorithms. Our hope is that Push lowers the barrier to experimenting with BDL by streamlining the scaling of particles across GPUs. We evaluate the scaling behavior of particles on single-node multi-GPU devices on vision and scientific machine learning (SciML) tasks.