Probabilistic Programming with Programmable Variational Inference (2406.15742v1)

Abstract: Compared to the wide array of advanced Monte Carlo methods supported by modern probabilistic programming languages (PPLs), PPL support for variational inference (VI) is less developed: users are typically limited to a predefined selection of variational objectives and gradient estimators, which are implemented monolithically (and without formal correctness arguments) in PPL backends. In this paper, we propose a more modular approach to supporting variational inference in PPLs, based on compositional program transformation. In our approach, variational objectives are expressed as programs, that may employ first-class constructs for computing densities of and expected values under user-defined models and variational families. We then transform these programs systematically into unbiased gradient estimators for optimizing the objectives they define. Our design enables modular reasoning about many interacting concerns, including automatic differentiation, density accumulation, tracing, and the application of unbiased gradient estimation strategies. Additionally, relative to existing support for VI in PPLs, our design increases expressiveness along three axes: (1) it supports an open-ended set of user-defined variational objectives, rather than a fixed menu of options; (2) it supports a combinatorial space of gradient estimation strategies, many not automated by today's PPLs; and (3) it supports a broader class of models and variational families, because it supports constructs for approximate marginalization and normalization (previously introduced only for Monte Carlo inference). We implement our approach in an extension to the Gen probabilistic programming system (genjax.vi, implemented in JAX), and evaluate on several deep generative modeling tasks, showing minimal performance overhead vs. hand-coded implementations and performance competitive with well-established open-source PPLs.

• The paper presents a modular approach that overcomes rigid VI objectives in current probabilistic programming languages.
• It introduces compositional transformations that compile generative models into differentiable programs for unbiased gradient estimation.
• Benchmarks confirm minimal runtime overhead, competitive convergence, and broader expressivity compared to existing systems.

Probabilistic Programming with Programmable Variational Inference

The paper "Probabilistic Programming with Programmable Variational Inference" addresses the limitations and opportunities in Variational Inference (VI) within probabilistic programming languages (PPLs). The authors propose a modular approach to integrating VI support into PPLs, moving beyond the monolithic and inflexible implementations found in contemporary systems. The goal of this work is to increase the expressiveness and modularity of VI methods by leveraging compositional program transformations.

Overview of Existing Limitations

The authors identify several key limitations in current PPLs that support VI:

1. Limited Options for VI Objectives: Existing PPLs provide a restricted set of predefined variational objectives and gradient estimators, which limits flexibility and customizability in probabilistic modeling tasks.
2. Duplicative Engineering Effort: Efforts to support new VI objectives or gradient estimators typically require significant code duplication and deep system-specific knowledge.
3. Difficulty of Reasoning: Monolithic implementations intertwine distinct concerns, making it challenging to reason about their correctness and to maintain the codebase.

Proposed Modular Approach

To address these limitations, the paper introduces a new framework that allows users to define VI objectives as programs. The framework is built on two new calculi for generative and differentiable probabilistic programming:

• Generative Probabilistic Programming Language ($\lambda_{\text{Gen}}$): This language allows users to encode probabilistic models and variational families. It includes constructs for sampling and observing random variables, transforming them into distributions over traces.
• Differentiable Probabilistic Programming Language ($\lambda_{\text{ADEV}}$): This is a lower-level language designed for differentiable programming and VI objectives. It extends $\lambda_{\text{Gen}}$ with constructs for computing expected values and differentiable scores, among others.

Program Transformations

The core innovation lies in two program transformations:

1. Sim: Transforms generative programs into procedures that simulate traces and compute their densities.
2. Density: Converts generative programs into density evaluators, which compute the density of given traces.

These transformations compile $\lambda_{\text{Gen}}$ programs into $\lambda_{\text{ADEV}}$ programs, facilitating the construction of unbiased gradient estimators for variational objectives.

Gradient Estimation via ADEV

For gradient estimation, the authors extend and utilize the ADEV framework, originally designed for automatic differentiation of probabilistic programs:

• Extensibility: The framework allows users to define new gradient estimation strategies for individual primitive distributions.
• Correctness: The transformations guarantee that the resulting gradient estimators are unbiased, under mild assumptions. This is formalized and proven using logical relations.

Extended Language Features

The full system further extends the formal model to support:

• Marginalization and Normalization: Constructs for marginalizing auxiliary variables and normalizing probabilistic programs to increase their expressiveness.
• Differentiable Stochastic Estimators: When exact densities are intractable, the system supports stochastic estimators that maintain the necessary differentiability properties.
• Reverse-Mode AD: Implementation enhancements to support reverse-mode automatic differentiation, making the framework more efficient for deep learning applications.

Evaluation

The system's performance is evaluated using several benchmarks, including variational autoencoders and the Attend-Infer-Repeat model. The results indicate that:

• Minimal Overhead: The modular transformations introduce negligible runtime overhead compared to hand-coded implementations.
• Competitiveness: The framework's performance is competitive with existing PPLs such as Pyro, often providing faster convergence.
• Expressivity: The framework supports a broader range of VI objectives and gradient estimators, demonstrating its flexibility.

Implications and Future Work

This work has significant implications for the field of probabilistic programming and variational inference:

• Modular Reasoning: By separating concerns and enabling compositional reasoning, the framework enhances maintainability and extensibility, making it easier to develop and experiment with new inference algorithms.
• Broader Applicability: The system’s support for a wide array of objectives and gradient strategies opens the door to more complex and expressive probabilistic models.
• Future Developments: The formal foundation and modular design pave the way for future research into low-variance gradient estimators and the application of more advanced stochastic programming techniques.

Conclusion

"Probabilistic Programming with Programmable Variational Inference" represents a substantial advancement in the field of probabilistic programming. By introducing modular transformations and extending the expressiveness of PPLs, the authors provide a robust framework that empowers users to define, reason about, and experiment with a wide variety of variational inference algorithms. This work is poised to significantly impact how complex probabilistic models are built and optimized in the future.

This summary captures the essence and implications of the paper's contribution to the development of flexible and modular VI support in probabilistic programming languages.