Optimization of Annealed Importance Sampling Hyperparameters

Annealed Importance Sampling (AIS) is a popular algorithm used to estimates the intractable marginal likelihood of deep generative models. Although AIS is guaranteed to provide unbiased estimate for any set of hyperparameters, the common implementations rely on simple heuristics such as the geometric average bridging distributions between initial and the target distribution which affect the estimation performance when the computation budget is limited. Optimization of fully parametric AIS remains challenging due to the use of Metropolis-Hasting (MH) correction steps in Markov transitions. We present a parameteric AIS process with flexible intermediary distributions and optimize the bridging distributions to use fewer number of steps for sampling. A reparameterization method that allows us to optimize the distribution sequence and the parameters of Markov transitions is used which is applicable to a large class of Markov Kernels with MH correction. We assess the performance of our optimized AIS for marginal likelihood estimation of deep generative models and compare it to other estimators.