GradInf: Gradient Estimation as Probabilistic Inference

2026-07-08

AI Summary

This work addresses the challenge of estimating gradients of expectations in probabilistic programs involving high-dimensional integrals, discrete random choices, and complex dependency structures. It introduces a novel paradigm—gradient inference—that rigorously reduces gradient estimation to a differentiable probabilistic inference problem. By integrating an information-flow type system, programmable couplings and factorization transforms, higher-order probabilistic semantics, and automatic differentiation, the authors develop a verifiable source-to-source transformation framework that automatically synthesizes efficient gradient estimators. This approach not only unifies and reproduces existing state-of-the-art estimators but also generates new estimators that outperform current best methods across several challenging tasks.

Abstract

Gradient estimation -- the task of computing the gradient of the expected value of a probabilistic program -- has diverse applications in scientific computing, but is notoriously difficult because of issues such as high-dimensional integration, discrete random choices, and complex stochastic dependencies. This article introduces gradient inference, a new approach to developing sound and efficient gradient estimators for probabilistic programs. Gradient inference rests on a formal reduction from a gradient estimation problem to a closely related probabilistic inference problem, whose solution can be differentiated to obtain a gradient estimator. This inference problem is obtained by applying two powerful statistical operations -- coupling and factorization -- to the input probabilistic program. Our reduction lets us leverage the rich toolkit of probabilistic inference algorithms to design novel gradient estimators that extend and improve upon existing methods. We introduce GradInf, a probabilistic programming system that facilitates the sound and automated implementation of gradient inference. GradInf is centered around programmable source-to-source transformations for coupling and factorizing higher-order probabilistic programs, whose soundness is proven in terms of a denotational semantics. Key to our development is the use of information-flow typing to allow random choices in a probabilistic program to be factored out and partially evaluated, which improves our ability to deploy sophisticated probabilistic inference algorithms. The resulting system offers practitioners a principled framework for designing gradient estimators. We apply GradInf to several challenging case studies, showing that it can express prominent gradient estimators from the literature and enables the construction of new state-of-the-art estimators that outperform the best existing baselines.

Problem

• gradient estimation
• probabilistic programs
• stochastic dependencies
• discrete random choices
• high-dimensional integration

Innovation

• gradient inference
• probabilistic programming
• coupling
• factorization
• information-flow typing