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

Yuanming Yang, Guoqing Ma, Bo Wang, Yuan Zhang, Wei Tang +3 more
6/22/2026
cs.LGcs.AI

No confident code implementation found for this paper.

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.

View on arXivView PDF

Code Implementations

No confident code match yet

We couldn't find an author-owned or strongly-evidenced community implementation for this paper. Any repos shown below are weak matches — verify before relying on them.

No code implementations found yet.

Know of an implementation? Let us know in the comments below!

Similar Papers

Cite this paper

@article{yang2026ditreward,
  title  = {DiT-Reward: Generative Representations for Text-to-Image Reward Modeling},
  author = {Yuanming Yang and Guoqing Ma and Bo Wang and Yuan Zhang and Wei Tang and Chenyi Li and Haoyang Huang and Nan Duan},
  year   = {2026},
  eprint = {2606.23626},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2606.23626v1}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2606.23626.svg)](https://codeofpaper.com/paper/2606.23626)

Discussion