Skip to content

README.md

Latest commit

 

History

History
158 lines (118 loc) · 4.93 KB

README.md

File metadata and controls

158 lines (118 loc) · 4.93 KB

RDSMC

Official codebase for the paper Reverse Diffusion Sequential Monte Carlo Samplers.

Installation

Prerequisites

  • Python 3.13
  • conda

Setup

  1. Create a conda environment named rdsmc with Python 3.13:

    conda create -n rdsmc python=3.13 -y
conda activate rdsmc

  2. Install the package with PyTorch:

    pip install -e .

  3. For evaluation in experiments, install pot package:

    pip install pot

Package Structure

rdsmc/
├── src/
│   ├── rdsmc/              # RDSMC sampler and utilities
│   │   ├── rdsmc_sampler.py    # Main RDSMC sampler with callback support
│   │   └── resample.py         # Resampling utilities
│   ├── diffusion/          # Diffusion noisers
│   │   └── noiser.py
│   ├── posterior_sampler/  # SMC, MCMC, and continual AIS samplers
│   │   ├── smc.py
│   │   ├── mcmc.py
│   │   └── continual_ais.py
│   └── target/             # Target distributions
│       ├── gmm.py
│       ├── logistic_regression.py
│       ├── rings.py
│       └── funnel.py
├── experiments/            # Experiment scripts and configurations
│   ├── run.py              # Main RDSMC experiment script (uses callbacks)
│   ├── run_ais.py          # Main AIS experiment script
│   ├── callbacks.py        # Callback implementations for logging/plotting
│   ├── eval.py             # Evaluation utilities
│   ├── configs/            # Configuration files
│   │   ├── experiment/     # Experiment-specific configs
│   │   ├── target/         # Target distribution configs
│   │   ├── noiser/         # Noiser configs
│   │   ├── wandb/          # Weights & Biases configs
│   │   ├── sampling_rdsmc.yaml  # Base RDSMC config
│   │   └── sampling_ais.yaml    # Base AIS config
│   └── data/               # Dataset files
└── pyproject.toml          # Package configuration

Dependencies

The package requires:

  • torch - PyTorch
  • numpy - Numerical computing
  • omegaconf - Configuration management
  • hydra-core - Configuration framework
  • wandb - Experiment tracking
  • matplotlib - Plotting
  • numba - JIT compilation for resampling

Usage

Running Experiments

RDSMC experiments:

cd experiments
python run.py experiment=gmm2d

You can override hyperparameters using command-line arguments:

python run.py experiment=gmm2d score_approx.ais.n_annealing_steps=10 plot_every=10 log_every=10

Configuration files for all experiments in the paper are available in the experiments/configs/ folder. For example:

python run.py experiment=rings

You can also create your own experiment configuration files following the same structure.

Note: For accurate performance timing, set plot_every=null and log_every=null to disable intermediate plotting and logging, which can add significant overhead.

AIS experiments:

We also provide a script to run baseline methods: Annealed Importance Sampling (AIS) and Sequential Monte Carlo (SMC), which use a geometric annealing schedule.

cd experiments
python run_ais.py experiment=ais_gmm # AIS
python run_ais.py experiment=ais_gmm ais.ess_threshold=1.0 # SMC with resampling at every step
python run_ais.py experiment=ais_gmm ais.ess_threshold=0.3 # SMC with adaptive resampling

AIS experiment configurations are available in experiments/configs/experiment/ (e.g., ais_gmm.yaml, ais_rings.yaml, ais_funnel.yaml).

Configuration

Experiments are configured using Hydra. Configuration files are in experiments/configs/:

  • experiment/ - Experiment-specific configs (e.g., gmm2d.yaml, ais_gmm.yaml)
  • target/ - Target distribution configs
  • noiser/ - Noiser configs (e.g., vpnoiser.yaml)
  • sampling_rdsmc.yaml - Base config for RDSMC experiments
  • sampling_ais.yaml - Base config for AIS experiments

Example:

python run.py experiment=gmm2d
python run_ais.py experiment=ais_gmm

Callbacks

The RDSMC sampler uses a callback-based architecture for extensibility:

  • WandbCallback - Logging to Weights & Biases
  • PlottingCallback - Generating plots during sampling
  • MetricsCallback - Computing and logging metrics
  • PrintCallback - Printing progress to console
  • SaveSamplesCallback - Saving samples to disk

See experiments/callbacks.py for implementation details.

Citation

If you use this codebase in your research, please cite our paper:

@inproceedings{wu2025reverse,
  title     = {Reverse Diffusion Sequential Monte Carlo Samplers},
  author    = {Wu, Luhuan and Han, Yi and Naesseth, Christian A. and Cunningham, John P.},
  booktitle = {Proceedings of the 39th Conference on Neural Information Processing Systems (NeurIPS 2025)},
  year      = {2025}
}