Probabilistic Programming with Programmable Divide-Conquer-Combine Inference on Modern Hardware

Universal probabilistic programming languages (PPLs) enable the specification of models with stochastic support structure. Posterior inference is notoriously hard for this class of models and remains difficult to accelerate on modern hardware. In response to these challenges, we introduce Upix - the first probabilistic programming system that realises the divide-conquer-combine (DCC) inference algorithm as a framework. In Upix, a model expressed in a universal PPL is automatically split into multiple sub-models with static support structure, which are then compiled with JAX for execution on accelerator hardware. The system allows extensive customisation of inference algorithms by incorporating established concepts from programmable inference literature. To evaluate our system, we implemented two existing DCC algorithms in Upix and instantiated three novel algorithms. We show that our implementation can result in better approximation quality compared to existing approaches by achieving up to 1070 times more computation within the same time budget. On machines with up to 64 CPU cores and 8 GPU devices, we demonstrate that Upix enables the scaling of inference algorithms to workloads that are impractically slow for CPUs and prior methods.