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

TL;DR

The paper introduces RNE, a unified framework for diffusion density estimation, inference-time control, and energy-based training, applicable to both continuous and discrete diffusion models.

Contribution

RNE provides a novel, flexible plug-and-play approach based on density ratios, connecting marginal densities and transition kernels for diffusion models.

Findings

RNE achieves strong inference-time control results in annealing and model composition.

RNE demonstrates promising scalability for inference tasks.

RNE offers an efficient regularization method for energy-based diffusion training.

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 \textit{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 demonstrate that RNE delivers strong results in inference-time control applications, such as annealing and model composition, with promising inference-time scaling performance, and achieves a 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.