Discrete Diffusion SMC

This repo contains PyTorch implementation of the paper "Inference-Time Scaling of Discrete Diffusion Models via Importance Weighting and Optimal Proposal Design"

by Zijing Ou, Chinmay Pani and Yingzhen Li.

We propose a Sequential Monte Carlo (SMC) framework that enables scalable inference-time control of discrete diffusion models through principled importance weighting and optimal proposal construction. Our approach derives tractable importance weights for a range of intermediate targets and characterises the optimal proposal, for which we develop two practical approximations: a first-order gradient-based approximation and an amortised proposal trained to minimise the log-variance of the importance weights.

In this repository, we provide the training and evaluation code for the experiements of language modelling, biology design, and image generation. For the toy synthetic experiments, please refer to the codebase.

Citation

:smile:If you find this repo is useful, please consider to cite our paper:

@article{ou2025inference,
  title={Inference-Time Scaling of Discrete Diffusion Models via Importance Weighting and Optimal Proposal Design},
  author={Ou, Zijing and Pani, Chinmay and Li, Yingzhen},
  journal={arXiv preprint arXiv:2505.22524},
  year={2025}
}