Principle:Pyro ppl Pyro Neural Transport
| Knowledge Sources | |
|---|---|
| Domains | MCMC, Normalizing Flows, Reparameterization |
| Last Updated | 2026-02-09 09:00 GMT |
Overview
Neural Transport (NeuTra) reparameterizes the target distribution using a normalizing flow to transform a geometrically challenging posterior into a simpler geometry where Hamiltonian Monte Carlo can sample efficiently.
Description
Hamiltonian Monte Carlo (HMC) and its adaptive variant NUTS are effective MCMC samplers for continuous distributions, but they struggle with pathological posterior geometries:
- Strong correlations: When parameters are highly correlated, the optimal trajectory length and step size differ across directions.
- Funnel-shaped posteriors: Common in hierarchical models, where the scale of one parameter depends on the value of another.
- Banana-shaped posteriors: Curved, non-linear correlations that prevent efficient exploration.
Mass matrix adaptation (a standard feature of modern HMC implementations) can partially address correlations by preconditioning with the posterior covariance. However, it cannot handle non-linear correlations or varying curvature.
NeuTra (Neural Transport) solves this more completely using a normalizing flow:
- Train a normalizing flow T to approximate the posterior. The flow learns a differentiable, invertible mapping from a simple base distribution (e.g., standard Normal) to the posterior.
- Reparameterize HMC: Instead of sampling in the original parameter space, run HMC in the transformed space where the posterior (as seen through the flow) is approximately standard Normal.
- Transform samples back: Apply the flow to map samples from the transformed space back to the original parameter space.
In the transformed space, the posterior is approximately isotropic Gaussian, which is the ideal geometry for HMC. This eliminates the need for careful step-size tuning and mass matrix adaptation.
Usage
Use Neural Transport when:
- HMC/NUTS is struggling with pathological posterior geometry (high R-hat, low ESS).
- The posterior has strong non-linear correlations (funnels, bananas).
- Standard mass matrix adaptation is insufficient.
- You can afford a preprocessing step (training the flow) to improve sampling efficiency.
- Working with hierarchical models that commonly exhibit funnel geometry.
Theoretical Basis
Problem: bad geometry for HMC:
# HMC proposes by simulating Hamiltonian dynamics:
# H(q, p) = U(q) + K(p)
# where U(q) = -log p(q | data) and K(p) = 0.5 * p^T M^{-1} p
# Efficiency depends on the geometry of U(q):
# - Isotropic: all directions have similar curvature -> ideal
# - Correlated: different curvatures -> step size limited by steepest direction
# - Funnel: curvature varies with position -> no single step size works
# Mass matrix M can address linear correlations: set M = Cov(q)
# But M is constant -- cannot adapt to position-dependent curvature
NeuTra reparameterization:
# Learn a normalizing flow T: z -> q
# such that if z ~ Normal(0, I), then T(z) approximately follows the posterior
# Transformed target density:
# p_z(z) = p(T(z) | data) * |det(dT/dz)|
# If T is a good approximation:
# p_z(z) is approximately Normal(0, I) -- ideal for HMC!
# Run HMC in z-space:
# - The target is nearly isotropic
# - Standard step size and identity mass matrix work well
# - No adaptation needed
# Transform samples:
# q_i = T(z_i) -- map back to original space
Training the flow:
# Option 1: Train by minimizing KL(q_flow || p_posterior)
# This is standard variational inference with a flow guide
# L = E_{z ~ N(0,I)}[log p(T(z) | data) + log |det(dT/dz)| - log N(z; 0, I)]
# Option 2: Train on MCMC samples from a preliminary run
# If initial MCMC samples are available (even poor quality):
# Minimize: E_{q in samples}[||T^{-1}(q) - z||^2] or KL(p_empirical || p_flow)
# The flow need not be perfect -- even a rough approximation
# significantly improves HMC geometry
Complete NeuTra algorithm:
# Phase 1: Train the flow
# 1. Run SVI with a normalizing flow guide to approximate the posterior
# 2. Or: run short MCMC, then fit flow to samples
# Phase 2: Reparameterized MCMC
# 1. Define transformed target:
# log p_z(z) = log p(T(z) | data) + log |det(dT/dz)|
# 2. Run NUTS in z-space targeting p_z
# (use identity mass matrix, default step size)
# 3. Collect samples z_1, ..., z_N
# 4. Transform: q_i = T(z_i)
# Phase 2 samples are exact (up to MCMC convergence)
# The flow only changes the parameterization, not the target distribution
# Correctness is guaranteed by the change-of-variables formula