Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Zijing Ou, Chinmay Pani, Yingzhen Li

TL;DR

The paper tackles the challenge of aligning outputs of discrete diffusion models with downstream constraints at inference time. It introduces a Sequential Monte Carlo (SMC) framework that yields tractable importance weights for intermediate targets and explores two practical optimal-proposal approximations: a first-order gradient-based method and an amortised proposal that minimises the log-variance of the weights. The authors derive tractable weights for product targets and reward-tilting, propose a twisted intermediate target with annealing, and empirically validate the approach across synthetic data, language modelling, biology design, and image generation, showing improved controllability and sample quality. This work provides a versatile, scalable recipe for test-time scaling and post-training alignment of discrete diffusion models, with broad applicability across domains and tasks.

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.

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

TL;DR

The paper tackles the challenge of aligning outputs of discrete diffusion models with downstream constraints at inference time. It introduces a Sequential Monte Carlo (SMC) framework that yields tractable importance weights for intermediate targets and explores two practical optimal-proposal approximations: a first-order gradient-based method and an amortised proposal that minimises the log-variance of the weights. The authors derive tractable weights for product targets and reward-tilting, propose a twisted intermediate target with annealing, and empirically validate the approach across synthetic data, language modelling, biology design, and image generation, showing improved controllability and sample quality. This work provides a versatile, scalable recipe for test-time scaling and post-training alignment of discrete diffusion models, with broad applicability across domains and tasks.

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.

Paper Structure

This paper contains 37 sections, 13 theorems, 62 equations, 23 figures, 10 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Discrete Diffusion Models
  4. Importance Sampling and Sequential Monte Carlo
  5. Sequence Monte Carlo for Discrete Diffusion Models
  6. Importance Weight: Tractability with Pretrained Diffusion Models
  7. Choices of Proposal: the Way to Reduce Variance
  8. Approximated Optimal Proposal via First-Order Approximation
  9. Amortised Optimal Proposal via Log-Variance Minimisation
  10. Sequential Monte Carlo Recipe: Practical Implementation
  11. Experiments
  12. Synthetic Experiemnts
  13. Language Modelling
  14. Biology Design
  15. Image Generation
  16. ...and 22 more sections

Key Result

Proposition 1

Let $R_t$ be the rate matrix generating the forward transition kernel $\gamma(x_t | x_{t-\Delta t})$, and $\hat{R}_t$ be its counterpart associated with the backward proposal kernel $q(x_{t-\Delta t} | x_t)$, where $\Delta t \to 0$ is the infinitesimal time increment. Then, the importance weight at

Figures (23)

  • Figure 1: SMC results for the reward-tilting and product target distributions.
  • Figure 2: Comparisons on reward-tilted discreteised MoGs. We consider the reward function as $r(X, Y) \!=\! -\hat{X}^2/100-\hat{Y}^2$, where $\hat{X} \!=\! 12(X/63 - 1/2)$ and $\hat{Y} \!=\! 12(Y/63-1/2)$.
  • Figure 3: Comparisons on reward-tiled binary MNIST. We train a classier $p_{\text{clf}}(y|x)$ on the clean data, and the reward is given by $r(x) = \log p_{\text{clf}}(y_{\text{target}}|x)$, where $y_{\text{target}}$ denotes the target digit.
  • Figure 4: Results of DNA sequence design. Both the pretrained discrete diffusion model and the reward models are adopted from wang2024fine.
  • Figure 5: The results of text-to-image generation across different reward models.
  • ...and 18 more figures

Theorems & Definitions (22)

  • Proposition 1: SMC for Continuous-Time Discrete Diffusion
  • Proposition 2: Locally Optimal Proposal
  • Corollary 1
  • Proposition 3
  • Proposition 3: Locally Optimal Proposal
  • proof
  • Corollary 1
  • proof
  • Proposition 3
  • proof
  • ...and 12 more