DiT-Reward: Generative Representations for Text-to-Image Reward Modeling
2026-06-22 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors explore whether models trained to create images can also help judge how good those images are. They develop DiT-Reward, which turns a text-to-image model into a system that predicts how much people like generated images by analyzing image features across different layers. DiT-Reward performs better than an existing model on several benchmarks and works well even if most of the model is kept frozen. The authors also find that certain parts of the model are more important for judging image quality, and using DiT-Reward improves image realism when guiding newer generation models.
text-to-image generationreward modelingdiffusion transformerslatent representationspretrained modelspreference predictionpolicy optimizationimage realisminference speedStable Diffusion
Authors
Yuanming Yang, Guoqing Ma, Bo Wang, Yuan Zhang, Wei Tang, Chenyi Li, Haoyang Huang, Nan Duan
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.