# StAD: Stein Amortized Divergence for Fast Likelihoods with Diffusion and Flow

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 63 of its 102 references that are in the map.
- First posted on arXiv 15 May 2026.

## Builds on (most cited first, 30 of 63)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [FiLM: Visual Reasoning with a General Conditioning Layer](https://mlsanity.com/paper/1709.07871/) (2017)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Normalizing Flows for Probabilistic Modeling and Inference](https://mlsanity.com/paper/1912.02762/) (2019)
- [Diffusion Models: A Comprehensive Survey of Methods and Applications](https://mlsanity.com/paper/2209.00796/) (2022)
- [Improved Techniques for Training Score-Based Generative Models](https://mlsanity.com/paper/2006.09011/) (2020)
- [One-step Diffusion with Distribution Matching Distillation](https://mlsanity.com/paper/2311.18828/) (2023)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Adversarial Diffusion Distillation](https://mlsanity.com/paper/2311.17042/) (2023)
- [Maximum Likelihood Training of Score-Based Diffusion Models](https://mlsanity.com/paper/2101.09258/) (2021)
- [FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models](https://mlsanity.com/paper/1810.01367/) (2018)
- [Mean Flows for One-step Generative Modeling](https://mlsanity.com/paper/2505.13447/) (2025)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Improved Techniques for Training Consistency Models](https://mlsanity.com/paper/2310.14189/) (2023)
- [Analyzing and Improving the Training Dynamics of Diffusion Models](https://mlsanity.com/paper/2312.02696/) (2023)
- [InstaFlow: One Step is Enough for High-Quality Diffusion-Based Text-to-Image Generation](https://mlsanity.com/paper/2309.06380/) (2023)
- [A Downsampled Variant of ImageNet as an Alternative to the CIFAR datasets](https://mlsanity.com/paper/1707.08819/) (2017)
- [Knowledge Distillation in Iterative Generative Models for Improved Sampling Speed](https://mlsanity.com/paper/2101.02388/) (2021)
- [Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation](https://mlsanity.com/paper/2303.00848/) (2023)
- [Flow Matching Guide and Code](https://mlsanity.com/paper/2412.06264/) (2024)
- [A note on the evaluation of generative models](https://mlsanity.com/paper/1511.01844/) (2015)
- [Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models](https://mlsanity.com/paper/2410.11081/) (2024)
- [A Kernelized Stein Discrepancy for Goodness-of-fit Tests and Model Evaluation](https://mlsanity.com/paper/1602.03253/) (2016)

## Abstract (arXiv)

Diffusion and flow-based models are ubiquitously used for generative modelling and density estimation. They admit a deterministic probability flow ordinary differential equation (PF-ODE), analogous to continuous normalizing flows (CNFs), which describes the transport of the probability mass. Obtaining the likelihood from these models is of interest to many workflows, especially Bayesian analysis, and requires solving the trace of the Jacobian to compute the divergence of the learned PF-ODE, which is either O(D2) to compute exactly or O(D) with a noisy estimate. We introduce StAD, a new distillation method to predict and learn the divergence of the PF-ODE using the Langevin-Stein operator without ever computing the Jacobian. We show that our method is competitive with the Hutchinson and Hutch++ on CIFAR-10, ImageNet and other density estimation tasks, consistently improving the variance and speed of the likelihood predictions compared to the Hutchinson. We additionally show our method will generalize to a varied class of generative models, and show that under some regularity conditions these learned vector fields can be made to satisfy the Stein class.