Adnan KhanAdnanProduct & AI
San Francisco Bay Area
Get in touch

Source note3 min read

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

Source

arXiv

arxiv.org/abs/2606.23626v1

Read the source

Source

Plain-Language Brief

The paper introduces DiT-Reward, a model that uses pre-trained text-to-image Diffusion Transformer to evaluate generated images by converting it into a reward model. DiT-Reward outperforms HPSv3 on four preference benchmarks and can extract meaningful preference predictions even with a lightweight learned head when the generative backbone is frozen. The performance is strongest in the middle-to-late layers and benefits from combining representations across different stages. DiT-Reward also achieves faster inference speed and better results when used to optimize Stable Diffusion 3.5 Large with Flow-GRPO.

What the Document Says

  • DiT-Reward is a text-to-image reward prediction model derived from a pre-trained Diffusion Transformer.
  • DiT-Reward processes near-clean image latents and aggregates text-conditioned image representations across transformer layers.
  • Under the same training data mixture as HPSv3, DiT-Reward outperforms HPSv3 on four preference benchmarks.
  • The highest performance achieved by DiT-Reward is 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.
  • Probing across depth reveals that downstream reward performance is strongest in the middle-to-late layers.
  • Combining representations across different stages benefits downstream reward performance.
  • DiT-Reward shows consistent positive scaling with generative backbone capacity.
  • When used to optimize Stable Diffusion 3.5 Large with Flow-GRPO, DiT-Reward outperforms HPSv3 along the matched training trajectory.

Key Concepts Mentioned

  • text-to-image reward prediction
  • DiT-Reward
  • Diffusion Transformer
  • HPSv3
  • preference benchmarks
  • generative backbone
  • learned head
  • Stable Diffusion 3.5 Large
  • Flow-GRPO

Claims Made by the Document

  • DiT-Reward outperforms HPSv3 on preference benchmarks
  • Downstream reward performance is strongest in the middle-to-late layers
  • Combining representations across different stages benefits downstream reward performance
  • DiT-Reward shows consistent positive scaling with generative backbone capacity
  • When used to optimize Stable Diffusion 3.5 Large with Flow-GRPO, DiT-Reward outperforms HPSv3

Evidence Used

  • 85.6% on HPDv2 and 77.6% on HPDv3 achieved by DiT-Reward
  • Results showing performance across different layers
  • Results comparing DiT-Reward and HPSv3 when used to optimize Stable Diffusion 3.5 Large with Flow-GRPO

Methods / Approach

  • Converting a pre-trained text-to-image Diffusion Transformer into a reward model by processing near-clean image latents and aggregating text-conditioned image representations across transformer layers
  • Freezing the generative backbone and using a lightweight learned head to extract preference predictions

Limitations Stated by the Authors

UNKNOWN

Limitations Noted by the Agent

Agent interpretation — not from the source.

UNKNOWN

Open Questions

UNKNOWN

Notes for Future Q&A

May include interpretation. Labeled where it does.

UNKNOWN

Raw Abstract / Excerpt

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.