Selected Publications

1. TMLR

   A Variational Perspective on Generative Flow Networks

   Zimmermann, Heiko, Lindsten, Fredrik, Meent, Jan-Willem, and Naesseth, Christian A

   Transactions on Machine Learning Research Apr 2023

   Abs HTML PDF Code

   Generative flow networks (GFNs) are a class of probabilistic models for sequential sampling of composite objects, proportional to a target distribution that is defined in terms of an energy function or a reward. GFNs are typically trained using a flow matching or trajectory balance objective, which matches forward and backward transition models over trajectories. In this work we introduce a variational objective for training GFNs, which is a convex combination of the reverse- and forward KL divergences, and compare it to the trajectory balance objective when sampling from the forward- and backward model, respectively. We show that, in certain settings, variational inference for GFNs is equivalent to minimizing the trajectory balance objective, in the sense that both methods compute the same score-function gradient. This insight suggests that in these settings, control variates, which are commonly used to reduce the variance of score-function gradient estimates, can also be used with the trajectory balance objective. We evaluate our findings and the performance of the proposed variational objective numerically by comparing it to the trajectory balance objective on two synthetic tasks.

2. NeurIPS

   Nested Variational Inference

   Zimmermann, Heiko, Wu, Hao, Esmaeili, Babak, and Meent, Jan-Willem

   In 35th Conference on Neural Information Processing Systems (NeurIPS) Dec 2021

   Abs PDF

   We develop nested variational inference (NVI), a family of methods that learn proposals for nested importance samplers by minimizing an forward or reverse KL divergence at each level of nesting. NVI is applicable to many commonly-used importance sampling strategies and provides a mechanism for learning intermediate densities, which can serve as heuristics to guide the sampler. Our experiments apply NVI to (a) sample from a multimodal distribution using a learned annealing path (b) learn heuristics that approximate the likelihood of future observations in a hidden Markov model and (c) to perform amortized inference in hierarchical deep generative models. We observe that optimizing nested objectives leads to improved sample quality in terms of log average weight and effective sample size.

3. UAI

   Learning Proposals for Probabilistic Programs with Inference Combinators

   Zimmermann, Heiko, Stites, Sam, Wu, Hao, Sennesh, Eli, and Meent, Jan-Willem

   In 37th Conference on Uncertainty in Artificial Intelligence (UAI) Jul 2021

   Abs PDF

   We develop operators for construction of proposals in probabilistic programs, which we refer to as inference combinators. Inference combinators define a grammar over importance samplers that compose primitive operations such as application of a transition kernels and importance resampling. Proposals in these samplers can be parameterized using neural networks, which in turn can be trained by optimizing variational objectives. The result is a framework for user-programmable variational methods that are correct by construction and can be tailored to specific models. We demonstrate the flexibility of this framework in applications to advanced variational methods based on amortized Gibbs sampling and annealing.

4. ICML

   Amortized Population Gibbs Samplers with Neural Sufficient Statistics

   Wu, Hao, Zimmermann, Heiko, Sennesh, Eli, Le, Tuan Anh, and Meent, Jan-Willem

   In Proceeding of the International Conference on Machine Learning (ICML) Jul 2020

   Abs PDF Code

   Amortized variational methods have proven difficult to scale to structured problems, such as inferring positions of multiple objects from video images. We develop amortized population Gibbs (APG) samplers, a class of scalable methods that frames structured variational inference as adaptive importance sampling. APG samplers construct high-dimensional proposals by iterating over updates to lower-dimensional blocks of variables. We train each conditional proposal by minimizing the inclusive KL divergence with respect to the conditional posterior. To appropriately account for the size of the input data, we develop a new parameterization in terms of neural sufficient statistics. Experiments show that APG samplers can train highly structured deep generative models in an unsupervised manner, and achieve substantial improvements in inference accuracy relative to standard autoencoding variational methods.