GradInf: Gradient Estimation as Probabilistic Inference

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.

• The paper demonstrates that gradient estimation can be reframed as programmable probabilistic inference through systematic coupling and factorization of probabilistic programs.
• It employs source-to-source transformations to construct unbiased and low-variance gradient estimators, achieving up to 370x variance reduction in empirical benchmarks.
• The methodology unifies classical estimators with modern AD techniques, enabling modular and competitive performance across combinatorial and continuous models.

GradInf: Gradient Estimation as Probabilistic Inference

Introduction

"GradInf: Gradient Estimation as Probabilistic Inference" (2607.07840) formulates gradient estimation for expectations under probabilistic programs as a programmable and compositional probabilistic inference problem. This reduction is made precise via two core statistical operations: coupling and factorization, which together enable automated transformations on higher-order probabilistic programs written in a language-agnostic core calculus. The authors present GradInf, a probabilistic programming system and transformation framework, which provides a denotational-semantic foundation for these operations. GradInf enables automatic construction of unbiased and low-variance gradient estimators as probabilistic programs themselves, leveraging existing probabilistic inference and AD techniques. The work unifies classical estimators with modern advances and demonstrates the modular construction of novel, high-performance gradient estimators across multiple domains.

Gradient Estimation as Probabilistic Inference

The foundational observation motivating GradInf is that the gradient of an expectation under a family of distributions—$\nabla_\theta \mathbb{E}_{x\sim\mu_\theta}[f(x)]$—cannot, in general, be computed directly by AD on the program text, due to intractable integrals and stochastic control flow. Existing approaches include the score function (REINFORCE), pathwise (reparameterization), smoothed perturbation analysis (SPA), and measure-valued derivatives (MVD), each with algorithmic, statistical, and variance tradeoffs contingent on the program's characteristics.

GradInf reframes gradient estimation as follows: rather than manipulating formulas directly, it performs program transformations to synthesize a "gradient inference program." This program expresses a gradient estimator as probabilistic inference over a derived target distribution, itself generated by formal transformations (see below). These inference targets can be efficiently differentiated using standard AD, provided the probabilistic dependencies are suitably factorized.

Programmable Couplings and Factorizations

GradInf introduces two syntactic and semantic program transformations:

1. Coupling Transformation ($\mathcal{C}$): Produces a probabilistic product program whose marginals correspond to runs of the original program at parameter values $\theta$ and $\theta'$, allowing generic common-randomness (CRN) strategies, maximal coupling, or other joint strategies. This transformation propagates couplings across arbitrary higher-order constructs.
2. Factorization and Partial Probability Evaluation ($\mathcal{E}$, $\mathcal{S}$): Refactors the coupled program so that a subset of random choices ("primal", e.g., at $\theta$) are first fixed and "partial evaluation" is applied, so only the residual randomness (in the second copy, e.g., at $\theta+\varepsilon$) is exposed to inference. This enables direct application of variance-reduction via Rao-Blackwellization and the use of sophisticated inference algorithms in the gradient estimation procedure.

This workflow is depicted as a sequence of program transformations ending in a differentiable Monte Carlo estimator. The soundness of this composition is proven using proof-relevant logical relations on the language's denotational semantics (QBS).

Figure 1: Applying two inference algorithms to $t_{\mathcal{S}}$, visualizing how difference estimation and probabilistic inference yield reduced variance.

System Design and Semantics

GradInf’s core language supports higher-order functions, general recursion, and both continuous and discrete random primitives, in the semantic setting of quasi-Borel spaces. Couplings and factorized programs are manipulated by programmable source-to-source transformations, which are sound if they preserve the requisite logical relations (ensuring the coupling/factorization is valid in the measure-theoretic sense). Crucially, GradInf decouples primitive random choice semantics from gradient estimator design, so the practitioner only specifies desired coupling/factorization (e.g., stratified CRN, maximal coupling).

GradInf can encode both classical estimators and new ones. For example, CRN-based SIPA/IPA, SPAs, MVD, DisARM, BitFlip, and straight-through estimators each correspond to coupling/inference scheme choices.

Empirical Results: Efficiency, Unbiasedness, and Compositionality

Variance Reduction

GradInf enables novel estimators optimized for variance: for example, in an M/M/c queueing model, applying variable elimination to the factorized program yields up to 110x variance reduction over REINFORCE and 16x over prior SPA estimators as shown empirically.

Figure 2: Variance scaling plots for gradient estimators expressed by GradInf on the M/M/c queue, showing dramatic variance reduction with GradInf-CRN-VE.

Combinatorial and continuous models such as trinomial option pricing and chemical reaction networks benefit from the ability to swap inference backends: stratified importance resampling (SIR), sequential Monte Carlo (SMC), or variable elimination. In gene transcription parameter estimation, GradInf's modular factorization and SMC estimator achieves up to 370x reduction in time-adjusted variance.

Figure 3: Generative model for trinomial option pricing (left).

Figure 4: Generative model for gene transcription and stochastic simulation via Gillespie's algorithm.

Empirical Unbiasedness

All estimators produced by GradInf are formally unbiased. Empirical validation shows statistical indistinguishability between means of novel estimators and gold-standard (but high-variance) baselines, verified by paired equivalence testing.

Figure 5: Gradient estimate samples from baseline and GradInf-derived estimators in several domains, with empirical means indicating unbiasedness.

System Runtime

Despite expressive semantics and modular transformations, GradInf’s runtime is competitive: when compared to ADEV and Storchastic, GradInf incurs only moderate (3x–6x) overhead, scaling favorably with program complexity.

Figure 6: Runtime scaling of the score-function estimator for several gradient estimation frameworks. GradInf achieves favorable runtime scaling relative to prior systems.

Implications, Limitations, and Future Directions

GradInf offers a compositional, semantically justified, and automated route to modular gradient estimator design, closing the gap between program semantics (e.g., stochastic control, higher-order structure) and statistical performance (variance, unbiasedness). The approach unifies classical, control-variated, and advanced SMC-based estimators.

Practically, this enables:

• Modular estimator construction for practitioners: inference algorithms, couplings, and factorization strategies can be independently improved and substituted.
• Theoretical progress: formal proofs of unbiasedness and variance properties for new estimators; a unified language for relating disparate estimation methods.
• Automated gradient estimator synthesis in PPLs, supporting higher-order programs and stochastic recursion.

Limitations include manual annotation of primitive couplings/factorizations, lack of support for unbounded recursion and infinite data structures, and absence of advanced variance reduction (e.g., control variates learned online). Higher-order gradient estimation remains an open direction; GPU- and batch-execution support would further broaden the framework's applicability. Automating primitive coupling selection, perhaps via symbolic disintegration or meta-learning, is a natural extension.

Conclusion

GradInf provides a denotationally justified transformation calculus for gradient estimation in probabilistic programs, reducing gradient estimation to programmable probabilistic inference. This perspective unifies standard estimators and facilitates novel, empirically superior gradient estimation schemes. The system is demonstrated to provide unbiased, lower-variance estimators, compositional construction, and competitive runtime performance. The compositional view of coupling, factorization, and inference portends continued advances in automatic, statistically optimal gradient estimation across scientific and ML applications.