# Differentiable Annealed Importance Sampling and the Perils of Gradient Noise

arXiv: https://arxiv.org/abs/2107.10211 · HTML: https://mlsanity.com/paper/2107.10211/

## Key facts

- Cited by 36 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 53 references that are in the map.
- First posted on arXiv 21 Jul 2021.

## Builds on (most cited first, 28 of 28)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Importance Weighted Autoencoders](https://mlsanity.com/paper/1509.00519/) (2015)
- [Gradient-based Hyperparameter Optimization through Reversible Learning](https://mlsanity.com/paper/1502.03492/) (2015)
- [Stochastic Gradient Hamiltonian Monte Carlo](https://mlsanity.com/paper/1402.4102/) (2014)
- [Stochastic Gradient Descent as Approximate Bayesian Inference](https://mlsanity.com/paper/1704.04289/) (2017)
- [Scalable Gradients for Stochastic Differential Equations](https://mlsanity.com/paper/2001.01328/) (2020)
- [Non-convex learning via Stochastic Gradient Langevin Dynamics: a nonasymptotic analysis](https://mlsanity.com/paper/1702.03849/) (2017)
- [Theoretical guarantees for approximate sampling from smooth and log-concave densities](https://mlsanity.com/paper/1412.7392/) (2014)
- [A Complete Recipe for Stochastic Gradient MCMC](https://mlsanity.com/paper/1506.04696/) (2015)
- [Underdamped Langevin MCMC: A non-asymptotic analysis](https://mlsanity.com/paper/1707.03663/) (2017)
- [Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model](https://mlsanity.com/paper/1907.04164/) (2019)
- [Variational Sequential Monte Carlo](https://mlsanity.com/paper/1705.11140/) (2017)
- [Noisy Natural Gradient as Variational Inference](https://mlsanity.com/paper/1712.02390/) (2017)
- [Filtering Variational Objectives](https://mlsanity.com/paper/1705.09279/) (2017)
- [Consistency and fluctuations for stochastic gradient Langevin dynamics](https://mlsanity.com/paper/1409.0578/) (2014)
- [Auto-Encoding Sequential Monte Carlo](https://mlsanity.com/paper/1705.10306/) (2017)
- [Hamiltonian Variational Auto-Encoder](https://mlsanity.com/paper/1805.11328/) (2018)
- [On the Quantitative Analysis of Decoder-Based Generative Models](https://mlsanity.com/paper/1611.04273/) (2016)
- [MCMC Variational Inference via Uncorrected Hamiltonian Annealing](https://mlsanity.com/paper/2107.04150/) (2021)
- [Monte Carlo Variational Auto-Encoders](https://mlsanity.com/paper/2106.15921/) (2021)
- [On the Convergence of Stochastic Gradient MCMC Algorithms with High-Order Integrators](https://mlsanity.com/paper/1610.06665/) (2016)
- [Faster Convergence of Stochastic Gradient Langevin Dynamics for Non-Log-Concave Sampling](https://mlsanity.com/paper/2010.09597/) (2020)
- [The Thermodynamic Variational Objective](https://mlsanity.com/paper/1907.00031/) (2019)
- [Reparameterization Gradients through Acceptance-Rejection Sampling Algorithms](https://mlsanity.com/paper/1610.05683/) (2016)
- [Sandwiching the marginal likelihood using bidirectional Monte Carlo](https://mlsanity.com/paper/1511.02543/) (2015)
- [Improving Lossless Compression Rates via Monte Carlo Bits-Back Coding](https://mlsanity.com/paper/2102.11086/) (2021)
- [Evaluating Lossy Compression Rates of Deep Generative Models](https://mlsanity.com/paper/2008.06653/) (2020)
- [SUMO: Unbiased Estimation of Log Marginal Probability for Latent Variable Models](https://mlsanity.com/paper/2004.00353/) (2020)

## Built on by (most cited first, 30 of 36)

- [Denoising Diffusion Samplers](https://mlsanity.com/paper/2302.13834/) (2023)
- [Flow Annealed Importance Sampling Bootstrap](https://mlsanity.com/paper/2208.01893/) (2022)
- [Continual Repeated Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2201.13117/) (2022)
- [Bayesian Model Selection, the Marginal Likelihood, and Generalization](https://mlsanity.com/paper/2202.11678/) (2022)
- [Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling](https://mlsanity.com/paper/2406.07423/) (2024)
- [Score-Based Diffusion meets Annealed Importance Sampling](https://mlsanity.com/paper/2208.07698/) (2022)
- [Underdamped Diffusion Bridges with Applications to Sampling](https://mlsanity.com/paper/2503.01006/) (2025)
- [Langevin Diffusion Variational Inference](https://mlsanity.com/paper/2208.07743/) (2022)
- [Liouville Flow Importance Sampler](https://mlsanity.com/paper/2405.06672/) (2024)
- [Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference](https://mlsanity.com/paper/2508.12511/) (2025)
- [Transport meets Variational Inference: Controlled Monte Carlo Diffusions](https://mlsanity.com/paper/2307.01050/) (2023)
- [Deep Generative Models for Detector Signature Simulation: A Taxonomic Review](https://mlsanity.com/paper/2312.09597/) (2023)
- [Surrogate Likelihoods for Variational Annealed Importance Sampling](https://mlsanity.com/paper/2112.12194/) (2021)
- [Bayesian inference via sparse Hamiltonian flows](https://mlsanity.com/paper/2203.05723/) (2022)
- [Designing Perceptual Puzzles by Differentiating Probabilistic Programs](https://mlsanity.com/paper/2204.12301/) (2022)
- [MixFlows: principled variational inference via mixed flows](https://mlsanity.com/paper/2205.07475/) (2022)
- [Annealing Flow Generative Models Towards Sampling High-Dimensional and Multi-Modal Distributions](https://mlsanity.com/paper/2409.20547/) (2024)
- [Provable Convergence and Limitations of Geometric Tempering for Langevin Dynamics](https://mlsanity.com/paper/2410.09697/) (2024)
- [Adaptive Annealed Importance Sampling with Constant Rate Progress](https://mlsanity.com/paper/2306.15283/) (2023)
- [Markov Chain Score Ascent: A Unifying Framework of Variational Inference with Markovian Gradients](https://mlsanity.com/paper/2206.06295/) (2022)
- [Differentiating Metropolis-Hastings to Optimize Intractable Densities](https://mlsanity.com/paper/2306.07961/) (2023)
- [Embracing the chaos: analysis and diagnosis of numerical instability in variational flows](https://mlsanity.com/paper/2307.06957/) (2023)
- [Neural Flow Samplers with Shortcut Models](https://mlsanity.com/paper/2502.07337/) (2025)
- [Variational Inference with Locally Enhanced Bounds for Hierarchical Models](https://mlsanity.com/paper/2203.04432/) (2022)
- [Ensemble-Based Annealed Importance Sampling](https://mlsanity.com/paper/2401.15645/) (2024)
- [Variational Learning of Gaussian Process Latent Variable Models through Stochastic Gradient Annealed Importance Sampling](https://mlsanity.com/paper/2408.06710/) (2024)
- [End-To-End Learning of Gaussian Mixture Priors for Diffusion Sampler](https://mlsanity.com/paper/2503.00524/) (2025)
- [Resampling Gradients Vanish in Differentiable Sequential Monte Carlo Samplers](https://mlsanity.com/paper/2304.14390/) (2023)
- [Reparameterized Variational Rejection Sampling](https://mlsanity.com/paper/2309.14612/) (2023)
- [Denoising Diffusion Variational Inference: Diffusion Models as Expressive Variational Posteriors](https://mlsanity.com/paper/2401.02739/) (2024)

## Abstract (arXiv)

Annealed importance sampling (AIS) and related algorithms are highly effective tools for marginal likelihood estimation, but are not fully differentiable due to the use of Metropolis-Hastings correction steps. Differentiability is a desirable property as it would admit the possibility of optimizing marginal likelihood as an objective using gradient-based methods. To this end, we propose Differentiable AIS (DAIS), a variant of AIS which ensures differentiability by abandoning the Metropolis-Hastings corrections. As a further advantage, DAIS allows for mini-batch gradients. We provide a detailed convergence analysis for Bayesian linear regression which goes beyond previous analyses by explicitly accounting for the sampler not having reached equilibrium. Using this analysis, we prove that DAIS is consistent in the full-batch setting and provide a sublinear convergence rate. Furthermore, motivated by the problem of learning from large-scale datasets, we study a stochastic variant of DAIS that uses mini-batch gradients. Surprisingly, stochastic DAIS can be arbitrarily bad due to a fundamental incompatibility between the goals of last-iterate convergence to the posterior and elimination of the accumulated stochastic error. This is in stark contrast with other settings such as gradient-based optimization and Langevin dynamics, where the effect of gradient noise can be washed out by taking smaller steps. This indicates that annealing-based marginal likelihood estimation with stochastic gradients may require new ideas.