DiT-Reward: 텍스트-이미지 보상 모델링을 위한 생성적 표현 (Generative Representations)

이미지 생성을 위해 학습된 표현(representations)이 생성된 이미지의 평가도 지원할 수 있을까요? 우리는 생성적 표현 학습(generative representation learning)의 다운스트림 태스크(downstream task)로서 텍스트-이미지 보상 예측(text-to-image reward prediction)을 연구합니다. 이를 위해, 우리는 사전 학습된 텍스트-이미지 Diffusion Transformer를 거의 깨끗한(near-clean) 이미지 잠재 변수(latents)를 처리하고 트랜스포머 레이어 전반에 걸쳐 텍스트 조건부 이미지 표현을 집계함으로써 보상 모델(reward model)로 변환하는 DiT-Reward를 소개합니다. HPSv3와 동일한 학습 데이터 혼합(training data mixture) 조건에서, DiT-Reward는 평가된 4가지 선호도 벤치마크 모두에서 HPSv3를 능가하며, HPDv2에서 85.6%, HPDv3에서 77.6%에 도달했습니다. 생성 백본(generative backbone)이 동결(frozen)된 상태에서도, 경량화된 학습된 헤드(learned head)를 통해 그 표현으로부터 의미 있는 선호도 예측을 여전히 추출할 수 있습니다. 깊이(depth)에 따른 프로빙(Probing) 결과, 다운스트림 보상 성능은 중간에서 후기 레이어(middle-to-late layers)에서 가장 강력하며, 서로 다른 단계의 표현을 결합할 때 이점을 얻는다는 것이 추가로 밝혀졌습니다. 또한 우리는 생성 백본의 용량(capacity)에 따라 일관된 양의 스케일링(positive scaling)을 관찰했습니다. 마지막으로, Flow-GRPO를 사용하여 Stable Diffusion 3.5 Large를 최적화할 때, DiT-Reward는 일치하는 학습 궤적(training trajectory)을 따라 HPSv3를 능가하며, 특히 사실주의(realism) 측면에서 뚜렷한 이득을 보였습니다. 직접적인 잠재 변수 점수 매기기(Direct latent scoring) 또한 유사한 피크 메모리(peak memory)를 유지하면서 HPSv3 대비 1.65배의 추론 속도 향상을 달성했습니다. 이러한 결과는 사전 학습된 생성적 DiT가 보상 모델링 및 정책 최적화(policy optimization)를 위한 전이 가능한 표현(transferable representations)을 제공함을 보여줍니다.