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