RNE: a plug-and-play framework for diffusion density estimation and inference-time control

Abstract: In this paper, we introduce the Radon-Nikodym Estimator (RNE), a flexible, plug-and-play framework for diffusion inference-time density estimation and control, based on the concept of the density ratio between path distributions. RNE connects and unifies a variety of existing density estimation and inference-time control methods under a single and intuitive perspective, stemming from basic variational inference and probabilistic principles therefore offering both theoretical clarity and practical versatility. Experiments demonstrate that RNE delivers strong results in diffusion density estimation, and offers broad applicability to inference-time control tasks -- such as annealing, diffusion model composition, and reward-tilting -- with promising inference-time scaling performance.