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
TL;DR
Flow-matching losses may poorly capture sample-level properties despite fitting training-time fields. DRL uses a pretrained-representation discriminator as a KL-regularized RL reward and consistently improves distributional alignment, preference metrics, and performance–drift tradeoffs.
Problem
Method
DRL trains a discriminator in a pretrained representation space and uses its logit as the reward for KL-regularized reinforcement learning.
Source: §1 · p. 2Results
DRL consistently improves distributional alignment, preference metrics, and performance–drift tradeoffs under subsequent preference-based reinforcement learning.
Source: §7 · p. 11Takeaways & Limitations
RL post-training can complement matching objectives by recovering structure in training data that flow-based models imperfectly capture.
Source: §1 · p. 2Because DRL corrects only representation-visible discrepancies, it generally cannot target the data distribution exactly.
Source: §4 · p. 6
Abstract
from arXiv · showhide
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- and score-matching objectives can undercapture data-grounded visual and structural properties, while introducing Discriminator-Guided RL (DRL) to recover them without preference data. DRL uses a pretrained representation-space discriminator as a reward and improves base models and subsequent preference-based post-training.
- RL post-training is commonly used to recover visual realism, coherent object structure, and physically correct motion already present in data, despite preference rewards being costly and subjective.This motivates asking why matching objectives miss these properties and whether RL can correct them without preference-based rewards or data.Source: §1 · p. 2
- FSM losses may fail to control data properties because they are evaluated under interpolation marginals q_t, whereas sampling follows the model’s rollout distribution p_t.Under regularity assumptions, the guarantee further worsens with the reward’s ℓ2 Lipschitz constant; RL avoids both obstructions by evaluating its own trajectories and following the reward landscape.Source: §1 · p. 2
- DRL uses a pretrained self-supervised representation-space discriminator’s logit as the reward for KL-regularized RL, estimating data-to-model density ratios without preference data.The representation restricts corrections to semantically meaningful discrepancies while making density-ratio estimation tractable.Source: §1 · p. 2
- DRL improves SiT, JiT, REPA, and RAE across multiple feature-space Fréchet distances, also improving held-out preference rewards without training on them.Applying preference-based RL after DRL produces a better reward–distortion Pareto frontier than applying it to the base model.Source: §1 · p. 2
- The broader conclusion is that RL post-training can complement matching-based training by recovering structure in the training data, not merely optimizing external preferences.The standard pipeline first fits a base model with flow or score matching and then uses RL to tilt it toward rewarded regions.Source: §1 · p. 2
2 Preliminaries
Flow and score matching transport Gaussian noise to the data distribution by learning velocity or score fields from interpolating marginals. KL-regularized RL then shifts a pretrained flow model toward high-reward samples, using path-distribution regularization because endpoint densities are unreliable.
Flow and Score Matching
- Flow and diffusion models transport Gaussian noise to the data distribution through an interpolation Xt = α(t)X1 + β(t)X0 with prescribed endpoint conditions.The induced marginals interpolate between q0 = N(0, I) and q1 = q.Source: §2 · p. 3
- Sampling reduces to learning conditional velocity or score fields, which admit simulation-free ℓ2 regression through conditional flow matching or denoising score matching, and are mutually recoverable under the interpolation.Thus, learning either field suffices, denoted LFSM when the distinction is unimportant.Source: §2 · p. 3
RL Post-training for Flow Models
- KL-regularized RL moves a pretrained model toward high-reward regions while limiting deviation, equivalently targeting p∗(x) ∝ exp(λr(x)) pbase(x).The parameter λ controls the reward–KL trade-off.Source: §2 · p. 3
- Because endpoint densities are unreliable for flow models, RL replaces endpoint regularization with analytically tractable path-distribution KL using Girsanov’s theorem.With a memoryless noise schedule, the path-KL minimizer has the same endpoint as the original objective and can be optimized with REINFORCE, PPO, GRPO, or adjoint-based methods.Source: §2 · p. 3
3 Motivation: Understanding the Limitations of Flow and Score Matching
Flow and score matching can achieve low training loss while poorly recovering reward-relevant properties because rollout errors compound off the training marginals and velocity-space errors may be misaligned with sample-space rewards. RL avoids these obstructions by evaluating generated samples and optimizing a reward directly, which is supported by both theoretical analysis and a teacher–student experiment.
Worst-case limitation
- Low FSM loss can be arbitrarily uninformative about the reward gap: fields with uniformly small training error can produce endpoint laws whose reward remains near its worst-region value.Proposition 3.1 establishes this for every ε > 0 and δ > 0, regardless of the data distribution’s reward.Source: §3 · p. 4
Distribution shift
- The core failure is train–test mismatch: early rollout errors move trajectories into low-training-density regions, where later errors compound and can miss high-reward states.Figure 2 depicts rollout paths drifting away from the noising band, while the learned field’s error grows farther from the training marginals.Source: §3 · p. 4, Figure 2 · p. 4
Regularity-based guarantee
- Under Lipschitz regularity, the reward gap is bounded by velocity error, error-compounding through the ODE, and reward sensitivity, with the dependence essentially tight and an on-policy analogue also holding.The guarantee assumes Lipschitz reward and optimal velocity-field regularity plus uniform control of the learned velocity error; the bound’s factors are described in the surrounding analysis.Source: §3 · p. 4
Geometry and optimization mismatch
- Velocity-space matching is poorly aligned with sample rewards: sharp rewards can change substantially under small sample-space perturbations, while irreducible regression noise can prevent FSM from reaching the required small velocity error.The analysis identifies the reward Lipschitz constant as the sensitivity factor and notes that the optimization signal vanishes near the optimum while conditional-variance noise remains.Source: §3 · pp. 4–5
RL’s edge
- RL evaluates generated samples without train–test mismatch and optimizes reward-improving directions directly; empirically, score-matching distillation from an ImageReward-trained RL teacher stalls below the teacher’s reward.The idealized KL-regularized analysis directly controls the property gap, while the Stable Diffusion 1.5 teacher–student comparison shows score matching does not recover the reward-relevant properties learned by RL.Source: §3 · p. 5
4 Method: Discriminator-Guided RL
Discriminator-Guided RL (DRL) derives a reward from data–model discrimination rather than human preferences, then uses it for KL-regularized flow-model reinforcement learning. A pretrained representation constrains correction to semantically visible structure, yielding the smallest KL change that matches target representation distributions.
Motivation
- DRL replaces expensive, potentially orthogonal human-preference rewards with a reward derived directly from the target data distribution q.Preference rewards may correlate with q-properties but can also optimize unrelated aesthetic dimensions and are difficult to collect or define broadly.Source: §4 · p. 5
From samples to rewards
- The ideal KL-regularized reward is the log density ratio between target q and reference pbase, and a discriminator estimates it from samples via its logit.The optimal discriminator satisfies D*(x)=q(x)/(q(x)+pbase(x)), whose logit recovers the ideal reward up to additive constants.Source: §4 · pp. 5–6
Using representations
- DRL estimates the reward in a frozen pretrained representation space, reducing statistical difficulty and restricting discrimination to structure visible through the encoder.Direct output-space discrimination can exploit artifacts that separate distributions without reflecting relevant properties.Source: §4 · p. 6
Constrained target
- Because of this representation restriction, DRL generally cannot target q exactly; instead, it makes the smallest KL change to pbase whose representation distribution matches q’s.Variation invisible to the encoder remains unchanged, so choosing the encoder determines which aspects of q DRL can correct.Source: §4 · p. 6
The DRL pipeline
- The 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 reward as the discriminator logit; Stage 2 applies it to the base model.Source: §4 · p. 6
5 Experiments
Across four ImageNet-pretrained flow models, DRL improves distribution alignment and held-out preference rewards without preference-data training. It also provides a stronger foundation for preference optimization, while ablations identify the effective guidance range and the importance of pretrained discriminator features.
Distribution Alignment
- Tuned DRL improves FD in all 16 model–feature pairs without CFG and 14/16 with best CFG, with especially large gains in DINOv2 and DINOv3.The only clear CFG exceptions are InceptionV3 for SiT and REPA; tuning λ consistently outperforms λ = 1.Source: Figure 5 · p. 7, §5 · p. 7
Guidance-Weight Analysis
- FD is minimized around λ = 5–10 but degrades at higher λ as coverage drops, so λ = 10 is adopted for downstream experiments.At λ = 40, the precision–coverage trajectory collapses off the frontier because the model overfits the discriminator reward.Source: Figure 6 · p. 7, §5 · p. 7
Image Quality Without Preference Data
- DRL improves every held-out preference reward on every architecture without seeing preference data, with the largest gains for SiT and REPA.The evaluated rewards are ImageReward, PickScore, Aesthetics v2.5, and HPSv2; samples also become sharper and more structurally coherent as λ increases.Source: Figure 7 · p. 8, §5 · p. 8
Foundation for Preference Alignment
- DRL+PRL achieves higher HPSv2 at comparable brightness, saturation, contrast, colorfulness, and whiteness drift than Base+PRL across the reward–drift frontier.Across fixed-seed samples, DRL+PRL remains structurally coherent and naturally colored as PRL strength increases, avoiding the base model’s brightness and whiteness drift.Source: §5 · pp. 8–9, Figure 9 · p. 9
Flow Matching Alone
- After more than 900k flow-matching steps and over 50 million samples, a student trained on DRL outputs still fails to match the DRL teacher.The student plateaus above the teacher across all four feature spaces, supporting the importance of the RL objective rather than only the teacher’s samples.Source: §5 · p. 9, Figure 10 · p. 10
Design Choices
- Frozen pretrained features outperform from-scratch or fine-tuned discriminators, while a small R1 penalty stabilizes λ = 10 training and DINOv2-L is the strongest tested embedder.Linear and MLP heads perform similarly, so the linear head is preferred; R1 helps at λ = 10 but degrades performance at λ = 1.Source: §5 · p. 10
6 Related Work
The paper connects DRL to three related-work threads: imitation learning, inverse reinforcement learning, and calibration. It extends these ideas to generative flow and score matching while emphasizing geometric obstructions and high-dimensional feature-space distribution alignment.
RL vs. Imitation Learning
- Flow and score matching can exhibit imitation-learning pathologies, with continuous-time training offering weaker guarantees than standard DAgger-style bounds (Ross et al., 2011).The paper extends the study of on-policy methods beyond demonstration-based imitation learning to generative models.Source: §6 · p. 10
Inverse RL
- DRL mirrors inverse RL by learning a reward from data-model distinctions, but its benefit appears tied to embedder-enabled geometric correction rather than solely closing the train–test gap (Ho and Ermon, 2016; Fu et al., 2018).Unlike prior GAIL and AIRL applications to standard control, DRL adapts the inverse-RL structure to generative models.Source: §6 · p. 11
Calibration
- DRL addresses calibration through a KL-constrained model whose high-dimensional self-supervised feature pushforward matches the data, using a discriminator logit as the RL reward rather than explicit moment constraints (Khalifa et al., 2021; Shen et al., 2023; Smith et al. (2025)).This differs from Smith et al. (2025) by targeting full feature-space distributions instead of user-specified moments.Source: §6 · p. 11
7 Conclusions … A Detailed Related Work
The paper frames matching objectives and RL as complementary tools: DRL uses on-policy evaluation and discriminator rewards to improve alignment, preference metrics, and performance–drift tradeoffs. Related work situates DRL against imitation-learning corrections, adversarial diffusion methods, reward-optimization algorithms, and calibration-based approaches, while the appendix provides proofs, implementation details, experimental setups, extended results, and qualitative samples.
7 Conclusions
- DRL consistently improves distributional alignment, preference metrics, and performance–drift tradeoffs by evaluating flow-based models on their own samples and exploiting reward geometry rather than ℓ2 distance.The conclusion presents matching objectives and RL as complementary tools for extracting more from the data.Source: §7 · p. 11
- The paper identifies frozen self-supervised representations, worst-case theoretical bounds, and uncertain sampler-level obstructions as limitations, while proposing jointly learned representations, sharper analyses, and alternatives such as MMD.It also presents RL as one instance of a broader strategy based on on-policy evaluation and non-ℓ2 geometry.Source: §7 · p. 11
Appendix Contents
- The appendix contains proofs for preliminaries, matching limitations, DRL, and representation spaces, plus adjoint-matching background, local-linear SDE integration, implementation details, and experimental setups.The contents list separate sections for distributional alignment, reward-transfer, preference-based RL, distillation, feature-space, and discriminator-training setups.Source: Appendix Contents · pp. 16–17
- The appendix reports extended λ sweeps, full distribution metrics, reward-versus-λ results, Pareto plots, discriminator and feature-space ablations, CFG effects, and qualitative comparisons of base, fine-tuned, and RL samples.These materials cover both quantitative and visual follow-up analyses, including DRL strength effects and RL fine-tuning samples.Source: Appendix Contents · pp. 16–17
A Detailed Related Work
- DRL connects flow-matching train–test mismatch to covariate shift, exposure bias, and imitation learning, replacing model-specific off-distribution corrections with an explicitly on-policy solution.The related-work discussion contrasts DRL with perturbations, shifted time steps, transition-norm changes, and short warm-started MMD rollouts that address only local mismatch.Source: A Detailed Related Work · p. 18
- Unlike GAIL, AIRL, and diffusion GAN methods, DRL trains one discriminator on actual model samples and data in a representation space, then uses its fixed logit reward to correct the velocity field directly.This avoids alternating min–max optimization and distinguishes DRL from posterior-mean discrimination, which does not generally match the model’s sample distribution.Source: A Detailed Related Work · pp. 18–19
- Effective reward post-training requires on-policy optimization that propagates reward information through ∇xr or score-function estimators, whereas reward-weighted regression and Tilt Matching retain matching geometry and may remain off-policy.The paper offers this geometry and on-policy distinction as a possible explanation for reported underperformance of RWR and Tilt Matching, while calling for further validation.Source: A Detailed Related Work · p. 19
- DRL is related to calibration and CGM-reward but targets the full representation pushforward rather than user-specified moments; the two exponential tilts coincide only in a special linear case and generally estimate different targets.The comparison emphasizes that DRL recovers its tilt from a discriminator logit, whereas CGM-reward solves an empirical maximum-entropy dual.Source: A Detailed Related Work · p. 19
B Limitations
DRL remains limited by practical RL costs and instability, dependence on the chosen representation, and a theory that simplifies sampler dynamics. Its reverse-KL behavior and training–inference mismatch further constrain how broadly its guarantees and gains should be interpreted.
Classifier-free guidance
- DRL cannot cleanly incorporate classifier-free guidance: training with CFG often became unstable, and its gains over the base model were smaller when both used their best CFG scales.The authors attribute this to current RL algorithms’ limited CFG compatibility and identify principled CFG-compatible RL training as future work; Domingo-Enrich et al. (2025).1 reported similar difficulties.Source: B Limitations · p. 20
Choice of λ
- λ=10 was a reliable default, but performance depends substantially on λ, requiring exploration that increases computational cost and may not generalize across settings.The authors propose improving robustness to λ and amortizing hyperparameter tuning, for example through a ControlNet-like mechanism (Zhang et al., 2023).Source: B Limitations · p. 20
Reverse KL behavior
- At large λ, DRL’s reverse-KL objective can make the model mode-seeking, although this was not a practical issue at the recommended λ values.The authors suggest this behavior partly explains why R1 gradient regularization helped at higher λ.Source: B Limitations · p. 20
Additional compute
- DRL requires additional post-training compute beyond base-model training, and that 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 more broadly.Source: B Limitations · p. 20
Role of the representation
- DRL can correct only discrepancies visible in its chosen pretrained representation, so its effectiveness depends on a suitable embedder and may partly reflect altered reward geometry rather than on-policy optimization alone.The representation focuses rewards on semantic discrepancies, but the empirical importance of the embedder suggests geometry changes may account for a substantial share of the benefit.Source: B Limitations · p. 20
Scope of the theory
- The theory uses a simplified probability-flow ODE, while training uses a memoryless SDE and evaluation an ODE; their distributions need not coincide during training, limiting direct guarantees for inference samples.The reward certificate may have SDE analogues under suitable drift control, but the ODE-specific worst-case construction may not extend to stochastic samplers; large practical discrepancies were not observed.Source: B Limitations · pp. 20–21
C Proofs for Section 2: Preliminaries … D.1 Distribution Shift and Error Accumulation
The appendix proves foundational identities for flow and score matching, including score–velocity equivalence and SDE–ODE marginal equivalence. It also introduces the proof sequence for Section 3’s limitations analysis, reward geometry, and RL regret bound.
C Proofs for Section 2: Preliminaries
- The appendix collects proofs supporting the statements in Section 2’s preliminaries.Source: C Proofs for Section 2: Preliminaries · p. 22
C.1 Score-velocity relationship
- The velocity field and score are related through conditional expectations under the interpolation marginal.The proof computes E[X1 | Xt = x] and E[X0 | Xt = x], using the interpolation and Tweedie’s identity.Source: §C.1 · p. 22
- The velocity at x is defined as vt(x) = ˙α(t)E[X1 | Xt = x] + ˙β(t)E[X0 | Xt = x], using conditional posterior means for X1 and X0.This is exactly (C.1).Source: §C.1 · p. 22
C.2 SDE and ODE marginal equivalence
- The true SDE and probability-flow ODE generate identical marginals qt at every time for any noise schedule σ(t) ≥ 0.Source: §C.2 · p. 22
- The marginal-equivalence proof verifies that qt satisfies the SDE’s Fokker–Planck equation using st qt = ∇qt.The diffusion and score terms cancel, reducing the equation to the continuity equation for the probability-flow ODE.Source: §C.2 · p. 23
- Because the SDE and ODE share the same continuity equation and initial condition q0, they have the same marginals throughout time.Source: §C.2 · p. 23
D Proofs for Section 3: Understanding the Limitations of Flow and Score Matching
- The Section 3 proofs proceed from the standard probability-flow ODE’s no-certificate result to reward geometry and the RL reward-regret bound.The appendix states that these materials follow the order of the main discussion and include proposition 3.2 and proposition D.7.Source: D Proofs for Section 3: Understanding… · p. 24
D.1.1 Proof of proposition 3.1: no reward certificate for the standard probability-flow sampler … D.3.1 RL reward regret bound
This section develops theoretical results clarifying when matching errors do—and do not—certify endpoint rewards for probability-flow sampling, including counterexamples, stronger rollout-control guarantees, and tightness results. It also analyzes reward-improving perturbations and shows how KL-regularized RL regret controls divergence from a reward-tilted target distribution.
D.1.1 Proof of proposition 3.1: no reward certificate for the standard probability-flow sampler
- Small flow-matching error can still produce arbitrarily poor reward: a perturbation with uniform error ε inflates rollout variance, and k(1) diverges as τ ↓ 0.The construction uses a bounded reward concentrated on high-likelihood regions of the target endpoint distribution, which the inflated rollout misses.Source: §D.1.1 · pp. 24, 26
- The Gaussian counterexample makes the mismatch explicit: the same perturbation has error ε under q_t but εk(t) under the rollout law p_bad,t.Thus training-time marginal error can underestimate the perturbation encountered along sampling trajectories.Source: §D.1.1 · pp. 24, 26
- A general radial-shell construction disproves reward certificates for both velocity and score parametrizations under the stated flow-map assumptions.The perturbation is localized to a thin moving shell, keeping training error small while steering trajectories into a low-reward region.Source: §D.1.1 · pp. 27–28
- Working in inverse-flow coordinates freezes the true dynamics, allowing the shell to catch trajectories and move them inward before being mapped back to x-space.Every trajectory initially inside the large ball is driven into the low-reward ball by time b, after which the perturbation is switched off.Source: §D.1.1 · pp. 28, 30, 31
- The constructed bad velocity field produces exactly the target bad endpoint law while retaining arbitrarily small training error under q_t, so no reward certificate follows from standard matching error alone.The construction is supported on a thin radial shell, whose mass can be controlled by shell thickness; the same endpoint reward bound and small-error argument applies to the score parametrization.Source: §D.1.1 · pp. 32–33
D.2.1 Reward geometry under stronger rollout control
- Under a uniformly bounded field mismatch and an L_v-Lipschitz reference field, every L_r-Lipschitz reward admits a global endpoint certificate obtained by Grönwall’s inequality.The bound propagates direct field error through the factor e^(L_v(1−s)) and converts endpoint displacement into reward error via L_r.Source: §D.2.1 · pp. 34, 36
- The same endpoint certificate holds under on-policy control when field mismatch is bounded only after averaging over rollout marginals, using an error function integrated through the reference dynamics.This extends the global result from pathwise mismatch control to expected on-policy mismatch control.Source: §D.2.1 · pp. 35–37
- The exponential weighting factor e^{Lv(t−s)} and its dependence on Lv, Lr, and the error function cannot generally be improved; explicit one-dimensional constructions attain equality.For bounded rewards, the amplification-tight construction is essentially tight. Theorem D.6 states a universal constant c > 0 for ε > 0 and L > 0 with εL < 1, using standard linear interpolation and a reward r : R → [0, 1] with Lip(r) ≤ L.Source: §D.2.1 · pp. 35, 38, 39
- A bounded L-Lipschitz triangular-wave reward and an uphill perturbation yield expected reward gain at least Lε/2 under stronger rollout control.The construction moves trajectories toward nearby reward peaks at speed ε; selecting an appropriate phase establishes the bound with constant c = 1/2.Source: §D.2.1 · p. 40
D.3.1 RL reward regret bound
- RL regret equals a trajectory-level KL gap, while data processing transfers this control to the final-sample marginal relative to the tilted target.The optimizer’s final-sample marginal is the reward-tilted target p*, and the proof uses the endpoint data-processing inequality.Source: §D.3.1 · pp. 40–41
E Proofs for Section 4: Method: Discriminator-Guided RL … G Local Linear Integrator for the Memoryless SDE
The proofs show that feature-based rewards reweight representation-space cells while preserving within-cell structure, with bounds controlled by RL suboptimality and feature-ratio error. The appendix then explains how memoryless-SDE adjoint matching removes endpoint bias and how a local linear integrator stabilizes sampling for flow-model RL.
E.1.1 Feature-space correction
- Feature-space RL uniquely minimizes KL under a target feature marginal by reweighting feature cells while preserving the base model’s within-cell conditional distributions.The result follows from the KL chain rule and extends beyond densities through regular conditional distributions.Source: §E.1.1 · p. 44
E.1.2 Feature-Space Test-Function Bound
- A feature-based reward targets only the induced feature-space distribution, so test-function error separates into RL suboptimality and feature-space density-ratio error.The bound applies when the learned policy is Δ-suboptimal for the KL-regularized objective and the test function is bounded.Source: §E.1.2 · p. 44
F Adjoint Matching: Background
- Adjoint matching formulates KL-regularized RL as stochastic optimal control, learning a drift correction whose quadratic control cost equals the path-KL penalty.For flow models, the control is parameterized through a learned velocity field.Source: F Adjoint Matching: Background · p. 45
- The memoryless noise schedule removes the initial-state-dependent endpoint bias, 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.Source: F Adjoint Matching: Background · p. 45
- Lean adjoints retain the terminal reward condition and the same optimizer as basic adjoint matching while dropping control-dependent terms to reduce variance and improve empirical performance.The resulting regression remains reward-informed despite the simplification.Source: F Adjoint Matching: Background · p. 46
- Flow-model specialization converts controlled-drift fine-tuning into learning a velocity field that plugs directly into the standard deterministic ODE at inference.The corresponding memoryless training SDE introduces numerical stiffness near t → 0, motivating the local linear integrator.Source: F Adjoint Matching: Background · p. 46
G Local Linear Integrator for the Memoryless SDE
- The local linear SDE integrator avoids the early-time norm blow-up seen with Euler SDE sampling and stays close to the Euler ODE reference across models.The comparison uses latent norm trajectories from 1000 sampled trajectories; the integrator is intended for memoryless-SDE sampling in path-based flow-model RL.Source: Figure 13 · p. 47, G Local Linear Integrator for the… · p. 47
- Because of its simplicity and superior performance, the local linear integrator is proposed as a standard component for methods sampling the memoryless SDE, including non-adjoint approaches such as Flow-GRPO (Liu et al., 2025a).The stated scope covers path-based RL methods on flow models, not only adjoint matching.Source: G Local Linear Integrator for the… · p. 47
G.1 The Memoryless SDE and the Local Linear Integrator … Backward adjoint ODE
The paper replaces unstable Euler–Maruyama discretization of the memoryless SDE with a local linear integrator that solves the stiff linear component exactly, improving sampling and training while preserving the desired terminal marginal. The same integrating-factor principle yields a backward-adjoint update that performs at least as well as Euler discretization.
G.1 The Memoryless SDE and the Local Linear Integrator
- The memoryless SDE guarantees the KL-regularized optimum has terminal marginal p*(x) ∝ pbase(x) exp(r(x)) while permitting standard-ODE sampling at test time.Training therefore uses the memoryless SDE even though inference can switch back to the standard ODE.Source: §G.1 · p. 47
- Near t = 0, the SDE’s linear drift and diffusion diverge, so naive explicit discretization struggles and can produce early-time trajectory blow-up.Euler–Maruyama’s noise increment can dominate the rapid stiff contraction, driving trajectories far from the training distribution and away from the ODE norm distribution.Source: §G.1 · pp. 47–48
- At K = 50, local linear produces coherent JiT images while Euler–Maruyama with offsets produces degraded samples under identical Brownian motion.The local linear method also matches the K = 100 JiT training reward, whereas Euler–Maruyama at K = 50 barely improves and underperforms at K = 100.Source: Figure 14 · p. 48, §G.1 · pp. 48–49, Figure 15 · p. 48
- The local linear integrator freezes the velocity over each interval and integrates the remaining time-dependent linear SDE in closed form with one network evaluation per step.Its contraction factor Φ_k = u_k/u_{k+1} stays in (0, 1), damping the stiff early-time dynamics without problematic singularities.Source: §G.1 · p. 48
- Across four ImageNet models and three feature spaces, local linear Pareto-dominates Euler–Maruyama with offsets, especially at low K, except for RAE under Inception-v3.Figure 16 evaluates endpoint Fréchet distances using 50,000 samples; the stated exception is that Euler–Maruyama slightly outperforms local linear for RAE under Inception-v3.Source: Figure 16 · p. 49, §G.1 · p. 49
G.2 Derivation of the Local Linear Integrator
- Lemma G.1 provides the integrating-factor solution for scalar linear SDEs, including Gaussian stochastic-integral variance, with the ODE recovered when h ≡ 0.The derivation applies Itô’s product rule and Itô’s isometry to obtain the closed-form solution and variance formula.Source: §G.2 · p. 50
Forward SDE
- For the forward SDE, freezing vθ at the left endpoint permits exact integration of the remaining drift and converts the noise term into a centered Gaussian.Setting δ = ε = Δt recovers the forward update used by the local linear method.Source: Forward SDE · p. 51
Backward adjoint ODE
- For the backward adjoint ODE, the method integrates the stiff scalar term exactly and freezes the known endpoint integrand at τ = t_{k+1}, yielding equation (38).In practice, the backward-adjoint step performs at least as well as Euler discretization, though its gains are smaller than those of the forward integrator.Source: §G.1 · pp. 49–50, Backward adjoint ODE · pp. 51–52
H Additional Implementation Details
The implementation uses buffered training for both discriminator and RL stages, with local linear forward and adjoint integration. Additional stabilizers include distinct training and inference samplers, stratified timestep sampling, and adaptive clipping of outlier adjoints.
Buffered discriminator training
- Buffered discriminator training refreshes paired real and base-model samples periodically, then applies logistic-loss updates with an image-space R1 penalty through the frozen encoder.The image-space penalty was chosen over a flow-space penalty for memory efficiency and effective performance.Source: H Additional Implementation Details · p. 53
Buffered adjoint matching
- Buffered adjoint matching stores stop-gradient trajectories and adjoints, then updates only current velocity evaluations; this off-policy scheme matched on-policy training, while importance weighting added no measurable improvement.Trajectories are generated with the local linear forward integrator and backward adjoint under the current model.Source: H Additional Implementation Details · p. 53
Training versus inference
- Training uses the memoryless SDE with a local linear integrator, whereas inference uses the standard deterministic ODE; sampler differences were small but non-negligible.The appendix also specifies a final deterministic denoising call after K −1 local linear SDE steps.Source: H Additional Implementation Details · p. 53
Other implementation details
- Each AM-loss update samples 40% of timesteps with balanced early and late coverage, while adaptive EMA-based clipping masks extreme or non-finite per-sample losses.The clipping settings vary by backbone, use spike cap c = 5, and are especially important for RAE because its loss spikes more frequently.Source: H Additional Implementation Details · p. 53
I Experimental Setup … I.7 Discriminator architecture and training ablation setup
The experiments evaluate DRL across four pretrained ImageNet generators using standardized sampling and distributional metrics, then test reward transfer, preference-based RL, distillation, feature-space choices, and discriminator designs. These studies reuse frozen-encoder discriminator rewards and controlled training protocols across alignment, transfer, and ablation settings.
I Experimental Setup
- Four pretrained 256 × 256 ImageNet generators span latent- and pixel-space architectures: SiT, JiT, REPA SiT, and RAE DiTDH.They differ in backbone, parameterization, and representation space, including SD VAE, raw pixels, REPA-aligned features, and a DINOv2-based semantic latent.Source: I Experimental Setup · p. 55
- Distributional evaluation uses 50,000 balanced ImageNet samples and reports FD, FDval, KD, Precision, Recall, Density, and Coverage across DINOv2, DINOv3, SigLIP, and Inception feature spaces.Unless otherwise indicated, sampling uses Heun with 50 steps; RAE follows its original scheduler.Source: I Experimental Setup · p. 55
I.1 Distributional alignment setup
- DRL uses a two-stage pipeline: a discriminator on frozen DINOv2-Large CLS features followed by KL-regularized adjoint-matching RL, with buffered training for both stages.The default discriminator is a class-conditional linear projection head with hidden dimension 512, and its logit supplies the RL reward.Source: §I.1 · pp. 55–56
I.2 Better image quality (reward-transfer evaluation) setup
- Image-quality transfer reuses the distributional-alignment checkpoints without additional training and evaluates ImageReward, PickScore, Aesthetics v2.5, and HPSv2 on 50,000 balanced generated images.Text-conditioned rewards use the ImageNet class label in the prompt template “a photo of a {class}”.Source: §I.2 · p. 56
I.3 Preference-based RL setup
- Preference-based RL compares Base+PRL with DRL+PRL across λ_PRL ∈ {1, 10, 40}, anchoring DRL+PRL’s KL reference to the DRL checkpoint.Both settings use ImageReward with the same class-prompt template, while PRL runs for 5,000 optimizer steps.Source: §I.3 · pp. 56–57
I.4 Distillation from RL teachers setup
- The Stable Diffusion 1.5 distillation study trains an ImageReward-optimized adjoint-matching teacher, then trains a fresh student with standard ε-prediction on teacher samples at CFG=1.The plotted student stopped after approximately 9,000 optimizer steps, despite seeing roughly 50× more prompt–teacher-image pairs than the RL teacher consumed.Source: §I.4 · pp. 57–58
I.5 Distillation from DRL Teachers
- DRL distillation initializes a student from the pretrained REPA SiT-XL/2 checkpoint and trains it with standard flow-matching velocity loss on repeatedly refreshed samples from a frozen DRL teacher.The 900,000-step run exposes the student to approximately 52M teacher samples, over 150× DRL’s post-training sample count and roughly 40× ImageNet’s size.Source: §I.5 · p. 58
I.6 Feature-space ablation setup
- The feature-space ablation varies only the reward embedder across DINOv2, DINOv3, SigLIP, and Inception variants, then evaluates every fine-tuned model in all seven feature spaces.The ablation is based on REPA SiT-XL/2 and otherwise follows the distributional-alignment discriminator and adjoint-matching settings.Source: §I.6 · p. 58
J Extended Results … RAE DiTDH-XL
Extended results show that DRL improves distributional alignment and held-out reward performance across model architectures, while preserving the main-text Pareto gains and revealing guidance, regularization, feature-space, and preference-RL tradeoffs. Qualitative samples further show that stronger preference-RL and autoguidance can increase proxy rewards while producing brighter, more distorted images.
J.1 Alignment: Quantitative λ Sweep / J.2 Alignment: Full Distribution Metrics (PRDC)
- DRL consistently improves the base model across seven distributional metrics, with λ ∈ {5, 10} generally strongest, λ=40 never selected, and gains persisting on FDval and KD.The results support genuine distributional alignment rather than overfitting to FD, which is optimized against the ImageNet training distribution.Source: J Extended Results · p. 60, Figure 17 · p. 62, Figure 18 · p. 63
J.3 Image Quality: Reward Improvement vs. λ
- Normalized held-out reward improvements are largely monotone through λ=20, strongest for ImageReward, HPSv2, and PickScore, but often dip at λ=40.Figure 19 covers the full λ sweep, while Table 4 reports the raw values for the headline λ=10 comparison.Source: J Extended Results · pp. 60–61, Figure 19 · p. 66, Table 4 · p. 66
J.4 Image Quality / Preference RL: Pareto Plots
- DRL achieves the best Pareto fronts for HPSv2, ImageReward, and Aesthetic v2.5 across no CFG, CFG=2, and autoguidance=2.Autoguidance raises proxy rewards but substantially worsens image-statistic distortion; on REPA, maximum brightness rises from ∼0.6 under base+CFG to ∼0.7 under autoguidance.Source: J Extended Results · p. 61, Figure 20 · p. 67, Figure 21 · p. 68
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 interpreting R1 as a stabilizer for aggressive λ rather than a standard vanishing-gradient remedy.Source: J Extended Results · p. 61, Figure 23 · p. 70
J.7 Effect of CFG on Reward Scores and Fréchet Distance
- DRL improves reward at every CFG level, but FD responses depend on the embedder: smaller λ favors Inception and SigLIP, larger λ favors DINOv2 and DINOv3, and CFG helps less as λ grows.The authors attribute these patterns to fine-tuning only the conditional branch, a constraint imposed by adjoint matching; RAE is the exception where small λ slightly lowers reward before larger λ removes the effect.Source: J Extended Results · p. 61, Figure 24 · p. 71, Figure 25 · p. 71
J.6 Ablations: Full Feature-Space Sweep
- Training the discriminator in a given feature space consistently improves alignment metrics in that space, while DINOv2-L provides the best cross-feature transfer.The full seven-by-seven feature-space sweep leaves the main conclusions unchanged.Source: J Extended Results · p. 61, Table 5 · p. 70
K Qualitative Samples / K.1 Base vs. Fine-tuned / K.2 Effect of DRL Strength / RAE DiTDH-XL
- Matched-seed qualitative grids compare base and DRL models across architectures and guidance settings, showing how DRL strength λ changes samples while isolating model-weight and guidance effects.The qualitative sections include base-versus-fine-tuned grids, DRL-strength sweeps, and the RAE DiTDH-XL samples.Source: K Qualitative Samples · p. 72, Effect of DRL Strength · p. 72, Figure 26 · p. 73
RL Fine-tuning Sample Images / Effect of DRL Strength. / RAE DiTDH-XL
- Increasing preference-RL strength λPRL makes images progressively brighter and more distorted, especially under autoguidance, which yields the highest reward but largest distortion.The grids compare Base, DRL+PRL, and Base+PRL at λPRL ∈ {1, 10, 40} under three guidance strategies.Source: J Extended Results · p. 60, RL Fine-tuning Sample Images · p. 72, Figure 29 · p. 76
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.
- RL fine-tuning samples are shown for JiT H-16.Source: Figure 34 · p. 81
- RL fine-tuning samples are shown for SiT-XL/2.Source: Figure 35 · p. 82
- RL fine-tuning samples are shown for REPA SiT-XL/2.Source: Figure 36 · p. 83
- RL fine-tuning samples are shown for RAE.Source: Figure 37 · p. 84
Related papers
This paper cites
- An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale2020
- Score-Based Generative Modeling through Stochastic Differential Equations2020
- Rethinking the Inception Architecture for Computer Vision2015
- Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow2022
- Adding Conditional Control to Text-to-Image Diffusion Models2023
- Sigmoid Loss for Language Image Pre-Training2023