PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2606.19162v1cs.LGcs.CV

TL;DR

Flow-matching losses may imperfectly capture data properties such as realism and structure, motivating preference-based RL for recovery. DRL uses a discriminator’s representation-space density-ratio logit as a KL-regularized RL reward, consistently improving distributional alignment, preference metrics, and preference-RL performance–drift tradeoffs.

  • Problem

    Matching-based objectives may imperfectly capture data properties such as visual realism and coherent structure, despite models being trained on data containing those properties.

  • Method

    DRL trains a discriminator in a pretrained self-supervised representation space and uses its density-ratio logit as the reward for KL-regularized RL.

  • Results

    DRL consistently improves distributional alignment, preference metrics, and performance–drift tradeoffs under subsequent preference-based RL.

  • Takeaways & Limitations

    RL post-training can complement flow-based matching objectives by recovering training-data structure that matching-based objectives imperfectly capture.

  • Takeaways & Limitations

    Because DRL restricts corrections to the pretrained representation space, it generally cannot target the data distribution exactly.

Abstract

from arXiv · show

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.

1 Introduction

The paper argues that flow/score-matching objectives can imperfectly capture data properties that RL recovers, because training-time matching marginals differ from model rollout distributions. It introduces Discriminator-Guided RL (DRL), which uses a pretrained-representation discriminator’s density-ratio logit as a KL-regularized reward and improves multiple image-generation architectures without preference-reward training.

  • Contribution: The introduction positions RL post-training as complementary to matching: it can recover training-data structure that standard objectives imperfectly capture, not merely optimize external preferences.This addresses the need for a reward proxy without relying on preference-based rewards or additional preference data.
  • Why matching fails: FSM losses are measured under interpolation marginals q_t, whereas sampling follows model rollouts p_t, so matching loss may not control gaps in properties evaluated on generated samples.In the worst case, no guarantee exists; under regularity assumptions, the guarantee worsens with the reward’s ℓ2 Lipschitz constant.
  • Motivation: The paper frames RL’s recovery of properties already present in q as evidence that the matching objective, rather than data or model capacity, is a plausible bottleneck.This motivates analyzing when flow-score-matching losses control a property gap in a simplified probability-flow ODE setting.
  • Method: DRL trains a discriminator in a pretrained self-supervised representation space to estimate the data-to-model density ratio, then uses its logit as the reward for KL-regularized RL.The representation space constrains discrimination to directions intended to capture meaningful data properties.
  • Empirical results: DRL improves SiT, JiT, REPA, and RAE over their base models across multiple feature-space Fréchet distances, while also improving held-out human-preference rewards without training on them.The reported gains use a simple linear discriminator and are summarized in Figure 1.

2 Preliminaries

Flow and score matching transport Gaussian noise to the data distribution by learning conditional velocity or score fields through simulation-free regression. KL-regularized RL then shifts a pretrained flow model toward high-reward regions, using path-distribution regularization because endpoint densities are unreliable.

  • Flow and Score Matching: Flow and diffusion models transport Gaussian noise to data by interpolating X_1 ∼ q and X_0 ∼ N(0, I) as X_t = α(t)X_1 + β(t)X_0.The schedule satisfies α(0)=0, β(0)=1, α(1)=1, and β(1)=0, producing marginals between Gaussian noise and q.
  • Flow and Score Matching: Sampling reduces to learning the conditional velocity v_t or score s_t field, each trainable with simulation-free ℓ2 regression through conditional flow matching or denoising score matching.Under the interpolation, velocity and score fields are recoverable from each other, so learning either objective suffices.
  • RL Post-training for Flow Models: KL-regularized RL moves a pretrained model p_base toward high-reward regions while limiting deviation, equivalently targeting p*(x) ∝ exp(λr(x))p_base(x).The parameter λ controls the trade-off between reward and KL regularization.
  • RL Post-training for Flow Models: Because flow-model endpoint densities are unreliable, RL replaces endpoint regularization with analytically tractable KL regularization over trajectory distributions under a shared noisy sampling SDE.This path-space objective uses Girsanov’s theorem; with a memoryless σ(t), its minimizer preserves the endpoint of the original KL-regularized objective.
  • RL Post-training for Flow Models: The path-distribution objective can be optimized with policy-gradient methods including REINFORCE, PPO, and GRPO, or with adjoint-based methods.The listed policy-gradient approaches include Williams (1992), Black et al. (2024), Fan et al. (2023), Schulman et al. (2017), Shao et al. (2024), and Liu et al. (2025a).

3 Motivation: Understanding the Limitations of Flow and Score Matching

Flow and score matching can achieve low training loss while failing to preserve reward-relevant properties because rollout errors compound off the training marginals and velocity/score geometry may misalign with sample rewards. RL instead evaluates generated samples and directly optimizes the reward landscape, providing a signal aligned with the target property.

  • Worst-case limitation: Low FSM loss can be arbitrarily uninformative about the reward gap: arbitrarily small training error may still produce an endpoint reward near its worst-case regional value.Proposition 3.1 constructs velocity and score fields with uniformly small training error but endpoint laws whose expected reward is at most η + δ, regardless of the data reward.
  • Worst-case limitation: Early rollout errors can push trajectories into low-training-density regions, where subsequent errors compound and cause distribution shift between training marginals q_t and rollout marginals p_t.Figure 2 depicts rollout paths drifting off the noising band and missing the high-reward endpoint region.
  • Quantitative guarantee: Under Lipschitz regularity and uniform velocity control, a reward certificate exists, but its quality depends on velocity error ε, field smoothness L_v, and reward sharpness L_r.The bound factorizes into ε, (e^L_v − 1)/L_v, and L_r; an on-policy variant under rollout marginals p_t also holds.
  • Geometry mismatch: Velocity-space accuracy may require impractically small ε because small velocity errors can produce large reward errors when the two geometries are misaligned.The bound implies ε ≈ δ/L_r to achieve reward gap δ; sharper rewards therefore demand smaller velocity errors, while matching-loss gradient signal-to-noise degrades in this regime.
  • RL’s edge: RL avoids these obstructions by evaluating generated samples without train–test mismatch and optimizing directly along reward-improving directions rather than velocity or score geometry.With a suitable reward, RL provides an optimization signal directly aligned with the property of interest, whereas FSM must recover it indirectly from regression targets.

4 Method: Discriminator-Guided RL

Discriminator-Guided RL (DRL) estimates the data-to-base-model log-density ratio from samples and uses it as a KL-regularized RL reward. A frozen pretrained representation makes estimation more reliable while restricting correction to representation-visible structure.

  • Motivation: DRL replaces human-preference rewards with a discriminator-derived reward designed to capture properties of q that flow matching may fail to learn.Preference rewards can optimize orthogonal attributes such as aesthetic appeal rather than data realism.
  • From samples to rewards: At λ=1, the ideal KL-regularized reward is the log-density ratio log(q/pbase), and logistic discrimination estimates this ratio through the discriminator logit.The optimal discriminator is D*(x)=q(x)/(q(x)+pbase(x)), whose logit recovers the ideal reward up to additive constants.
  • Using representations: DRL estimates the reward in a frozen pretrained representation space, reducing statistical difficulty and preventing the discriminator from relying on semantically irrelevant artifacts.The reward takes the form r_hat(x)=h(phi(x)), restricting correction to structure visible through the encoder.
  • Target characterization: Because representation restriction prevents exact targeting of q in general, DRL makes the smallest KL change to pbase that matches target and model distributions after projection through phi.Variation invisible to phi remains unchanged, so choosing the representation determines which aspects of q can be corrected.
  • The DRL pipeline: The two-stage pipeline trains the representation-space discriminator first, then fine-tunes the flow model with KL-regularized RL using adjoint matching to exploit reward gradients.Stage 1 defines the discriminator-logit reward from q and pbase samples; Stage 2 optimizes the base model under that reward.

5 Experiments

Across four flow-model backbones, DRL improves distributional alignment and held-out preference rewards without preference training, while λ=10 provides the default trade-off. Starting from DRL also improves subsequent preference alignment, yielding higher rewards at comparable low-level drift.

  • Experimental setup: Across four ImageNet-pretrained backbones spanning latent- and pixel-space architectures, DRL uses a frozen DINOv2-Large discriminator trained for 10k steps followed by 3k RL steps.The evaluated models are SiT, JiT, REPA, and RAE.
  • Distributional alignment: DRL improves all 16 model–feature FD pairs with tuned λ⋆ without CFG, while λ=1 improves all but one; gains are often greater than half in DINOv2 and DINOv3.With best CFG, base models are already closer to data, but DRL still improves alignment in most cases.
  • Guidance trade-off: FD decreases through λ=5–10 and degrades at higher λ as coverage drops, so λ=10 is adopted as the default downstream setting.At λ=40, the precision–coverage trajectory collapses off the frontier because the model overfits the discriminator reward.
  • Held-out preference rewards: DRL improves every held-out human-preference reward without seeing preference data, across ImageReward, PickScore, Aesthetics v2.5, and HPSv2.Improvements are evaluated with and without classifier-free guidance using normalized reward gains.
  • Preference alignment: DRL+PRL achieves higher HPSv2 than Base+PRL at comparable brightness, saturation, contrast, colorfulness, and whiteness drift across the reward–drift frontier.DRL+PRL samples remain structurally coherent and naturally colored across λPRL, avoiding the brighter and whiter outputs seen with aggressive Base+PRL.
  • Ablations: Pretrained frozen features are essential: from-scratch discriminators underperform despite 95% validation accuracy, while small R1 stabilizes λ=10 training but degrades performance at λ=1.Linear and MLP heads perform similarly, motivating the simpler linear head.

6 Related Work

The paper situates DRL within three related threads: imitation learning, inverse reinforcement learning, and calibration. It extends these ideas to flow- and score-matching generative models and frames calibration as a distribution-level question.

  • RL vs. Imitation Learning: Flow and score matching may exhibit imitation-learning pathologies, with continuous-time training offering weaker guarantees than standard DAgger-style bounds (Ross et al., 2011).
  • Inverse RL: DRL adapts inverse reinforcement learning by recovering a reward from data and base-model samples, then training a generative model against it.This extends the structure used by GAIL (Ho and Ermon, 2016) and AIRL (Fu et al., 2018) from standard RL control to generative models.
  • Calibration: The paper asks whether small flow- or score-matching loss guarantees similar expectations under the data and model distributions, relating DRL to direct distributional constraints and KL-based calibration.Closest is Smith et al. (2025), which finds the KL-closest model satisfying user-specified moment constraints.

7 Conclusions … A Detailed Related Work

The paper concludes that DRL complements matching objectives by using on-policy reward geometry to improve distributional alignment, preference metrics, and performance–drift tradeoffs, while related work positions it as an explicitly on-policy correction distinct from adversarial, matching-based, and diffusion-specific alternatives.

  • 7 Conclusions: DRL consistently improves distributional alignment, preference metrics, and performance–drift tradeoffs by evaluating models on their own samples and exploiting reward geometry rather than ℓ2 distance.The conclusion frames matching objectives and RL as complementary tools for extracting more value from the data.
  • 7 Conclusions: The paper identifies frozen self-supervised representations, worst-case bounds, and unclear links to practical SDE/ODE samplers as directions for further investigation.It also asks when an RL stage is helpful and whether comparable representations can be learned jointly with the model.
  • Appendix Contents: The appendix organizes proofs, implementation details, experimental setups, extended results, ablations, and qualitative samples supporting the paper’s theoretical and empirical analyses.The listed materials include proofs for preliminaries, matching limitations, and DRL; setup sections; extended metric and Pareto analyses; feature-space and discriminator ablations; and qualitative comparisons.
  • A Detailed Related Work: The related-work discussion connects flow- and score-matching pathologies to covariate shift, exposure bias, and DAgger-style imitation-learning failures, motivating explicitly on-policy correction.Prior diffusion approaches often use model-specific off-distribution perturbations, while the closest cited MMD method uses short, warm-started rollouts that address local rather than full trajectory occupancy mismatch.
  • A Detailed Related Work: DRL differs from adversarial imitation because its discriminator supplies a fixed reusable reward, whereas GAIL matches occupancy measures through an adversarial loop and AIRL decomposes logits into reward and shaping terms.The cited IRL lineage includes classical reward inference and maximum-entropy formulations.
  • A Detailed Related Work: DRL corrects an existing flow rather than learning an implicit model or few-step sampler, and it trains on model samples instead of the posterior-mean MSE denoiser used by earlier diffusion adversarial training.The distinction is principled because the posterior mean is not generally a sample from the model.
  • A Detailed Related Work: For reward-based post-training, DRL contributes the reward while effective Stage 2 optimization must remain on-policy and propagate reward information through the landscape via gradients or score-function estimators.Tilt Matching and reward-weighted regression instead convert rewards into matching objectives, leaving them exposed to matching-geometry limitations.
  • A Detailed Related Work: The calibration discussion relates DRL to distribution-level constraints and shows that, although CGM-reward and DRL can coincide in a linear representation case, they generally target different objects and DRL also works with MLP discriminators.CGM-reward matches user-specified moments through an empirical maximum-entropy dual, whereas DRL recovers a discriminator-logit tilt under KL regularization.

B Limitations

DRL remains limited by practical RL fine-tuning costs, CFG incompatibility, sensitivity to KL regularization, and dependence on the chosen representation. Its theoretical justification also uses a simplified ODE setting, while training and inference dynamics may differ in practice.

  • Classifier-free guidance: DRL cannot cleanly incorporate classifier-free guidance during training, which often causes instability and reduces gains over the base model at best CFG scales.These limitations reflect current RL algorithms’ difficulty handling CFG; principled CFG-compatible RL training remains future work.
  • Choice of λ: The KL weight λ strongly affects performance, requiring exploration of multiple values and adding computational cost despite λ=10 being a reliable default in these experiments.The default may not transfer across settings, motivating methods more robust to λ.
  • Reverse KL behavior: Reverse-KL optimization can make DRL mode-seeking at large λ, although this was not a practical issue at the recommended values.R1 gradient regularization was beneficial at higher λ, partly explaining its use there.
  • Additional compute: DRL requires additional post-training compute beyond base-model training, and this expense may not be justified in every scenario.The added cost is only a small fraction of pretraining but remains an extra burden shared by RL methods.
  • Role of the representation: Because DRL relies on a pretrained representation, it can correct only distributional differences visible to the chosen embedder and therefore requires a suitable representation.This restriction focuses rewards on semantic discrepancies but also limits the method’s scope.
  • Scope of the theory: The theory uses a simplified probability-flow ODE setting, while training uses a memoryless SDE and evaluation an ODE sampler whose rollout distributions may differ during training.The analysis captures proposed geometric obstructions but does not fully model every practical setup; reward-gradient effects are hypothesized to affect both samplers similarly.

C Proofs for Section 2: Preliminaries … D.1 Distribution Shift and Error Accumulation

The appendix proves foundational equivalences for score, velocity, SDE, and probability-flow ODE formulations, then introduces the proof sequence for Section 3’s limitations analysis. The supplied passages do not provide substantive proof details for D.1 Distribution Shift and Error Accumulation.

  • C Proofs for Section 2: Preliminaries: The appendix collects proofs supporting the statements in Section 2.
  • C.1 Score-velocity relationship: Score and velocity fields are related through conditional expectations under the interpolation and Tweedie’s identity.
  • C.1 Score-velocity relationship: The proof derives velocity by expressing the noise through the interpolation, computing conditional means, and substituting them into the velocity definition.
  • C.2 SDE and ODE marginal equivalence: The SDE and probability-flow ODE generate identical marginals q_t at every time for any noise schedule σ(t) ≥ 0.
  • C.2 SDE and ODE marginal equivalence: This equivalence follows because both dynamics satisfy the same continuity equation with the same initial condition after diffusion and score terms cancel.
  • D Proofs for Section 3: Understanding the Limitations of Flow and Score Matching: The Section 3 proofs proceed from the no-certificate result for the standard probability-flow ODE to reward geometry and the RL reward-regret bound.
  • D.1 Distribution Shift and Error Accumulation: The supplied passages do not state a substantive result for D.1 Distribution Shift and Error Accumulation.

D.1.1 Proof of proposition 3.1: no reward certificate for the standard probability-flow sampler … D.3.1 RL reward regret bound

The appendices show that small matching error alone cannot certify reward quality under standard probability-flow sampling, while stronger rollout control yields Lipschitz-reward guarantees and RL admits a KL-regularized regret bound. The negative constructions and tightness examples identify both the limits of matching-based guarantees and the conditions under which reward control is possible.

  • D.1.1 Proof of proposition 3.1: no reward certificate for the standard probability-flow sampler: Arbitrarily small flow- or score-matching error can still produce an arbitrarily large reward gap under the standard probability-flow sampler.The Gaussian construction has training error exactly ε, while rollout variance inflation can grow without bound; the radial-shell construction applies to both velocity and score parametrizations.
  • D.1.1 Proof of proposition 3.1: no reward certificate for the standard probability-flow sampler: The general counterexample hides a harmful perturbation in a thin, inward-moving shell that captures trajectories and transports them into a low-reward region while keeping training error small.The perturbation is constructed in inverse-flow coordinates, where true trajectories are stationary, then mapped back to x-space.
  • D.2.1 Reward geometry under stronger rollout control: Under stronger rollout control, coupling yields endpoint reward certificates for every Lipschitz reward, with global and on-policy bounds sharing the same Grönwall amplification factor.Global control bounds field mismatch pathwise, whereas on-policy control bounds it after averaging over rollout marginals.
  • D.2.1 Reward geometry under stronger rollout control: The exponential amplification factor and its dependence on dynamics and reward Lipschitz constants are tight, including for bounded rewards.The tightness construction attains equality, and the bounded-reward extension remains essentially tight.
  • D.3.1 RL reward regret bound: KL-regularized RL targets the exponentially reward-tilted endpoint distribution induced by the base trajectory distribution.The optimized trajectory law P* has final marginal p*; the objective is L_RL(P) = 1/λ KL(P ∥ P_base) − E_p[r].
  • D.3.1 RL reward regret bound: The RL reward regret bound follows by comparing trajectory laws to the tilted optimum and applying endpoint data processing to KL divergence.Because the tilt depends only on the terminal state, the optimal path-measure construction induces exactly the desired final-sample marginal.

E Proofs for Section 4: Method: Discriminator-Guided RL … G Local Linear Integrator for the Memoryless SDE

The proofs formalize how feature-based KL-regularized RL corrects feature marginals while preserving within-feature conditionals, and bound the resulting test-function error. The appendix then explains memoryless-schedule adjoint matching and motivates a local linear SDE integrator for stable flow-model training.

  • E.1.1 Feature-space correction: The feature-space correction is obtained by tilting the base density with a function of the encoder representation, leaving within-cell structure unchanged.The proposition separately establishes the tilted feature marginal and unchanged conditional distribution.
  • E.1.1 Feature-space correction: Feature-based RL uniquely reweights feature cells to the target marginal while preserving the base model’s conditional distribution within each cell.The KL characterization identifies this distribution as the unique minimizer, and the result extends beyond the density setting through regular conditional distributions.
  • E.1.2 Feature-Space Test-Function Bound: The feature-space test-function bound decomposes approximation error into KL-RL suboptimality and error in the feature-space density ratio.The RL term is controlled through KL and Pinsker’s inequality, while the second term arises because the test function depends only on the representation.
  • F Adjoint Matching: Background: Adjoint matching formulates KL-regularized RL as stochastic optimal control, learning a drift correction whose quadratic control energy equals the path-KL regularizer.For flow models, the control is parameterized through a learned velocity field that can be used with the standard deterministic ODE at inference.
  • F Adjoint Matching: Background: The memoryless schedule removes endpoint-tilt bias by making the initial and terminal states independent, recovering the desired reward tilt and permitting inference under other noise schedules, including the deterministic ODE.Non-memoryless training schedules instead lock the fine-tuned model to their specific inference schedule.
  • F Adjoint Matching: Background: Lean adjoint matching drops control-dependent adjoint terms, retaining reward anchoring through the terminal condition while reducing variance and improving empirical performance.The resulting loss has the same minimizer as the basic adjoint-matching objective.
  • G Local Linear Integrator for the Memoryless SDE: The local linear SDE integrator addresses numerical stiffness near t →0 in memoryless flow-model training and applies to path-based RL methods beyond adjoint matching.The appendix motivates it as a training-time integrator for memoryless-SDE sampling, including methods such as Flow-GRPO (Liu et al., 2025a).

G.1 The Memoryless SDE and the Local Linear Integrator … Backward adjoint ODE

The local linear integrator exactly handles the stiff linear component of the memoryless SDE at unchanged per-step cost, preserving the desired tilted terminal marginal while enabling standard ODE testing. Across models and step counts it improves sample quality and training efficiency, and the same splitting yields a competitive backward-adjoint discretization.

  • G.1 The Memoryless SDE and the Local Linear Integrator: The memoryless SDE guarantees the optimal terminal marginal p*(x) ∝ pbase(x) exp(r(x)) while allowing standard ODE sampling at test time.This property motivates using the memoryless SDE during KL-regularized RL training.
  • G.1 The Memoryless SDE and the Local Linear Integrator: The stiffness arises because the drift coefficient −1/t and diffusion coefficient 2(1 − t)/t diverge as t → 0+, making naive explicit steps unstable.Offset Euler–Maruyama shifts the singularity, but early-time noise can still dominate and drive trajectories far from the training distribution.
  • Forward SDE: The integrator freezes vθ at each interval’s left endpoint, then integrates the remaining linear SDE in closed form using an integrating factor.It has the same computational cost as Euler–Maruyama—one network evaluation per step—while exactly handling the stiff linear part.
  • G.1 The Memoryless SDE and the Local Linear Integrator: At K = 50, local linear produces coherent JiT images under matched Brownian motion, whereas offset Euler–Maruyama yields visibly degraded samples.The same comparison demonstrates stability at identical noise and discretization settings.
  • G.1 The Memoryless SDE and the Local Linear Integrator: Local linear with δ = ε = ∆t consistently matches or improves Euler–Maruyama across four ImageNet models and three feature spaces, especially at low K.The sole exception is RAE under Inception-v3, where Euler–Maruyama slightly outperforms local linear.
  • G.2 Derivation of the Local Linear Integrator: Lemma G.1 solves linear SDEs with deterministic time-varying drift, forcing, and noise via an integrating factor, yielding Gaussian transitions with an explicit variance.The ODE case is recovered when h ≡ 0.
  • Backward adjoint ODE: The backward-adjoint variant integrates the stiff scalar term exactly and freezes the known endpoint integrand at τ = tk+1, producing equation (38).In practice it performed at least as well as Euler discretization, though its gains were smaller than for the forward integrator.

H Additional Implementation Details

DRL uses buffered discriminator and adjoint-matching updates with a local linear integrator, reducing sample-generation cost while retaining explicit control over training-time offsets and stability. Additional choices include separate training and inference samplers, stratified timesteps, and adaptive outlier clipping.

  • Training versus inference: Training uses the memoryless SDE with the local linear scheme, whereas inference uses the standard deterministic ODE; their differences were small but non-negligible.The algorithm uses shifted-time offsets δ and ε, with δ = ε = ∆t in practice.
  • Buffered discriminator training: Discriminator training refreshes a buffer of paired real and base-model samples, then performs multiple logistic-loss updates with an R1 penalty through the frozen encoder.The buffer contains ND pairs and is reused for ED epochs before refreshes.
  • Buffered adjoint matching: RL similarly buffers trajectories and adjoints as stop-gradient tensors, updating only current velocity evaluations; this off-policy procedure matched on-policy training, while importance weighting added no measurable improvement.Trajectories use the local linear forward integrator, and adjoints use its backward counterpart.
  • Other implementation details: The adjoint-matching loss samples 40% of timesteps per update, split evenly between early and late trajectory regions so both ends remain represented.The default split is t < 0.6 versus t ≥0.6.
  • Other implementation details: Adaptive outlier clipping masks samples with non-finite or above-threshold per-sample norms, using EMA thresholds and a spike cap c = 5; the cap is especially important for RAE.Parameters are α = 0.9, ρ = 0.9 for SiT-XL/2, REPA SiT-XL/2, and JiT-H/16, versus α = 0.75, ρ = 0.995 for RAE DiTDH-XL.

I Experimental Setup … I.7 Discriminator architecture and training ablation setup

Experiments span four pretrained ImageNet generators and evaluate DRL through distributional alignment, reward transfer, preference-based RL, distillation, feature-space choices, and discriminator-training ablations. The studies use frozen pretrained representations by default, standardized sampling and evaluation protocols, and controlled variations in downstream training setups.

  • I Experimental Setup: Experiments use four pretrained 256 × 256 ImageNet generators spanning latent- and pixel-space flow/diffusion architectures: SiT, JiT, REPA SiT, and RAE.The models range from 0.72B to 1.25B parameters and include VAE-latent, pixel-space, REPA-regularized, and semantic-latent designs.
  • I Experimental Setup: Evaluation reports distributional metrics in DINOv2, DINOv3, SigLIP, and Inception-v3 spaces using 50 000 balanced ImageNet samples against an equally sized reference set.Sampling generally uses Heun-50 with a linear schedule, while RAE follows its original scheduler.
  • I.1 Distributional alignment setup: DRL uses a frozen-encoder discriminator followed by KL-regularized RL via adjoint matching, with shared checkpoints supporting distributional-alignment and image-quality-transfer evaluations.The default discriminator uses frozen DINOv2-Large CLS features and a class-conditional linear projection head whose logit supplies the RL reward.
  • I.2 Better image quality (reward-transfer evaluation) setup: Image-quality transfer reuses distributional-alignment checkpoints without additional training and evaluates four held-out preference rewards over 50 000 balanced generated images per checkpoint.The evaluated rewards are ImageReward, PickScore, Aesthetics v2.5, and HPSv2.
  • I.3 Preference-based RL setup: Preference-based RL compares Base+PRL with DRL+PRL, sweeping λ_PRL ∈ {1, 10, 40}; for DRL+PRL, the DRL checkpoint is both initialization and KL reference.PRL uses the same adjoint-matching setup as distributional alignment but runs for 5 000 optimizer steps.
  • I.4 Distillation from RL teachers setup: The Stable Diffusion 1.5 distillation study trains an adjoint-matching RL teacher, then fits a fresh student with standard DDPM ε-prediction on teacher samples at CFG=1.The plotted student stopped after approximately 9 000 optimizer steps, seeing roughly 50× more prompt–teacher-image pairs than the RL teacher used.
  • I.5 Distillation from DRL Teachers: The DRL distillation study initializes a REPA SiT-XL/2 student from its base checkpoint and trains standard flow matching on frozen-teacher samples for 900k steps.The student receives approximately 52M teacher samples, over 150× the samples consumed during DRL post-training and roughly 40× ImageNet’s size.
  • I.6 Feature-space ablation setup: Ablations vary the reward embedder across seven feature spaces and test MLP, fine-tuned, and from-scratch discriminator variants while preserving the class-conditional projection wrapper.Embedder updates use 10^-5 for fine-tuning and 10^-4 from scratch; all other settings follow the distributional-alignment setup.

J Extended Results … RAE DiTDH-XL

Extended results show that DRL consistently improves distributional alignment and held-out reward across architectures, while revealing guidance, discriminator, and preference-RL tradeoffs. Qualitative grids further show how DRL strength and preference fine-tuning affect image appearance, including increasing brightness and distortion at higher preference-RL strength.

  • J Extended Results: Autoguidance in RL produces mixed results: low w can improve DRL at high λ, preference alignment favors w>1 for reward, and both generally underperform interval CFG or worsen image quality.Consequently, autoguidance is omitted from the main text and presented only in extended Pareto plots and sample grids.
  • J.1 Alignment: Quantitative λ Sweep; J.2 Alignment: Full Distribution Metrics (PRDC): DRL consistently improves distributional metrics over the base model, with λ ∈ {5, 10} generally strongest, FID favoring smaller λ, and gains persisting on FDval and KD.λ=40 is never selected as best, and improvements on unoptimized metrics suggest genuine distributional alignment rather than overfitting.
  • J.3 Image Quality: Reward Improvement vs. λ: Held-out reward improvements are largely monotone through λ=20, strongest for ImageReward, HPSv2, and PickScore, but often dip at λ=40.Figure 19 reports normalized DRL-minus-Base reward improvements across λ under no guidance and best CFG.
  • J.4 Image Quality / Preference RL: Pareto Plots: DRL achieves the best Pareto fronts for HPSv2, ImageReward, and Aesthetic v2.5 across nearly every plot and remains robust under no CFG, CFG=2, and autoguidance=2.Autoguidance increases reward but substantially worsens image-statistic distortion, such as REPA maximum brightness rising from approximately 0.6 to 0.7.
  • J.5 Ablations: Discriminator Full R1 Sweep: At λ=1, every R1>0 value hurts DINOv2-L FD, whereas at λ=10 a small R1 rescues frozen-feature heads from collapse.This supports R1 as a stabilizer for aggressive λ, and additional embedders leave the conclusions unchanged.
  • J.7 Effect of CFG on Reward Scores and Fréchet Distance: DRL raises reward at every CFG level, sometimes without CFG, but FD responses vary by embedder: Inception and SigLIP favor smaller λ, while DINOv2 and DINOv3 favor larger λ.CFG becomes less effective as λ grows, likely because only the conditional branch is fine-tuned; RAE’s small-λ reward decrease disappears at larger λ.
  • J.6 Ablations: Full Feature-Space Sweep: Training the discriminator in a feature space consistently improves alignment metrics evaluated in that embedder, while DINOv2-L provides the best cross-feature transfer.The result comes from the full feature-space ablation across training and evaluation embeddings.
  • K Qualitative Samples; Effect of DRL Strength.; RL Fine-tuning Sample Images.; K.1 Base vs. Fine-tuned; K.2 Effect of DRL Strength; RAE DiTDH-XL: Matched-seed qualitative grids compare base, DRL, and preference-RL variants across architectures and guidance settings; increasing λPRL makes images progressively brighter and more distorted.The grids include base-versus-fine-tuned samples, DRL-strength sweeps, and RAE DiTDH-XL examples, with visual differences attributable to weights and guidance.

K.3 RL Fine-tuning Sample Images

The section presents RL fine-tuning sample images across JiT H-16, SiT-XL/2, REPA SiT-XL/2, and RAE.

  • K.3 RL Fine-tuning Sample Images: RL fine-tuning samples are shown for JiT H-16.
  • K.3 RL Fine-tuning Sample Images: RL fine-tuning samples are shown for SiT-XL/2.
  • K.3 RL Fine-tuning Sample Images: RL fine-tuning samples are shown for REPA SiT-XL/2.
  • K.3 RL Fine-tuning Sample Images: RL fine-tuning samples are shown for RAE.
Loading 2606.19162v1…