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

2025/06/06 by Jiajun He, He, Jiajun, José Miguel Hernández-Lobato +5 · 2 citations
Computer Science · Medicine · Physics and Astronomy · #Advanced Neuroimaging Techniques and Applications #FOS: Computer and information sciences #Generative Adversarial Networks and Image Synthesis #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Model Reduction and Neural Networks

paper · pdf · doi:10.48550/arxiv.2506.05668

openalex publication_date 2025/06/06 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

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 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.

Citations

• Consistent Sampling and Simulation: Molecular Dynamics with Energy-Based Diffusion Models
• Learning normalized image densities via dual score matching
• Solving Inverse Problems via Diffusion-Based Priors: An Approximation-Free Ensemble Sampling Approach
• FEAT: Free energy Estimators with Adaptive Transport
• Feynman-Kac Correctors in Diffusion: Annealing, Guidance, and Product of Experts
• Scalable Equilibrium Sampling with Sequential Boltzmann Generators
• Composition and Control with Distilled Energy Diffusion Models and Sequential Monte Carlo
• LEAPS: A discrete neural sampler via locally equivariant networks
• Accelerated Parallel Tempering via Neural Transports
• Debiasing Guidance for Discrete Diffusion with Sequential Monte Carlo
• No Trick, No Treat: Pursuits and Challenges Towards Simulation-free Training of Neural Samplers
• Inference-Time Alignment in Diffusion Models with Reward-Guided Generation: Tutorial and Review
• A General Framework for Inference-time Scaling and Steering of Diffusion Models
• From discrete-time policies to continuous-time diffusion samplers: Asymptotic equivalences and faster training
• The Superposition of Diffusion Models Using the Itô Density Estimator
• Sequential Controlled Langevin Diffusions
• Diffusion Models as Cartoonists: The Curious Case of High Density Regions
• Reprogramming Pretrained Target-Specific Diffusion Models for Dual-Target Drug Design
• Learned Reference-based Diffusion Sampling for multi-modal distributions
• Solvation free energies from neural thermodynamic integration
• Diffusion Density Estimators
• NETS: A Non-Equilibrium Transport Sampler
• Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control
• Diffusion Soup: Model Merging for Text-to-Image Diffusion Models
• Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data
• Simplified and Generalized Masked Diffusion for Discrete Data
• Neural Thermodynamic Integration: Free Energies from Energy-Based Diffusion Models
• DEFT: Efficient Fine-Tuning of Diffusion Models by Learning the Generalised h-transform
• Learning Diffusion Priors from Observations by Expectation Maximization
• Diffusion Models as Constrained Samplers for Optimization with Unknown Constraints
• Target Score Matching
• Particle Denoising Diffusion Sampler
• Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design
• SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers
• OpenMM 8: Molecular Dynamics Simulation with Machine Learning Potentials
• Compositional Foundation Models for Hierarchical Planning
• Accurate transition state generation with an object-aware equivariant elementary reaction diffusion model
• Stochastic Interpolants: A Unifying Framework for Flows and Diffusions
• 3D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Prediction
• Reduce, Reuse, Recycle: Compositional Generation with Energy-Based Diffusion Models and MCMC
• Improving and generalizing flow-based generative models with minibatch optimal transport
• Two for One: Diffusion Models and Force Fields for Coarse-Grained Molecular Dynamics
• Structure-based Drug Design with Equivariant Diffusion Models
• Flow Matching for Generative Modeling
• Diffusion Posterior Sampling for General Noisy Inverse Problems
• Towards Healing the Blindness of Score Matching
• Score-Based Diffusion meets Annealed Importance Sampling
• Flow Annealed Importance Sampling Bootstrap
• Classifier-Free Diffusion Guidance
• Diffusion probabilistic modeling of protein backbones in 3D for the motif-scaffolding problem
• Compositional Visual Generation with Composable Diffusion Models
• Elucidating the Design Space of Diffusion-Based Generative Models
• A Continuous Time Framework for Discrete Denoising Models
• Video Diffusion Models
• Equivariant Diffusion for Molecule Generation in 3D
• High-Resolution Image Synthesis with Latent Diffusion Models
• High-Resolution Image Synthesis with Latent Diffusion Models
• Structured Denoising Diffusion Models in Discrete State-Spaces
• Variational Diffusion Models
• AutoDock Vina 1.2.0: New Docking Methods, Expanded Force Field, and Python Bindings
• A Variational Perspective on Diffusion-Based Generative Models and Score Matching
• Diffusion Models Beat GANs on Image Synthesis
• Improved Denoising Diffusion Probabilistic Models
• Score-Based Generative Modeling through Stochastic Differential Equations
• Denoising Diffusion Implicit Models
• Denoising Diffusion Probabilistic Models
• Neural Ordinary Differential Equations
• Variational Inference: A Review for Statisticians
• Mean Field Forward-Backward Stochastic Differential Equations
• A Connection Between Score Matching and Denoising Autoencoders
• Annealed Importance Sampling
• Annealed importance sampling
• Dynamical Theories of Brownian Motion
• High-Temperature Equation of State by a Perturbation Method. I. Nonpolar Gases
• Statistical Mechanics of Fluid Mixtures

Cited by

• Reinforced sequential Monte Carlo for amortised sampling
• Learning Hamiltonian Flow Maps: Mean Flow Consistency for Large-Timestep Molecular Dynamics

Related

• Improving Adversarial Energy-Based Model via Diffusion Process
• Improved off-policy training of diffusion samplers
• Composition and Control with Distilled Energy Diffusion Models and Sequential Monte Carlo
• Diffusion Models: A Mathematical Introduction
• ρ-Diffusion: A diffusion-based density estimation framework for computational physics