Optimising Density Computations in Probabilistic Programs via Automatic Loop Vectorisation

2025-11-14

AI Summary

In probabilistic programming, inference over large-scale data—such as long sequences or massive datasets—incurs high computational overhead due to iterative data traversal; existing vectorization techniques struggle with general repetitive structures (e.g., nested loops, data-dependent control flow), require manual intervention, and exhibit poor generalizability. This paper proposes a fully automatic loop vectorization method: it formalizes the translation from imperative probabilistic programming languages (PPLs) to a vector-oriented low-level language, integrating speculative parallel execution with semantics-preserving fixed-point verification. The approach requires no user annotation, fully supports complex control flow, and is implemented in Pyro. Experiments show speedups of 1.1–6× over state-of-the-art baselines, reduced GPU memory consumption in most scenarios, and 100% coverage across all tested models—demonstrating significant improvements in generality, robustness, and end-to-end inference throughput.

Technology Category

• Reasoning under Uncertainty: Probabilistic Programming
• Machine Learning: Probabilistic Circuits and Graphical Models
• Natural Language Processing: Code Generation / Program Synthesis from Natural Language

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Large language models for search

Abstract

Probabilistic programming languages (PPLs) are a popular tool for high-level modelling across many fields. They provide a range of algorithms for probabilistic inference, which analyse models by learning their parameters from a dataset or estimating their posterior distributions. However, probabilistic inference is known to be very costly. One of the bottlenecks of probabilistic inference stems from the iteration over entries of a large dataset or a long series of random samples. Vectorisation can mitigate this cost, but manual vectorisation is error-prone, and existing automatic techniques are often ad-hoc and limited, unable to handle general repetition structures, such as nested loops and loops with data-dependent control flow, without significant user intervention. To address this bottleneck, we propose a sound and effective method for automatically vectorising loops in probabilistic programs. Our method achieves high throughput using speculative parallel execution of loop iterations, while preserving the semantics of the original loop through a fixed-point check. We formalise our method as a translation from an imperative PPL into a lower-level target language with primitives geared towards vectorisation. We implemented our method for the Pyro PPL and evaluated it on a range of probabilistic models. Our experiments show significant performance gains against an existing vectorisation baseline, achieving $1.1$--$6 imes$ speedups and reducing GPU memory usage in many cases. Unlike the baseline, which is limited to a subset of models, our method effectively handled all the tested models.

Problem

• Automatically vectorizing loops in probabilistic programs to improve performance
• Overcoming limitations of manual and existing automatic vectorization techniques
• Enabling efficient density computations through speculative parallel execution of loops

Innovation

• Automatic loop vectorisation for probabilistic programs
• Speculative parallel execution with fixed-point checks
• Translation from imperative PPL to vectorisation-primed language