Paper: Push: Concurrent Probabilistic Programming for Bayesian Deep Learning (ICLR-2024)

Reviewer comment (Writing):
Is there an intuition you can provide for what problem (if any) the program illustrated in Figs 1 and 2 is solving? (E.g., why has a user decided to write `_gather`?)

Author response:
We thank the reviewer for this question. The program illustrated in Figures 1 and 2 would be needed to implement the all-to-all communication pattern between neural networks required by SVGD. We use SVGD as an example of the densest communication pattern between neural networks that could be exhibited by a BDL algorithm. Deep ensembles are an example of a lack of communication between neural networks in a BDL context. We hope to investigate BDL algorithms and their characteristics  with sparser and dynamic communication patterns, i.e., communication dependent on the state of the neural networks, in future work. To more easily and systematically explore this, we were motivated to build a library that could easily express different communication patterns and heterogeneous computations across an ensemble of neural networks.

Reviewer comment (Reproducibility):
It seems somewhat unnatural that in the implementation of SVGD, the "leader" process needs to itself be a particle, rather than just a process that coordinates the other particles. At several points in the code, there is special-case handling of the leader's parameters vs. everyone else's parameters. It is unclear why Push requires each "worker" in the concurrent algorithm to itself be a particle that is storing one set of weights for a neural network. It might be more useful to present Push as a general-purpose library for concurrent programming on GPUs, with an application to BDL.

Author response:
We thank the reviewer for this question. It is possible to design push as a process + particles system so that we do not need to special case the "leader" process. The tradeoff in this case is that the process would need to introduce primitives itself for communicating with particles and particles would need special case handling communication from the process. We choose the design where everything is a particle so that communication can be handled uniformly. It might be possible to present Push as a general-purpose library for concurrent programming with applications to BDL, although the multi-threaded execution model operates at the level of particles/neural networks and not tensors. As we mention in the paper, it is possible that Push is useful for expressing other algorithms such as meta-heuristic algorithms, although we leave these investigations for future work.

Reviewer comment (Novelty):
The reviewer asks why straight torch is insufficient for the task, noting that with primitives like model.to(torch.device("cuda:0")), model.to(torch.device("cuda:1")), and so on, it is already possible to control GPU model placement and data transfers without an additional library.

Author response:
We thank the reviewer for this question. The advantage of using the library is primarily speed and convenience. The library automates the underlying threading mechanism that coordinates concurrent particle communication and execution on different GPUs for the user. This is why in Figure 4, the library implementation of Stein Variational Gradient Descent (which has all-to-all communication) is faster than the PyTorch only implementation even on 1 GPU across a range of particles. The library additionally creates views of parameters sets to help the user with memory management in cases like Stein Variational Gradient Descent where multiple neural network parameter sets are considered for the parameter update. The underlying multi-threaded implementation of the library also enables the library to utilize all GPU devices concurrently, whereas a straight PyTorch implementation does not. This can lead to speedups since the library implementation can access the GPU devices in any order and concurrently, whereas a PyTorch only implementation would access the GPU devices in the order specified in the original program and sequentially. While it is certainly possible to utilize multiple GPUs using PyTorch alone, one might need to write additional multi-threaded code or other modules to get similar performance.

Reviewer comment (Reproducibility):
The reviewer asks why not use jax.pmap, noting that all the BDL proposals tested could be implemented as pure functions of weights, and that JAX provides collectives such as psum and all-gather which meet most of push's parallel communication needs, while also offering a JIT compiler to manage memory and hide latency effectively.

Author response:
We thank the reviewer for this question. While this may change in the future, our understanding is that "Jax has limited support for Python concurrency" as stated in the documentation https://jax.readthedocs.io/en/latest/concurrency.html. Thus, it would be difficult to implement a concurrent library on top of Jax. It would be an interesting direction of future work to try to implement the library in Jax given its many useful features (such as JIT compilation) and nice engineering. Pmap and all-gather are useful constructs, but they operate at a lower level of abstraction compared to the level that particles are represented at. The concurrent execution semantics of particles in the library can provide the illusion of parallelism to scale across additional GPUs, but it also makes it easier to implement BDL algorithms where the communication patterns involved are not uniform and the computations that each particle perform may depend on the state of the particle. For example, an all-to-all communication pattern as in Stein Variational Gradient Descent is not scalable, and so we may want to explore BDL algorithms that have sparser and dynamic communication patterns instead that would not be easily expressed with pmap or all-gather.