English

Inference-Time Scaling of Discrete Diffusion Models via Importance Weighting and Optimal Proposal Design

Machine Learning 2026-03-17 v4

Abstract

Discrete diffusion models have become highly effective across various domains. However, real-world applications often require the generative process to adhere to certain constraints. To this end, 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. Specifically, 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. Empirical results across synthetic tasks, language modelling, biology design, and text-to-image generation demonstrate that our framework enhances controllability and sample quality, highlighting the effectiveness of SMC as a versatile recipe for scaling discrete diffusion models at inference time.

Keywords

Cite

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