[Paper Review] Amortized Rejection Sampling in Universal Probabilistic Programming

Saeid Naderiparizi, Adam Ścibior+10 more|arXiv (Cornell University)|Oct 20, 2019

Machine Learning and Algorithms18 references4 citations

TL;DR

This paper proposes a novel amortized importance sampling estimator for universal probabilistic programming languages that handles rejection sampling loops with unbounded iterations, ensuring finite variance where naive approaches fail. By reweighting samples based on conditional acceptance probabilities and integrating over the rejection loop's structure, the method achieves stable, finite-variance inference without requiring model rewrites.

ABSTRACT

Naive approaches to amortized inference in probabilistic programs with unbounded loops can produce estimators with infinite variance. This is particularly true of importance sampling inference in programs that explicitly include rejection sampling as part of the user-programmed generative procedure. In this paper we develop a new and efficient amortized importance sampling estimator. We prove finite variance of our estimator and empirically demonstrate our method's correctness and efficiency compared to existing alternatives on generative programs containing rejection sampling loops and discuss how to implement our method in a generic probabilistic programming framework.

Motivation & Objective

• To address the instability of amortized importance sampling in probabilistic programs with unbounded rejection sampling loops, where naive estimators can have infinite variance.
• To develop a method that ensures finite variance in importance sampling estimators even when the number of rejection iterations is unbounded.
• To provide a generic, efficient, and correct inference solution compatible with existing universal probabilistic programming frameworks.
• To demonstrate empirically that the proposed estimator outperforms naive inference compilation in terms of convergence stability and accuracy.
• To formalize the theoretical conditions under which importance sampling with rejection loops remains valid and finite-variance.

Proposed method

• The method introduces a reweighting scheme that accounts for the conditional acceptance probability of each rejection loop iteration, treating the loop as a stochastic process conditioned on acceptance.
• It redefines the importance weight as a product of conditional densities, incorporating the probability of acceptance at each step, thereby ensuring finite variance.
• The estimator is derived by transforming the original generative program into an equivalent form that explicitly models the acceptance condition, allowing for proper normalization of weights.
• The approach leverages inference compilation by training a neural network to approximate the proposal distribution, but modifies the weight computation to account for the loop structure.
• The method ensures that the final importance weight is computed as a product of conditional densities: w ∝ p(z|x, c(x,z)) / q(z|x,y, c(x,z)), where c(x,z) indicates acceptance.
• Theoretical analysis proves that under mild regularity conditions, the resulting estimator has finite variance, even when the number of loop iterations is unbounded.

Experimental results

Research questions

• RQ1 Can amortized importance sampling in probabilistic programs with unbounded rejection sampling loops be made stable with finite variance?
• RQ2 What modifications to the importance weight computation are necessary to ensure finite variance when loops have unbounded iterations?
• RQ3 How can the structure of rejection sampling loops be formally incorporated into inference compilation without requiring model rewrites?
• RQ4 Does the proposed estimator converge to the correct posterior distribution, even when naive estimators fail silently?
• RQ5 Can the method be implemented generically in a universal probabilistic programming framework with minimal changes to existing inference pipelines?

Key findings

• The proposed estimator guarantees finite variance in importance sampling for programs with unbounded rejection sampling loops, unlike naive approaches that can have infinite variance.
• The method achieves stable convergence and correct posterior estimation, even when the number of rejection iterations is unbounded and the proposal distribution is learned via inference compilation.
• Empirical evaluation shows that the proposed estimator outperforms naive inference compilation in terms of convergence speed and accuracy on models with rejection sampling loops.
• Theoretical analysis confirms that the estimator remains valid and finite-variance under mild regularity conditions, even when individual importance weights have finite variance but their product does not.
• The method enables correct inference in models that previously suffered from silent failure due to infinite variance in importance sampling, particularly in simulators and constrained distributions.
• The approach is compatible with existing universal probabilistic programming frameworks and can be integrated with neural proposal networks in inference compilation pipelines.