AdaAnn: Adaptive Annealing Scheduler for Probability Density Approximation

TL;DR

AdaAnn is an adaptive annealing scheduler that improves the efficiency of probability density approximation by dynamically adjusting temperature increments based on KL divergence, enhancing existing sampling methods.

Contribution

It introduces AdaAnn, an automatic, adaptive annealing scheduler that optimizes temperature increments for density approximation tasks, compatible with various sampling approaches.

Findings

AdaAnn improves computational efficiency in density approximation.

It integrates seamlessly with normalizing flows and MCMC methods.

Demonstrated effectiveness on dynamical systems and variational inference.

Abstract

Approximating probability distributions can be a challenging task, particularly when they are supported over regions of high geometrical complexity or exhibit multiple modes. Annealing can be used to facilitate this task which is often combined with constant a priori selected increments in inverse temperature. However, using constant increments limit the computational efficiency due to the inability to adapt to situations where smooth changes in the annealed density could be handled equally well with larger increments. We introduce AdaAnn, an adaptive annealing scheduler that automatically adjusts the temperature increments based on the expected change in the Kullback-Leibler divergence between two distributions with a sufficiently close annealing temperature. AdaAnn is easy to implement and can be integrated into existing sampling approaches such as normalizing flows for variational inference and Markov chain Monte Carlo. We demonstrate the computational efficiency of the AdaAnn scheduler for variational inference with normalizing flows on a number of examples, including density approximation and parameter estimation for dynamical systems.