Effective Test-Time Scaling of Discrete Diffusion through Iterative Refinement

TL;DR

This paper introduces IterRef, a novel test-time scaling method for discrete diffusion models that iteratively refines states guided by rewards, significantly improving generation quality especially under low compute budgets.

Contribution

We propose IterRef, a new reward-guided refinement technique for discrete diffusion models that explicitly refines states during inference, formalized within a MTM framework.

Findings

IterRef improves reward-guided generation quality across text and image domains.

It achieves significant gains under low compute budgets.

Outperforms prior state-of-the-art baselines.

Abstract

Test-time scaling through reward-guided generation remains largely unexplored for discrete diffusion models despite its potential as a promising alternative. In this work, we introduce Iterative Reward-Guided Refinement (IterRef), a novel test-time scaling method tailored to discrete diffusion that leverages reward-guided noising-denoising transitions to progressively refine misaligned intermediate states. We formalize this process within a Multiple-Try Metropolis (MTM) framework, proving convergence to the reward-aligned distribution. Unlike prior methods that assume the current state is already aligned with the reward distribution and only guide the subsequent transition, our approach explicitly refines each state in situ, progressively steering it toward the optimal intermediate distribution. Across both text and image domains, we evaluate IterRef on diverse discrete diffusion models and observe consistent improvements in reward-guided generation quality. In particular, IterRef achieves striking gains under low compute budgets, far surpassing prior state-of-the-art baselines.