The Reward Was in Your Data All Along: Correcting Flow Matching with Discriminator-Guided RL

Nicolas Beltran-Velez; Felix Friedrich; Zhang Xiaofeng; Reyhane Askari-Hemmat; Xiaochuang Han; Adriana Romero-Soriano; Michal Drozdzal

title

authors

Nicolas Beltran-Velez, Felix Friedrich, Zhang Xiaofeng, Reyhane Askari-Hemmat, Xiaochuang Han, Adriana Romero-Soriano, Michal Drozdzal

date

2026-06-17

venue

arXiv (preprint, not peer reviewed)

arXiv:2606.19162

affiliation

AI at Meta

upvotes

21 (Hugging Face Daily Papers, 2026-11-01)

links

abs pdf hf

abstract

Score- and flow-matching models often rely on preference-based reinforcement learning for two purposes: aligning with subjective preferences and, surprisingly, recovering properties such as visual realism and coherent object structure that matching-based training is intended to learn from the data itself. We argue that this reflects a structural mismatch. Matching losses measure ell_2 regression error on the velocity or score field under training-time marginals, a proxy poorly aligned with the visual and semantic properties that determine sample quality at inference. Given a reward aligned with these properties, RL sidesteps the mismatch by evaluating the model on its own samples and following the reward landscape directly. The challenge is to obtain such a reward without relying on human preferences, which are expensive and conflate data realism with annotator inclinations. We propose Discriminator-Guided RL (DRL). DRL trains a discriminator to separate data from base-model samples in a pretrained representation space and uses its logit as the reward in KL-regularized RL. The pretrained space restricts the discriminator to perceptually meaningful directions, and the logit estimates the log-likelihood ratio between data and model, which is the optimal reward for targeting the data distribution. Across SiT, JiT, REPA, and RAE, DRL reduces guidance-free FID (e.g., 9.38 to 2.62 on SiT) and semantic-space FD (e.g., 88.2 to 19.3 on DINOv3 for SiT), with consistent gains across all backbones, and improves human-preference rewards without training on them. It also yields a better Pareto frontier between preference reward and image fidelity under subsequent preference-based post-training, increasing alignment while reducing low-level artifacts such as oversaturation and excessive brightness.

@misc{beltranvelez2026reward,
  title={The Reward Was in Your Data All Along: Correcting Flow Matching with Discriminator-Guided RL},
  author={Nicolas Beltran-Velez and Felix Friedrich and Zhang Xiaofeng and Reyhane Askari-Hemmat and Xiaochuang Han and Adriana Romero-Soriano and Michal Drozdzal},
  year={2026},
  eprint={2606.19162},
  archivePrefix={arXiv},
  url={https://arxiv.org/abs/2606.19162}
}