DRM: Diffusion-based Reward Model With Step-wise Guidance

Abstract: Current mainstream methods of aligning diffusion models with human preferences typically employ VLM-based reward models. However, these reward models, pre-trained for semantic alignment, struggle to capture the essential perceptual qualities-such as aesthetics, composition, and visual harmony. In this work, we argue that a model capable of high-fidelity generation must possess a profound understanding of these visual attributes. Based on this insight, we introduce the Diffusion-based Reward Model (DRM), a novel paradigm that use the pre-trained diffusion model as a powerful evaluative backbone. A key advantage of the DRM is its unique ability to assess not only the final image but also the noisy intermediate latents at any stage of the generative process. We leverage this step-wise evaluative capacity in two ways. First, we propose Step-wise GRPO, a reinforcement learning algorithm that provides dense, per-step rewards to resolve the imprecise credit assignment problem in GRPO algorithm, leading to more stable and effective alignment. Second, we introduce Step-wise Sampling, a novel inference strategy that employs the DRM as a dynamic guide to evaluate multiple generation paths at each step, steering the process towards higher-quality outcomes. Extensive experiments confirm that our approach significantly enhances the final quality of generated images. Code: https://github.com/jjaxonx/DRM.

• The paper introduces a diffusion-based reward model that leverages pre-trained diffusion features for fine-grained reward assignment in text-to-image alignment.
• Using step-wise reinforcement learning (GRPO) and guided sampling, the method achieves superior convergence and improved compositional fidelity.
• Empirical results show enhanced alignment metrics and compositional generalization, marking a significant shift from traditional CLIP-style reward models.

Diffusion-based Reward Modeling and Step-wise Guidance for Text-to-Image Alignment

Motivation and Background

Conventional reward models (RMs) for aligning diffusion models with human preferences predominantly employ vision-LLM (VLM) backbones, typically using CLIP-style image encoders. These models are optimized for semantic alignment, but this paradigm fundamentally neglects key perceptual qualities—such as aesthetics, composition, and visual harmony—that characterize human judgments of image quality. The resulting reward signals are further limited to terminal outputs, obscuring the contributions of intermediate latent states in multi-step generative processes. This leads to coarse credit assignment in RL-based alignment protocols, hindering both convergence and the fine-tuning of generative fidelity.

The paper "DRM: Diffusion-based Reward Model With Step-wise Guidance" (2605.25661) proposes a paradigm shift: leveraging the inherent perceptual sensitivity of pre-trained diffusion models as the RM backbone. This approach is grounded in the hypothesis that models capable of high-fidelity synthesis must internalize a nuanced understanding of visual quality, including compositional attributes not captured by semantic encoders.

Methodological Advances

Diffusion-based Reward Model (DRM)

The DRM architecture repurposes features from a pre-trained Diffusion Transformer, truncated to match parameter counts with VLM-based RMs for fair comparison. At each timestep $t$, a noisy latent $x_t$ is extracted and processed through the modified backbone, generating a spatial feature map via MLP and convolutional layers, culminating in a preference score. Training is supervised using human preference triplets $(I_{\text{win}}, I_{\text{lose}}, p)$, where noisy latents at randomly sampled timesteps are compared using a Bradley-Terry loss.

The critical innovation is DRM’s capacity to assess latent states throughout the denoising trajectory, not solely at $t=0$. This enables fine-grained reward assignment and circumvents the information bottleneck inherent to VLMs.

Step-wise GRPO: Fine-Grained Reinforcement Learning Alignment

Existing RL-based alignment methods (e.g., GRPO [23]) uniformly distribute terminal rewards across all steps, resulting in imprecise credit assignment. Step-wise GRPO, enabled by DRM, evaluates $k$ candidate latents at each timestep via SDE branching and computes immediate advantages based on local rewards:

$$Â_i = \frac{R(x_i, c) - \text{mean}(\{R(x_j, c)\}_{j=1}^k)}{\text{std}(\{R(x_j, c)\}_{j=1}^k)}$$

This step-wise signal facilitates stable and efficient policy optimization, as evidenced by empirical convergence rates and reward curves.

Step-wise Sampling: Guided Inference

Beyond training, DRM is deployed at inference to dynamically guide sampling. At each timestep, $k$ candidate paths are explored, scored, and the trajectory is greedily advanced along the highest-reward branch:

$$x_{t-1} = \arg\max_{x'_j \in \{x'_1,\ldots, x'_k\}} R(x'_j, c)$$

This "explore-and-select" mechanism mitigates the rigidity of deterministic samplers, corrects early trajectory errors, and enhances final output quality without retraining.

Experimental Results

Numerical Results

DRM achieves competitive accuracies across standard reward benchmarks: PickScore (64.1%), HPDv2 (73.4%), HPDv3 (82.2%), and ImageReward (74.0%). Notably, DRM outperforms similarly-sized VLM-based RMs (e.g., HPSv3-2B), demonstrating superior parameter efficiency. However, due to its domain-general design (training over noisy latents), a slight performance trade-off exists on benchmarks limited to clean final images.

Ablation studies confirm the necessity of pre-trained diffusion weights for both convergence and maximum accuracy. Increasing image resolution during training further enhances DRM's evaluative capacities.

Step-wise GRPO improves both final reward and convergence speed by ~2.5x compared to standard GRPO, while Step-wise Sampling systematically increases alignment metrics (e.g., ImageReward, PickScore, HPSv3) with higher $k$ values, balancing quality and computational cost.

Qualitative Analysis

Images generated using DRM + Step-wise GRPO exhibit marked improvements in fine-grained detail, compositional accuracy, and absence of artifacts, outperforming baselines in prompt fidelity, attribute binding, and object counting tasks. Visualizations corroborate the model’s ability to handle complex prompts—particularly those demanding accurate counting and compositional reasoning.

Compositional Generalization

Despite DRM being trained solely on subjective preferences, it achieves substantial gains on GenEval benchmarks measuring compositional understanding—a notable 60% relative improvement in counting and 35% in attribute binding. This suggests strong generalization to objective prompt alignment.

Theoretical and Practical Implications

The central claim—that generation requires an implicit understanding of perceptual qualities—grounds the rationale for diffusion-based reward modeling. DRM’s step-wise evaluative capacity directly enables granular credit assignment, resolving longstanding obstacles in RL-based alignment for generative models.

Practically, DRM allows for both more efficient training (via Step-wise GRPO) and improved inference (via Step-wise Sampling), leading to higher fidelity, reduced artifacts, and greater compositional robustness in generated images. The plug-and-play nature of Step-wise Sampling yields immediate quality improvements without retraining, broadening deployment applicability.

Theoretically, the findings indicate that reward model architecture deeply impacts the sample efficiency and attainable alignment, positioning diffusion-based RMs as superior to scaled VLMs for multimodal preference learning. The generalization from subjective preference alignment to objective compositional tasks also suggests further avenues for integrating reward modeling into denoising trajectories.

Future Perspectives

Future developments may explore:

• Scaling DRM architectures on larger diffusion backbones to further enhance evaluative granularity.
• Leveraging step-wise reward signals in video and multimodal generative contexts, extending beyond still images.
• Investigating active learning protocols that dynamically select informative noisy latents for reward modeling.
• Formalizing DRM-guided sampling within theoretical frameworks for trajectory optimization, potentially linking to branching RL or stochastic control.

Conclusion

DRM establishes a principled paradigm for reward modeling in generative alignment, leveraging diffusion models to capture perceptual qualities overlooked by semantic encoders. Step-wise GRPO and Step-wise Sampling harness DRM’s unique capabilities, yielding superior convergence, sample efficiency, and output quality compared to prevailing baselines. The approach’s generalization to compositional tasks, combined with strong empirical results and plug-and-play inference strategies, positions diffusion-based reward modeling as a robust foundation for future preference alignment research in generative AI (2605.25661).