ADEV: Sound Automatic Differentiation of Expected Values of Probabilistic Programs

Abstract

Optimizing the expected values of probabilistic processes is a central problem in computer science and its applications, arising in fields ranging from artificial intelligence to operations research to statistical computing. Unfortunately, automatic differentiation techniques developed for deterministic programs do not in general compute the correct gradients needed for widely used solutions based on gradient-based optimization. In this paper, we present ADEV, an extension to forward-mode AD that correctly differentiates the expectations of probabilistic processes represented as programs that make random choices. Our algorithm is a source-to-source program transformation on an expressive, higher-order language for probabilistic computation, with both discrete and continuous probability distributions. The result of our transformation is a new probabilistic program, whose expected return value is the derivative of the original program's expectation. This output program can be run to generate unbiased Monte Carlo estimates of the desired gradient, which can then be used within the inner loop of stochastic gradient descent. We prove ADEV correct using logical relations over the denotations of the source and target probabilistic programs. Because it modularly extends forward-mode AD, our algorithm lends itself to a concise implementation strategy, which we exploit to develop a prototype in just a few dozen lines of Haskell (https://github.com/probcomp/adev).

arXiv commentto appear at POPL 2023

References in corpus (21)

• Pyro: Deep Universal Probabilistic Programming
• Gradient Estimation Using Stochastic Computation Graphs
• Learning Disentangled Representations with Semi-Supervised Deep Generative Models
• A Convenient Category for Higher-Order Probability Theory
• Monte Carlo Gradient Estimation in Machine Learning
• Denotational validation of higher-order Bayesian inference
• Pathwise Derivatives Beyond the Reparameterization Trick
• A Simple Differentiable Programming Language
• Automatic Differentiation in PCF
• A Domain Theory for Statistical Probabilistic Programming
• Commutative monads as a theory of distributions
• Measurable Cones and Stable, Measurable Functions
• Backpropagation in the Simply Typed Lambda-calculus with Linear Negation
• Densities of Almost Surely Terminating Probabilistic Programs are Differentiable Almost Everywhere
• Automatic Differentiation of Programs with Discrete Randomness
• Reparameterization Gradient for Non-differentiable Models
• ${\lambda }_S$: Computable Semantics for Differentiable Programming with Higher-Order Functions and Datatypes
• Denotational Correctness of Forward-Mode Automatic Differentiation for Iteration and Recursion
• Learning Proposals for Probabilistic Programs with Inference Combinators
• Automating Involutive MCMC using Probabilistic and Differentiable Programming
• Smoothness Analysis for Probabilistic Programs with Application to Optimised Variational Inference

Cited by in corpus (4)

• Smoothing Methods for Automatic Differentiation Across Conditional Branches
• Probabilistic Programming with Programmable Variational Inference
• Tuning Random Generators: Property-Based Testing as Probabilistic Programming
• GradInf: Gradient Estimation as Probabilistic Inference