Inference algorithms

Inference algorithms.

aft(targets, flows, *, num_targets=None)[source]

Annealed Flow Transport.

  1. Annealed Flow Transport Monte Carlo, Michael Arbel, Alexander G. D. G. Matthews, Arnaud Doucet https://arxiv.org/abs/2102.07501

Parameters:

  • targets – a list of target programs

  • flows – a list of flows

  • num_targets – the number of targets

Returns:

the inference program

Return type:

q

apgs(target, kernels, *, num_sweeps=1)[source]

Amortized Population Gibbs Sampler.

  1. Amortized Population Gibbs Samplers with Neural Sufficient Statistics, Hao Wu, Heiko Zimmermann, Eli Sennesh, Tuan Anh Le, Jan-Willem van de Meent https://arxiv.org/abs/1911.01382

Parameters:

  • target – the target program

  • kernels – the Gibbs kernels

  • num_sweeps – the number of sweeps

Returns:

the inference program

Return type:

q

dais(targets, momentum, leapfrog, refreshment, *, num_targets=None)[source]

Differentiable Annealed Importance Sampling.

  1. MCMC Variational Inference via Uncorrected Hamiltonian Annealing, Tomas Geffner, Justin Domke https://arxiv.org/abs/2107.04150

  2. Differentiable Annealed Importance Sampling and the Perils of Gradient Noise, Guodong Zhang, Kyle Hsu, Jianing Li, Chelsea Finn, Roger Grosse https://arxiv.org/abs/2107.10211

Parameters:

  • targets – a list of target programs

  • momentum – the momentum program which calculates kinetic energy

  • leapfrog – the program which performs leapfrog update

  • refreshment – the momentum refreshment program

  • num_targets – the number of targets

Returns:

the inference program

Return type:

q

nasmc(targets, proposals, *, num_targets=None)[source]

Neural Adaptive Sequential Monte Carlo.

  1. Neural Adaptive Sequential Monte Carlo, Shixiang Gu, Zoubin Ghahramani, Richard E. Turner https://arxiv.org/abs/1506.03338

Parameters:

  • targets – a list of target programs

  • proposals – a list of proposal programs

  • num_targets – the number of targets

Returns:

the inference program

Return type:

q

nvi_avo(targets, forwards, reverses, *, num_targets=None)[source]

AIS with Annealed Variational Objective.

  1. Improving Explorability in Variational Inference with Annealed Variational Objectives, Chin-Wei Huang, Shawn Tan, Alexandre Lacoste, Aaron Courville https://arxiv.org/abs/1809.01818

Parameters:

  • targets – a list of target programs

  • forwards – a list of forward kernels

  • reverses – a list of reverse kernels

  • num_targets – the number of targets

Returns:

the inference program

Return type:

q

nvi_fkl(targets, proposals, *, num_targets=None)[source]

Nested Variational Inference with forward KL objective.

Note: The implementation assumes that targets are smoothing distributions. This is different from nasmc, where we assume that the targets are filtering distributions. We also assume that the final target does not have parameters.

  1. Nested Variational Inference, Heiko Zimmermann, Hao Wu, Babak Esmaeili, Jan-Willem van de Meent https://arxiv.org/abs/2106.11302

Parameters:

  • targets – a list of target programs

  • proposals – the proposal for the initial target

  • num_targets – the number of targets

Returns:

the inference program

Return type:

q

nvi_rkl(targets, forwards, reverses, *, num_targets=None)[source]

Nested Variational Inference with reverse KL objective.

If targets is a callable which takes an integer index and returns the i-th taget, we will use fori_loop combinator to improve the compiling time. This requires num_targets to be a concrete value.

Note: In nested VI, we typically assume that the final target does not have parameters. This allows us to optimize intermediate KLs to bridge from the initial target to the final target. Here we use ELBO loss in the last step to also maximize likelihood in case there are parameters in the final target.

  1. Nested Variational Inference, Heiko Zimmermann, Hao Wu, Babak Esmaeili, Jan-Willem van de Meent https://arxiv.org/abs/2106.11302

Parameters:

  • targets – a list of target programs

  • forwards – a list of forward kernels

  • reverses – a list of reverse kernels

  • num_targets – the number of targets

Returns:

the inference program

Return type:

q

rws(target, proposal)[source]

Reweighted Wake-Sleep.

  1. Reweighted Wake-Sleep, Jörg Bornschein, Yoshua Bengio https://arxiv.org/abs/1406.2751

  2. Revisiting Reweighted Wake-Sleep for Models with Stochastic Control Flow, Tuan Anh Le, Adam R. Kosiorek, N. Siddharth, Yee Whye Teh, Frank Wood https://arxiv.org/abs/1805.10469

Parameters:

  • target – the target program

  • proposal – the proposal program

Returns:

the inference program

Return type:

q

svi(target, proposal)[source]

Stochastic Variational Inference.

  1. Auto-Encoding Variational Bayes, Diederik P Kingma, Max Welling https://arxiv.org/abs/1312.6114

  2. Stochastic Backpropagation and Approximate Inference in Deep Generative Models, Danilo Jimenez Rezende, Shakir Mohamed, Daan Wierstra https://arxiv.org/abs/1401.4082

Parameters:

  • target – the target program

  • proposal – the proposal program

Returns:

the inference program

Return type:

q

svi_iwae(target, proposal)[source]

SVI with Important Weighted Autoencoder objective.

  1. Importance Weighted Autoencoders, Yuri Burda, Roger Grosse, Ruslan Salakhutdinov https://arxiv.org/abs/1509.00519

Parameters:

  • target – the target program

  • proposal – the proposal program

Returns:

the inference program

Return type:

q

svi_stl(target, proposal)[source]

SVI with Sticking-the-Landing objective.

  1. Sticking the Landing: Simple, Lower-Variance Gradient Estimators for Variational Inference, Geoffrey Roeder, Yuhuai Wu, David Duvenaud https://arxiv.org/abs/1703.09194

Parameters:

  • target – the target program

  • proposal – the proposal program

Returns:

the inference program

Return type:

q

vsmc(targets, proposals, *, num_targets=None)[source]

Variational Sequential Monte Carlo.

Note: Here, we assume that the dimension of variables is constant (modulo masking) during SMC steps. The targets can be filtering distributions or smoothing distributions (as in [4]).

  1. Filtering Variational Objectives, Chris J. Maddison, Dieterich Lawson, George Tucker, Nicolas Heess, Mohammad Norouzi, Andriy Mnih, Arnaud Doucet, Yee Whye Teh https://arxiv.org/abs/1705.09279

  2. Auto-Encoding Sequential Monte Carlo, Tuan Anh Le, Maximilian Igl, Tom Rainforth, Tom Jin, Frank Wood https://arxiv.org/abs/1705.10306

  3. Variational Sequential Monte Carlo, Christian A. Naesseth, Scott W. Linderman, Rajesh Ranganath, David M. Blei https://arxiv.org/abs/1705.11140

  4. Twisted Variational Sequential Monte Carlo, Dieterich Lawson, George Tucker, Christian A Naesseth, Chris J Maddison, Ryan P Adams, Yee Whye Teh http://bayesiandeeplearning.org/2018/papers/111.pdf

Parameters:

  • targets – a list of target programs

  • proposals – a list of proposal programs

  • num_targets – the number of targets

Returns:

the inference program

Return type:

q