Radon-Nikodym Estimator (RNE)

This repo contains code for RNE: plug-and-play diffusion inference-time control and energy-based training (ICLR 2026).

This repo is currently under construction, and we will release the complete code ASAP. Sorry for the delay.

🗓️ Plans

• Tutorial notebook for density estimator;

• Tutorial notebook for diffusion control;

• Tutorial notebook for energy-based training;

• Annealing application;

• Energy-based training application;

• Reward-tilting application.

🎓 Notebook and Tutorial

To aid understanding of how to implement our method, we provide several notebooks in Notebooks, covering:

(1) density estimation;

(2) diffusion control;

(3) energy-based training.

💻 Experiments

Applications
├── anneal/      # RNE tempering
├── energy_based_training/     # Energy-based training with RNE regularization
├── stitch/     # (under construction)
└── ctmc/       # (under construction)

👉🏻Test-time Annealing on Boltzmann distributions

👉🏻Test-time Reward-tilting with CTMC

👉🏻Test-time stitching for maze navigation

👉🏻Energy-based diffusion training on Boltzmann distributions (Alanine Dipeptide)

Please email jh2383@cam.ac.uk if you have more questions.

📑 Reference

If you find this repository useful, please consider citing:

@inproceedings{he2025rne,
  title={RNE: plug-and-play diffusion inference-time control and energy-based training},
  author={He, Jiajun and Hern{\'a}ndez-Lobato, Jos{\'e} Miguel and Du, Yuanqi and Vargas, Francisco},
  booktitle={ICLR},
  year={2025}
}