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

2026-05-15

AI Summary

This work addresses the challenge of efficiently and accurately estimating the divergence of the probability flow ordinary differential equation (PF-ODE) in diffusion and flow-based generative models, where existing approaches are either computationally expensive or suffer from high variance. The authors propose StAD, a novel method that, for the first time, incorporates the Langevin–Stein operator into divergence distillation, enabling accurate learning of the PF-ODE divergence without explicit Jacobian computation. They theoretically show that the learned vector field belongs to the Stein class under suitable conditions. By combining function approximation with regularization techniques, StAD significantly reduces estimation variance and accelerates likelihood evaluation on benchmarks such as CIFAR-10 and ImageNet, while demonstrating broad applicability across diverse generative modeling frameworks.

Abstract

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 $\mathcal{O}(D^2)$ to compute exactly or $\mathcal{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.

Problem

• diffusion models
• flow-based models
• likelihood estimation
• Jacobian trace
• divergence computation

Innovation

• Stein Amortized Divergence
• PF-ODE
• Jacobian-free divergence estimation
• diffusion models
• likelihood distillation