GradInf: Gradient Estimation as Probabilistic Inference

Gradient estimation—calculating the gradient of an expected value in a probabilistic program—is essential for fields like reinforcement learning and scientific simulation. However, it is notoriously difficult due to high-dimensional integration, discrete random choices, and complex stochastic dependencies. The authors ask: can we treat gradient estimation as a modular probabilistic inference problem to improve the efficiency and automation of these estimators?

The authors introduce gradient inference, a framework that decomposes gradient estimation into three distinct steps: coupling, factorization, and inference. First, the system applies a coupling transformation to the input program, creating a joint distribution over two parameter settings that share randomness. Second, it factorizes this joint distribution into a primal trace and a residual trace, allowing for partial evaluation. Finally, it applies standard probabilistic inference algorithms (such as variable elimination or sequential Monte Carlo) to the residual trace. The authors implement this in GradInf, a system that uses information-flow typing to automate these transformations and ensure they are sound.

GradInf provides a principled, compositional framework for constructing gradient estimators. By reinterpreting existing methods (like the reparameterization trick or score function estimation) as specific combinations of couplings and inference algorithms, the authors unify disparate techniques. Furthermore, they demonstrate that their framework can generate novel, state-of-the-art estimators. In case studies involving queuing theory, mathematical finance, and genetics, GradInf-generated estimators achieved significant variance reduction—up to 370x in some instances—compared to existing baselines.

This work shifts the focus of gradient estimation from manual, ad-hoc derivation to a systematic, programmable workflow. By leveraging the existing, mature toolkit of probabilistic inference, researchers can develop more efficient and robust estimators for complex stochastic models without needing to derive them from scratch for every new problem.