RNE: plug-and-play diffusion inference-time control and energy-based training
For practitioners using diffusion models, RNE offers a unified method for density estimation, inference-time control, and energy-based training, but the improvements are incremental over existing techniques.
The paper introduces the Radon-Nikodym Estimator (RNE), a plug-and-play framework that estimates marginal densities along diffusion trajectories, enabling inference-time control and energy-based training. RNE achieves strong results in annealing and model composition with promising inference-time scaling, and provides efficient regularization for energy-based diffusion models.
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.