DiT-Reward: Generative Representations for Text-to-Image Reward Modeling

2026-06-22

AI Summary

This work proposes an efficient reward modeling approach based on pretrained diffusion Transformers (DiT) for evaluating the quality and text-image alignment of generated images. By freezing the DiT backbone and introducing only a lightweight trainable head, the method aggregates text-conditioned image representations from intermediate-to-late DiT layers and performs scoring directly in the latent space. This is the first demonstration that intermediate features from generative DiTs can effectively support reward prediction without fine-tuning. Combined with cross-layer feature fusion and Flow-GRPO policy optimization, the approach achieves 85.6% and 77.6% accuracy on HPDv2 and HPDv3 benchmarks, respectively—surpassing HPSv3—while offering 1.65× faster inference at comparable memory cost, significantly enhancing the perceptual realism of generated images.

Abstract

Can representations learned for image generation also support the evaluation of generated images? We study text-to-image reward prediction as a downstream task of generative representation learning. To this end, we introduce DiT-Reward, which converts a pretrained text-to-image Diffusion Transformer into a reward model by processing near-clean image latents and aggregating text-conditioned image representations across transformer layers. Under the same training data mixture as HPSv3, DiT-Reward outperforms HPSv3 on all four evaluated preference benchmarks, reaching 85.6% on HPDv2 and 77.6% on HPDv3. When the generative backbone is frozen, a lightweight learned head can still extract meaningful preference predictions from its representations. Probing across depth further reveals that downstream reward performance is strongest in the middle-to-late layers and benefits from combining representations across different stages. We also observe consistent positive scaling with generative backbone capacity. Finally, when used to optimize Stable Diffusion 3.5 Large with Flow-GRPO, DiT-Reward outperforms HPSv3 along the matched training trajectory, with particularly clear gains in realism. Direct latent scoring also achieves a 1.65x inference speedup over HPSv3 with comparable peak memory. These results show that pretrained generative DiTs provide transferable representations for reward modeling and policy optimization.

Problem

• text-to-image reward modeling
• generative representations
• pretrained diffusion models
• image evaluation
• preference prediction

Innovation

• Diffusion Transformer
• reward modeling
• generative representation
• pretrained transfer
• latent scoring