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

Computer Science > Machine Learning arXiv:2506.05668 (cs) [Submitted on 6 Jun 2025 (v1), last revised 2 Mar 2026 (this version, v5)] Title:RNE: plug-and-play diffusion inference-time control and energy-based training Authors:Jiajun He, José Miguel Hernández-Lobato, Yuanqi Du, Francisco Vargas View a PDF of the paper titled RNE: plug-and-play diffusion inference-time control and energy-based training, by Jiajun He and 3 other authors View PDF 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...