Inference-Time Scaling of Discrete Diffusion Models via Importance Weighting and Optimal Proposal Design
This work addresses the need for constrained generative processes in applications such as biology design and text-to-image generation, representing an incremental improvement with a novel method for a known bottleneck.
The paper tackles the problem of enabling scalable inference-time control for discrete diffusion models to adhere to constraints in real-world applications, proposing a Sequential Monte Carlo framework that enhances controllability and sample quality across tasks like language modelling and text-to-image generation.
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.