Implementation:Pyro ppl Pyro NeuTra

Property	Value
Implementation Type	Pattern Doc
Source File	examples/neutra.py
Module	examples
Pyro Features	NeuTraReparam, AutoDiagonalNormal, AutoNormalizingFlow, poutine.reparam, block_autoregressive transforms, MCMC (NUTS), SVI, custom TorchDistribution
Paper	Hoffman et al. (2019), "NeuTra-lizing Bad Geometry in Hamiltonian Monte Carlo Using Neural Transport"

Overview

This file demonstrates Neural Transport (NeuTra) reparameterization, a technique that improves HMC sampling by learning a transport map from a simple distribution to the complex posterior geometry. The example uses a banana-shaped bivariate distribution that is difficult for standard HMC.

The workflow:

1. Define a custom BananaShaped distribution with correlated, non-Gaussian geometry
2. Train an autoguide (either AutoDiagonalNormal or AutoNormalizingFlow with BNAF) using SVI
3. Use NeuTraReparam to reparameterize the model so HMC runs in the "warped" (simplified) space
4. Transform samples back to the original space using the learned transport map

The example compares: vanilla HMC, DiagNormal + NeuTra HMC, and BNAF + NeuTra HMC, showing progressive improvement in posterior sampling quality.

Code Reference

class BananaShaped(dist.TorchDistribution):
    def __init__(self, a, b, rho=0.9):
        self.mvn = dist.MultivariateNormal(torch.tensor([0., 0.]),
            covariance_matrix=torch.tensor([[1., rho], [rho, 1.]]))
        super().__init__(event_shape=(2,))

    def log_prob(self, x):
        x, y = x[..., 0], x[..., 1]
        u0 = x / self.a
        u1 = (y - self.b * (u0**2 + self.a**2)) * self.a
        return self.mvn.log_prob(torch.stack([u0, u1], dim=-1))

def model(a, b, rho=0.9):
    pyro.sample("x", BananaShaped(a, b, rho))

# Train NeuTra guide
guide = AutoNormalizingFlow(model,
    partial(iterated, args.num_flows, block_autoregressive))
fit_guide(guide, args)

# Reparameterize model and run HMC in warped space
neutra = NeuTraReparam(guide.requires_grad_(False))
neutra_model = poutine.reparam(model, config=lambda _: neutra)
mcmc = run_hmc(args, neutra_model)

# Transform samples back to original space
zs = mcmc.get_samples()["x_shared_latent"]
samples = neutra.transform_sample(zs)

I/O Contract

Parameter	Type	Description
-n / --num-steps	int	SVI steps for guide training (default: 10000)
-lr / --learning-rate	float	Learning rate (default: 1e-2)
--num-warmup	int	NUTS warmup steps (default: 500)
--num-samples	int	NUTS samples (default: 1000)
--param-a	float	BananaShaped parameter a (default: 1.15)
--param-b	float	BananaShaped parameter b (default: 1.0)
--num-flows	int	Number of BNAF layers (default: 1)

Output:

• PDF figure with 8 subplots comparing: density, vanilla HMC, DiagNormal guide, DiagNormal + NeuTra (warped and transformed), BNAF guide, BNAF + NeuTra (warped and transformed)

Usage Examples

# Run with default parameters
# python neutra.py -n 10000 --num-samples 1000

# Customize banana shape and flows
# python neutra.py --param-a 2.0 --param-b 0.5 --num-flows 2

Related Pages

• Pyro_ppl_Pyro_SparseGammaDEF - Uses AutoDiagonalNormal guide