RNE: plug-and-play diffusion inference-time control and energy-based training

Jiajun He, José Miguel Hernández-Lobato, Yuanqi Du, Francisco Vargas

Keywords: diffusion generative models, SDE, SMC, sequential monte carlo

Abstract:

Diffusion models generate data by removing noise gradually, which corresponds to the time-reversal of a noising process.
However, access to only the denoising kernels is often insufficient.
In many applications, we need the knowledge of the marginal densities along the generation trajectory, which enables tasks such as inference-time control.
To address this gap, in this paper, we introduce the Radon-Nikodym Estimator (RNE).
Based on the concept of the density ratio between path distributions, it reveals a fundamental connection between marginal densities and transition kernels, providing a flexible plug-and-play framework that unifies (1) diffusion density estimation, (2) inference-time control, and (3) energy-based diffusion training under a single perspective.
Experiments demonstrated that RNE delivers strong results in inference-time control applications, such as annealing and model composition, with promising inference-time scaling performance, and achieves simple yet efficient regularisation for training energy-based diffusion models.
Additionally, our proposed RNE is modality-agnostic and applicable not only to continuous diffusion models but also to their discrete diffusion counterparts.

Primary Area: generative models

Submission Number: 2224