Improving Single Noise Level Diffusion Samplers with Restricted Gaussian Oracles

ICLR 2025 DeLTa Workshop Poster

---

Abstract

Diffusion models and diffusion Monte-Carlo schemes that sample from unnormalized log-densities, both rely on denoisers ( or score estimates) at different noise scales. This complicates the sampling process as denoising schedules require careful tuning and nested inner-MCMC loops. In this work, we propose a single noise level sampling procedure that only requires a single low-noise denoiser. Our framework results from improvements we bring to the multimeasurement Walk-Jump sampler of Saremi et al. 2021 by mixing in ideas from the proximal sampler of Shen et al. 2020. Our analysis shows that annealing (or multiple noise scales) is unnecessary if one is willing to pay an increased memory cost. We demonstrate this by proposing an entirely log-concave sampling framework.

Instructions

Obtaining the code

Clone this repository by running:

git clone https://github.com/Bani57/multi-prox-diffusion-iclr-delta-2025.git

Dependencies

This repository extends the implementation of "DiGress: Discrete Denoising diffusion models for graph generation", originally available at https://github.com/cvignac/DiGress.

To install the Python programming language and the dependent Python packages, you need to execute the following steps:

• Download anaconda/miniconda if needed

• Create a rdkit environment that directly contains rdkit:

  conda create -c conda-forge -n digress rdkit=2023.03.2 python=3.9

• conda activate digress

• Check that this line does not return an error:

  python3 -c 'from rdkit import Chem'

• Install graph-tool (https://graph-tool.skewed.de/):

  conda install -c conda-forge graph-tool=2.45

• Check that this line does not return an error:

  python3 -c 'import graph_tool as gt'

• Install the nvcc drivers for your cuda version. For example:

  conda install -c "nvidia/label/cuda-11.8.0" cuda

• Install a corresponding version of pytorch, for example:

  pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118

• Install other packages using the requirement file:

  pip install -r requirements.txt

• Run:

  pip install -e .

• Navigate to the src/analysis/orca directory and compile orca.cpp:

  g++ -O2 -std=c++11 -o orca orca.cpp

Reproducing results

To generate molecules with our Multi-prox sampler and obtain example chains like the one displayed in our work, run the following command from inside the src directory:

CUDA_VISIBLE_DEVICES=<GPU_ID> python main.py dataset="qm9" general.name="c_gibbs_startNone_N10_M100_t0.5_0.1" model="continuous" general.test_only="<REPO_ABSPATH>/checkpoints/qm9_c.ckpt" +model.gibbs=True +model.gibbs_N=10 +model.gibbs_M=100 +model.gibbs_fixed_t_2=0.5 +model.gibbs_fixed_t_1=0.1 +model.gibbs_chain_freq=0.1 ++general.final_model_samples_to_generate=1000

where you would replace <GPU_ID> with the integer ID of one of your available GPU devices and replace <REPO_ABSPATH> with the absolute path of the main repo directory (where you cloned it).

The batch of full molecule chains will be saved to <REPO_ABSPATH>/chains/c_gibbs_startNone_N10_M100_t0.5_0.1_resume/epoch0/chains/, while the batch of just the final molecules will be saved to <REPO_ABSPATH>/graphs/c_gibbs_startNone_N10_M100_t0.5_0.1_resume/epoch0_b0/.

To reproduce our results with the denoising diffusion GAN for images, run these commands from inside the src directory:

• CIFAR 10: CUDA_VISIBLE_DEVICES=<GPU_ID> python main_cv.py --dataset cifar10 --exp ddgan_cifar10_exp1 --num_channels 3 --num_channels_dae 128 --num_timesteps 4 --num_res_blocks 2 --nz 100 --z_emb_dim 256 --n_mlp 4 --ch_mult 1 2 2 2 --epoch_id 1200 --n 1 --m 10 --fixed_t 0.25 --sample_freq 1

• CelebA-256: CUDA_VISIBLE_DEVICES=<GPU_ID> python main_cv.py --dataset celeba_256 --image_size 256 --batch_size 50 --exp ddgan_celebahq_exp1 --num_channels 3 --num_channels_dae 64 --ch_mult 1 1 2 2 4 4 --num_timesteps 2 --num_res_blocks 2 --epoch_id 550 --n 1 --m 10 --fixed_t 0.5 --sample_freq 1

The generated image chains will be saved to <REPO_ABSPATH>/generated_chains and the individual images to <REPO_ABSPATH>/generated_samples.

To reproduce our Mixture of Gaussians RGO example, run the MoG40.ipynb notebook.

---

Citation

If you use this code for your projects, please cite:

@inproceedings{
dadi2025improving,
title={Improving Single Noise Level Denoising Samplers with Restricted Gaussian Oracles},
author={Leello Tadesse Dadi and Andrej Janchevski and Volkan Cevher},
booktitle={ICLR 2025 Workshop on Deep Generative Model in Machine Learning: Theory, Principle and Efficacy},
year={2025},
url={https://openreview.net/forum?id=xkiI5tou6J}
}

---

Authors

• Leello Tadesse Dadi, EPFL STI IEM LIONS, Lausanne, Switzerland, leello.dadi@epfl.ch
• Andrej Janchevski, EPFL STI IEM LIONS, Lausanne, Switzerland, andrej.janchevski@epfl.ch
• Volkan Cevher, EPFL STI IEM LIONS, Lausanne, Switzerland, volkan.cevher@epfl.ch