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

2025-06-06

AI Summary

Diffusion models face limitations in density estimation and inference-time control. To address this, we propose the Radon–Nikodym Estimator (RNE), a plug-and-play framework grounded in path-distribution density ratios. RNE unifies variational objectives and path distributions from first probabilistic principles, ensuring theoretical consistency while enabling diverse inference-time controls—including annealing, model composition, and reward-based reweighting—without architectural modification or retraining. Unlike existing approaches, RNE achieves flexible, controllable generation solely through density ratio estimation, decoupling control from sampling dynamics. Experiments demonstrate that RNE significantly outperforms baselines across multiple density estimation benchmarks and diverse control tasks. Notably, it is the first method to jointly model density estimation and general-purpose inference-time control within a single, theoretically coherent, and computationally efficient framework.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Intelligent Robots: State Estimation
• Machine Learning: Deep Generative Models & Autoencoders

Application Category

• Economics, Online Markets and Human Computation: LLM based quality controls for crowd work
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

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 achieves promising performances in diffusion density estimation and inference-time control tasks, including annealing, composition of diffusion models, and reward-tilting.

Problem

• Estimating diffusion model density ratios
• Unifying density estimation and control methods
• Enhancing inference-time tasks like annealing and reward-tilting

Innovation

• Plug-and-play framework for diffusion density estimation
• Unifies density estimation and control methods
• Based on density ratio between path distributions