Risk-sensitive Diffusion

This repository contains a Pytorch implementation of the risk-sensitive diffusion model as introduced in this ICLR-2025 paper.

Running

This codebase is built on EDM, with almost the same use cases. Take the model training as an example:

export CUDA_VISIBLE_DEVICES=1,2
torchrun --standalone --nproc_per_node=2 train.py \
--outdir=${SAVE_DIR} \
--data=${DATA_DIR} \
--cond=0 \
--arch=ncsnpp \
--precond=ve \
--augment=0.0 \
--tick=20 \
--snap=100 \
--dump=200 \
--risk_type=gauss \
--cbase=64 \
--cres=2 \
--blocks=8 \
--duration=20 \
--batch=1024 \

Citation

@inproceedings{
li2025risksensitive,
title={Risk-Sensitive Diffusion: Robustly Optimizing Diffusion Models with Noisy Samples},
author={Yangming Li and Max Ruiz Luyten and Mihaela van der Schaar},
booktitle={The Thirteenth International Conference on Learning Representations},
year={2025},
url={https://openreview.net/forum?id=b0WpXBABdu}
}