Review, Remask, Refine (R3): Process-Guided Block Diffusion for Text Generation

TL;DR

The paper introduces R3, a process-guided framework for iterative text generation that improves output quality by reviewing, remasking, and refining generated text blocks without additional training.

Contribution

R3 is a novel, training-free framework that leverages a Process Reward Model to guide targeted correction in pre-trained masked diffusion models.

Findings

Enhanced text generation quality demonstrated

Applicable to various pre-trained models

No additional training required

Abstract

A key challenge for iterative text generation is enabling models to efficiently identify and correct their own errors. We propose Review, Remask, Refine (R3), a relatively simple yet elegant framework that requires no additional model training and can be applied to any pre-trained masked text diffusion model (e.g., LLaDA or BD3-LM). In R3, a Process Reward Model (PRM) is utilized for the Review of intermediate generated blocks. The framework then translates these PRM scores into a Remask strategy: the lower a block's PRM score, indicating potential mistakes, the greater the proportion of tokens within that block are remasked. Finally, the model is compelled to Refine these targeted segments, focusing its efforts more intensively on specific sub-optimal parts of past generations, leading to improved final output.