AdaAnn: Adaptive Annealing Scheduler for Probability Density Approximation
This work addresses a computational bottleneck in probability density approximation for researchers and practitioners using methods like variational inference and MCMC, though it is incremental as it builds on existing annealing techniques.
The paper tackles the challenge of approximating complex probability distributions by introducing AdaAnn, an adaptive annealing scheduler that automatically adjusts temperature increments based on expected KL divergence changes, resulting in improved computational efficiency for tasks like density approximation and parameter estimation in dynamical systems.
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.