Principle:Pyro ppl Pyro Inference Utilities
| Knowledge Sources | |
|---|---|
| Domains | Probabilistic Programming, Inference, Software Design |
| Last Updated | 2026-02-09 09:00 GMT |
Overview
Inference utilities provide the foundational abstractions, helper functions, and diagnostic tools that support the implementation of diverse inference algorithms in a probabilistic programming framework.
Description
A probabilistic programming system must support multiple inference algorithms (SVI, MCMC, importance sampling, etc.) that share common needs: inspecting model structure, initializing parameters, computing log-probabilities, and resampling weighted particles. The inference utilities layer provides these shared capabilities.
Abstract inference base class: Defines the common interface that all inference algorithms implement. This typically includes methods for initialization, running inference steps, and extracting results. The abstract base class ensures that different inference algorithms can be used interchangeably and composed with other components (e.g., loss functions, optimizers).
Inference helper functions: A collection of utility functions used across inference algorithms:
- Computing log-joint probabilities from execution traces.
- Extracting posterior samples and converting them to usable formats.
- Initializing variational parameters at reasonable starting points.
- Computing predictive distributions from posterior samples.
- Handling plate/batch dimensions correctly in loss computation.
Model inspection: Tools for analyzing the structure of a probabilistic program before running inference. This includes:
- Discovering all sample sites and their properties (distributions, shapes, dependencies).
- Determining which sites are observed vs. latent.
- Checking compatibility between model and guide (matching site names and shapes).
- Detecting potential issues (missing sites, shape mismatches, non-reparameterizable sites).
Resampling: Algorithms for drawing new particle sets from weighted particle collections, used in SMC and importance sampling. Key resampling strategies include multinomial, systematic, and stratified resampling, each with different variance properties.
Usage
Use inference utilities when:
- Implementing a new inference algorithm and needing standard building blocks.
- Inspecting a model's structure to verify correctness before running inference.
- Computing predictive distributions after inference.
- Debugging shape errors or missing sample sites in model-guide pairs.
- Performing resampling operations in particle-based inference methods.
Theoretical Basis
Log-joint computation from trace:
# Given an execution trace T with sites {s_1, ..., s_n}:
log_joint(T) = sum over sites s in T:
s.scale * s.mask * s.log_prob(s.value)
# For ELBO computation:
# model_log_joint = sum of model trace log-probs
# guide_log_joint = sum of guide trace log-probs (latent sites only)
# ELBO = model_log_joint - guide_log_joint
Model inspection protocol:
# Trace the model once with dummy data to discover structure:
# 1. Run model under trace handler with initialization strategy
# 2. For each site, record:
# - name, distribution type, event_shape, batch_shape
# - is_observed, has_rsample (reparameterizable)
# - plate context (conditional independence)
# Compatibility check:
# For each latent site in model:
# assert site exists in guide
# assert model.event_shape == guide.event_shape
# assert guide.has_rsample or score_function_gradient_configured
Resampling algorithms:
# Given N particles with normalized weights w_1, ..., w_N:
# Multinomial resampling:
# Draw N indices i.i.d. from Categorical(w_1, ..., w_N)
# Simple but high variance: Var = O(N)
# Systematic resampling:
# u ~ Uniform(0, 1/N)
# For k = 0, ..., N-1:
# select particle i where cumsum(w)[i] >= u + k/N
# Low variance: ensures each particle sampled ~ N*w_i times
# Stratified resampling:
# For k = 0, ..., N-1:
# u_k ~ Uniform(k/N, (k+1)/N)
# select particle i where cumsum(w)[i] >= u_k
# Variance between multinomial and systematic
Predictive distribution:
# Given posterior samples theta_1, ..., theta_S:
# Predictive density for new observation x*:
# p(x* | data) approx (1/S) * sum_s p(x* | theta_s)
# Algorithm:
# For each posterior sample theta_s:
# Run model forward with theta_s to get x*_s
# Return {x*_1, ..., x*_S} as predictive samples