TLDR: IterRef is a novel test-time scaling method for discrete diffusion models that uses iterative reward-guided refinement. It leverages a noising-denoising process within a Multiple-Try Metropolis framework to progressively align intermediate states with a reward function, ensuring convergence to an optimal distribution. Experiments across text and image generation tasks show IterRef consistently outperforms prior methods, achieving significant gains in generation quality, especially with low computational budgets, and is particularly effective in later denoising stages and for safety alignment.
In the rapidly evolving world of artificial intelligence, large language models and diffusion models have made incredible strides in generating text and images. However, as these models grow in complexity, researchers are constantly looking for ways to improve their performance, especially during the ‘test-time’ or inference phase, where the model generates its output. While continuous diffusion models have seen significant advancements in this area, discrete diffusion models, which work with distinct tokens like words or image pixels, have faced unique challenges.
A new research paper introduces a novel method called Iterative Reward-Guided Refinement (IterRef) to address these challenges for discrete diffusion models. This approach focuses on enhancing the quality of generated content by iteratively refining the model’s intermediate steps, guided by a reward function. Unlike previous methods that try to guide the next step based on the current state, IterRef actively refines the current state itself, steering it closer to the desired outcome.
The core idea behind IterRef is inspired by a ‘predictor-corrector’ mechanism. It involves a noising-denoising process: adding a small amount of noise to an intermediate state allows the model to explore different possibilities, and then denoising helps it restore consistency with the target, but now in a more reward-aligned way. This entire process is formalized within a Multiple-Try Metropolis (MTM) framework, which provides a theoretical guarantee that the iterative refinement will converge to the desired, reward-aligned distribution.
The researchers conducted extensive experiments across various discrete diffusion models and tasks. For language generation, they used models like MDLM and LLaDA-8B, guiding them with reward functions for tasks such as controlling toxicity, sentiment, perplexity (a measure of language fluency), and grammatical correctness (CoLA). For image generation, they applied IterRef to MaskGIT, using CLIPScore as a reward to align images with text prompts.
The results were consistently positive. IterRef demonstrated superior performance compared to existing reward-guided diffusion methods, often achieving better results with significantly less computational effort. For instance, on toxicity control with MDLM, IterRef was up to 8 times faster than prior state-of-the-art methods. A notable finding was that increasing the number of refinement iterations proved more effective than simply generating more candidate samples, highlighting the power of the iterative approach.
Interestingly, the study also revealed that for discrete diffusion models, applying IterRef at later stages of the denoising process (closer to the final output) was more effective for most tasks. This contrasts with continuous diffusion models, where early stages often determine much of the content. This insight suggests that the entire denoising trajectory plays a crucial role in shaping the final generation in discrete models.
A compelling case study showcased IterRef’s potential in safety alignment. When applied to LLaDA-8B for detoxification, IterRef significantly reduced the proportion of toxic generations, even with limited computational resources. This demonstrates its practical utility in making large language models safer and more reliable.
Also Read:
- Lookahead Unmasking: A New Strategy for Accurate Text Generation in Diffusion Language Models
- Improving Image Generation with Smarter Diffusion Model Guidance
In conclusion, IterRef offers a robust and theoretically sound method for test-time scaling in discrete diffusion models. By iteratively refining intermediate states through a reward-guided noising-denoising process, it consistently improves generation quality across different modalities and tasks, especially under low compute budgets. This work paves the way for more effective and controllable discrete generative AI. You can read the full research paper here.
