The Reward Was in Your Data All Along: Correcting Flow Matching with Discriminator-Guided RL Nicolas Beltran-Velez1,2 , Felix Friedrich1 , Zhang Xiaofeng1,3,5 , Reyhane Askari-Hemmat1 , Xiaochuang Han1 , Adriana Romero-Soriano1,3,4,6,∗ , Michal Drozdzal1,∗
FAIR at Meta, 2 Columbia University, 3 Mila – Québec AI Institute, 4 McGill University, 5 Université de Montréal, 6 Canada CIFAR AI Chair
1
arXiv:2606.19162v1 [cs.LG] 17 Jun 2026
∗
Joint senior authors.
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 ℓ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 → 2.62 on SiT) and semantic-space FD (e.g., 88.2 → 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. Date: June 18, 2026 Correspondence: Nicolas Beltran-Velez, [email protected]
(a) Base model (REPA SiT-XL/2)
(b) After DRL (ours)
Figure 1 Seed-matched samples from REPA before and after DRL post-training with CFG = 1. Our DRL model
produces sharper, more coherent images even without using any CFG.
1
1
Introduction
Flow-based models have become the dominant paradigm for generative modeling in continuous domains due to their simulation-free, regression-based training objectives (Lipman et al., 2022; Albergo and Vanden-Eijnden, 2023; Ho et al., 2020; Song et al., 2020). In practice, however, these models are rarely trained only with these losses. Instead, they follow a multi-stage pipeline: (1) a flow or score matching (FSM) objective is used to fit a base model pbase to a data distribution q; and (2) reinforcement learning (RL) is used to tilt pbase towards desirable regions under a reward r(x). The standard motivation for RL post-training is that it enables fitting a model to the implicit distribution p∗ (x) ∝ pbase (x) exp(r(x)) for which no data exist but which is easy to specify via a reward (Xu et al., 2023; Black et al., 2024; Fan et al., 2023; Schuhmann et al., 2022). Nevertheless, in practice, RL post-training is also used to enforce properties that are already present in the data, like visual realism, coherent object structure (Wallace et al., 2024; Domingo-Enrich et al., 2025), or physically correct motion (Ye et al., 2025; Liu et al., 2025b). This is both puzzling—why can the model not learn these real-world properties from the data directly but recover them with RL?—and undesirable, as preference data, which is the norm for training the rewards in these applications, is expensive to collect and conflates the real-world properties we want to preserve with subjective concerns. This raises two questions that we focus on in this paper: (1) Why do matching-based objectives fail to capture properties that are present in the data but that RL helps recover? (2) Can we use RL to correct these failures without relying on preference-based rewards or data? First, we argue that the fact that RL recovers properties already present in q points to the matching objective, rather than data or capacity, as a plausible bottleneck. To build intuition, we study a simplified setting with probability-flow ODEs and ask when FSM losses control the property gap Epbase [r] − Eq [r] on a property r of interest. The guarantees we find are weak. In the worst case, none exist: the FSM loss is measured under the interpolation marginals qt , but sampling is governed by the model’s rollout distribution pt (Proposition 3.1). Under regularity assumptions a guarantee does exist. However, it degrades with the reward’s ℓ2 Lipschitz constant (Proposition 3.2), which we expect to be large for many properties of interest. Although these bounds are worst-case, they suggest how RL can help: by construction it avoids both obstructions, evaluating the model on its own trajectories rather than qt and following the reward landscape directly instead of the geometry of velocity space. However, RL is only useful if we optimize a reward that is a good proxy for the properties of q we want to preserve, and a priori we do not have one. To this end, motivated by the analysis above, we introduce Discriminator-Guided RL (DRL). DRL trains a discriminator to estimate the density ratio between the data and the model’s distributions in a pretrained self-supervised learning (SSL) representation space. The logit of this discriminator then serves as the reward for KL-regularized RL. Working in an SSL space is the central design choice: it restricts the reward to use discrepancies visible through the representation, both making density-ratio estimation tractable and confining corrections to semantically meaningful axes—all without ever relying on preference data. We apply DRL to state-of-the-art image generation architectures, SiT (Ma et al., 2024), JiT (Li and He, 2025), REPA (Yu et al., 2024), and RAE (Zheng et al., 2025), and show that even with a simple linear discriminator, DRL delivers large improvements over the base model across the board (Figure 1), measured by Fréchet distances in multiple feature spaces. These gains directly translate into better scores under held-out human-preference rewards, despite DRL never being trained on such rewards. Beyond improving the base model, DRL also strengthens the post-training pipeline that follows. Applying preference-based RL (PRL) on top of DRL rather than the base model yields a better reward–distortion Pareto frontier, since DRL absorbs the corrections recoverable from q and leaves PRL to handle genuine subjective preferences. Full experimental details are in Section I. Overall, our results suggest viewing RL post-training of flow-based models not only as a way to optimize external preferences, but also as a complementary mechanism for recovering structure in the training data that is imperfectly captured by standard matching-based objectives.
2
2
Preliminaries
Flow and Score Matching. Flow and diffusion models generate data by transporting a simple base distribution,
typically Gaussian noise, to the data distribution q (Ho et al., 2020; Song et al., 2020; Lipman et al., 2022; Albergo and Vanden-Eijnden, 2023; Liu et al., 2022). A standard training construction introduces independent random variables X1 ∼ q and X0 ∼ N (0, I), and defines the auxiliary interpolation (Albergo and Vanden-Eijnden, 2023; Lipman et al., 2022; Liu et al., 2022) Xt = α(t)X1 + β(t)X0 ,
where t ∈ [0, 1], α(0) = 0, β(0) = 1, α(1) = 1, β(1) = 0.
(1)
This induces a family of marginals qt interpolating between q0 = N (0, I) and q1 = q. Sampling from q then reduces to learning the velocity field vt (x) := E[α̇(t)X1 + β̇(t)X0 | Xt = x] or the score field st (x) = ∇x log qt (x), since the SDE h i dXt = vt (Xt ) + 21 σ(t)2 st (Xt ) dt + σ(t) dWt , X0 ∼ N (0, I), (2) has marginals qt for any noise schedule σ(t) (Song et al., 2020; Lipman et al., 2022; Albergo and VandenEijnden, 2023); see Section C.2. Both fields are conditional expectations, and so admit simulation-free ℓ2 regression targets (Hyvärinen and Dayan, 2005; Vincent, 2011; Ho et al., 2020; Song et al., 2020; Lipman et al., 2022; Albergo and Vanden-Eijnden, 2023): conditional flow matching (CFM) and denoising score matching (DSM) use, respectively, h i 2 2 E vθ (Xt , t) − α̇(t)X1 + β̇(t)X0 and E sθ (Xt , t) + Xt − α(t)X1 /β(t)2 , (3) with expectation over (t, X0 , X1 ). Under Eq. (1), vt and st are recoverable from each other (Karras et al., 2022; Domingo-Enrich et al., 2025) (see Section C.1), so learning one suffices; we write LFSM for either objective when the distinction is unimportant. RL Post-training for Flow Models. Given a pretrained model pbase and a reward r : X → R, KL-regularized RL (Ziebart, 2010) aims to move probability mass toward high-reward regions without drifting too far from the base model by solving maxp Ex∼p [r(x)] − (1/λ) KL(p ∥ pbase ) ,
or equivalently minimizing the reverse KL, KL(p ∥ p∗ ), where p∗ (x) ∝ exp(λr(x)) pbase (x) and λ > 0 is a hyper-parameter controlling the reward-KL trade-off.
Unfortunately, directly optimizing this objective is infeasible for flow models as the endpoint densities required by the KL term are unreliable. Instead, the standard approach is to use a nonzero noise schedule σ(t) in the sampling SDE (Eq. (2)) and replace the endpoint regularization by a KL over path distributions (Fan et al., 2023; Domingo-Enrich et al., 2025), which can be analytically evaluated using Girsanov’s theorem (Girsanov, 1960; Domingo-Enrich et al., 2025). If Pθ and Pbase denote the trajectory distributions induced by the fine-tuned and base models under the same training SDE, the new objective is then LRL (Pθ ) := −Ex∼pθ [r(x)] + (1/λ) KL(Pθ ∥ Pbase ) .
(4)
If training uses a specific, so-called memoryless σ(t), the minimizer of Eq. (4) has the same endpoint as the original KL-regularized objective, and admits sampling after training with any noise schedule (DomingoEnrich et al., 2025). This objective can then be optimized with standard policy gradient methods such as REINFORCE, PPO, or GRPO (Williams, 1992; Black et al., 2024; Fan et al., 2023; Schulman et al., 2017; Shao et al., 2024; Liu et al., 2025a), or with adjoint-based methods if the reward is differentiable (Domingo-Enrich et al., 2025); details are deferred to Section F.
3
Motivation: Understanding the Limitations of Flow and Score Matching
To motivate our method, we first ask why FSM losses might fail to capture properties of the data distribution that RL is able to recover. With infinite capacity and perfect optimization, the minimizer of either matching loss should reproduce the data distribution q exactly. Therefore, a natural starting point is to consider 3
how these losses behave in the setting where optimization is only approximate. We analyze the ODE-based formulation for simplicity. Let r : X → [0, 1] quantify a property of interest in the data—say whether an object in an image is well formed. A model that preserves this property adequately satisfies Epv [r] ≈ Eq [r]. Writing LFSM for either matching objective and v ∗ for its minimizer (with overloaded notation for the score), we ask: when does LFSM (v) − LFSM (v ∗ ) ≤ ε imply Epv [r] − Eq [r] → 0 as ε → 0, and at what rate? Our first result shows that, in the worst case, no such rate exists: low FSM loss can be arbitrarily uninformative about the reward gap. Proposition 3.1 (No reward certificate from FSM). Fix a (sufficiently regular) q and a bounded reward r ∈ [0, 1] with r(x) ≤ η on some region B ⊂ X . For every ε > 0 and every δ > 0 there is a velocity field v and a score field s, whose training error under qt is at most ε uniformly in t, but whose probability-flow ODE endpoint law p satisfies Ep [r] ≤ η + δ, regardless of Eq [r]. A precise statement and proof is given in Proposition D.3.
This counterexample is, however, adversarial. Under regularity assumptions on q and v ∗ a quantitative reward guarantee does exist — though, as we will see, a loose one. Proposition 3.2 (Reward certificate under uniform velocity control). Let r : X → [0, 1] be Lr -Lipschitz in that |r(x) − r(y)| ≤ Lr ∥x − y∥ for all x, y, and let v ∗ (·, t) be Lv -Lipschitz in x uniformly in t ∈ [0, 1]. If supt,x ∥v(x, t) − v ∗ (x, t)∥ ≤ ε, then the endpoint laws p and q of the probability-flow ODE satisfy Ep [r] − Eq [r] ≤ εLr ((eLv − 1)/Lv ) . (5) This dependence on Lr , Lv , and ε is essentially tight, and an on-policy variant under the rollout marginals pt holds as well; see Section D.2.1 for a formal statement and proof.
1.0
0.5
x
The construction exploits a train–test mismatch. An early error pushes the trajectory into regions of low qt -mass, so subsequent errors occur at states unseen during training and compound along the rollout (Figure 2). The same pathology motivates DAgger in behavioral cloning (Ross et al., 2011), with qt and pt playing the roles of expert and learner.
0.0
qt band true paths rollout
−0.5 −1.0
0
t
1
Figure 2 Distribution shift. Rollout paths (or-
ange) drift off the noising band qt (gray) and miss the high-reward region at t=1. The further from qt , the larger the error of true field (gray) vs. learned field (orange). Construction in Proposition D.3.
The bound factorizes into three terms: the velocity error ε, which FSM directly minimizes; the factor (eLv −1)/Lv , which controls how errors compound as the ODE is integrated from t=0 to 1; and the reward Lipschitz constant Lr , which measures how sharply r responds to small ℓ2 changes in the sample.
reward loss large reward change small loss change
FSM loss
v∗
vbad
Figure 3 Geometry mismatch. A small step in velocity space can produce large reward changes.
The last term exposes a second, more practically relevant obstruction: FSM controls errors in velocity space, while the reward is a function of samples, and when the two geometries are misaligned small velocity errors can map to large reward errors (Figure 3). The constant Lr quantifies this: rearranged, the bound says that to achieve a reward gap of size δ we need ε ≈ δ/Lr , so the sharper the reward, the smaller the velocity error must be. For rewards of practical interest, this is unfavorable. For example, in pixel space, whether a hand or animal face looks well-formed can flip under a few edge pixels, making Lr effectively very large and the required ε correspondingly small. Similar phenomena will occur in any representation space not directly aligned with the reward. For matching losses, this is doubly unfavorable as the gradient signal-to-noise ratio degrades precisely in the small-ε regime 4
the bound demands. With ξ := Y − v ∗ (Xt , t) and E[ξ | Xt , t] = 0, the per-sample gradient of the CFM loss splits as gθ = 2Jθ⊤ (vθ − v ∗ ) − 2Jθ⊤ ξ, Jθ := ∇θ vθ (Xt , t).
The first term is the optimization signal and vanishes as vθ → v ∗ ; the second is irreducible noise from the conditional variance of the regression target and does not. The achievable ε can therefore plateau well above what Eq. (5) requires. A similar decomposition can be derived for score matching.
RL’s edge. These obstructions suggest how RL helps. First, it evaluates the model on the samples it generates,
so there is no train–test mismatch. Second, and more importantly, with a reward in hand, we can use the reward landscape directly, focusing optimization on the directions of reward improvement and bypassing the geometry of velocity or score space. Moreover, while generally loose in practice, when the reward is such that the KL-regularized optimum is the desired target q and the noise schedule is fixed, the RL objective controls the property gap directly: a Pinsker-style bound (Section D.3.1) gives r λ |Ep [r] − Eq [r]| ≤ (LRL (p) − LRL (q)). 2
This verifies that we optimize the right quantity in the idealized setting. The implication is that even with perfect data, RL post-training can be beneficial when a good reward is available: the reward provides an optimization signal directly aligned with the property of interest, whereas FSM must recover that property indirectly from the regression target alone.
−0.25
ImageReward
To show that these distinctions matter in practice, we compare RL training against supervised distillation from an RL teacher. Starting from Stable Diffusion 1.5 (Rombach et al., 2022), we first train an RL teacher on ImageReward using adjoint matching (Domingo-Enrich et al., 2025), and then use its samples to fine-tune another Stable Diffusion student using score matching. If score matching could recover the reward-relevant properties learned by RL, the student should match the teacher. In practice it does not: the student stalls well below the teacher’s reward and fails to close the gap (Figure 4). See Section I.4 for the full setup.
RL–FSM gap
−0.50 −0.75 −1.00
FSM from pRL pRL (teacher)
−1.25 0
4
Method: Discriminator-Guided RL
2k
4k
6k
8k
Optimization steps
Figure 4 Distillation gap on Stable Diffusion 1.5: Unfortunately, RL is only beneficial if we have a reward that a score-matching student trained on samples captures the aspects of q that FSM may have failed to learn. from an RL teacher plateaus well below the Finding such a reward is non-trivial. The standard approach, teacher’s reward. learning a reward from human preferences, is only incidentally aligned with this goal. While preference rewards can correlate with q-properties (e.g., annotators prefer well-formed faces), they can also drag optimization along orthogonal axes like aesthetic appeal, so optimizing them is not necessarily desirable. Moreover, preference data is expensive to collect, and hard to define for many domains beyond images and video. We would instead like a reward derived directly from q. To this end, we introduce Discriminator-Guided RL (DRL).
From samples to rewards. The starting point of DRL is the observation that, in KL-regularized RL, there
is a reward whose optimum is exactly q. Recall from Section 2 that, at λ = 1, the unique maximizer of the KL-regularized objective is p∗ (x) ∝ exp(r(x)) pbase (x). Forcing p∗ = q identifies the ideal reward, up to additive constants, as the log density ratio between target and reference, r∗ (x) = log
q(x) . pbase (x)
(6)
This ratio is intractable in closed form, but it can be estimated from samples. Training a discriminator D : X → (0, 1) to separate q from pbase under the logistic loss has a well-known optimum, D∗ (x) = q(x)/(q(x)+pbase (x)). 5
Hence, its logit recovers r∗ . This motivates the reward estimator r̂(x) := log
D(x) 1 − D(x)
(7)
Algorithm 1 Discriminator-Guided RL (DRL) Stage 1: Reward Estimation Require: Samples x ∼ q, encoder ϕ, base model pbase 1: repeat
2: Sample xreal ∼ q, xfake ∼ pbase 3: ℓdisc (ψ) ← − log Dψ (ϕ(xreal )) − log 1−Dψ (ϕ(xfake )) 4: Update ψ by descending ℓdisc (ψ) 5: until convergence 6: Define r̂(x) ← logit Dψ (ϕ(x))
In practice, we parametrize the logit r̂ directly and use a sigmoid to define the discriminator. Using representations. Estimating r ∗ directly in
the flow output space is, however, both statistically hard and semantically unreliable. For example, a discriminator might separate q from pbase using artifacts that have no bearing on the relevant properties, achieving low classification error without producing a useful estimate of the density ratio (see design choices). We therefore constrain the reward to a pretrained representation space. Given a frozen encoder ϕ : X → Z, we set r̂(x) = h(ϕ(x)) for a learned head h. This both reduces the dimensionality of the estimation problem and restricts the discriminator to ϕ-visible structure.
Stage 2: KL-Reg. RL via Adjoint Matching
Require: Reward r̂, base velocity vbase , KL weight λ Require: Schedule σ(t), grid 0=t0 < · · · <tK =1
1: Initialize vθ ← vbase 2: repeat 3: Sample {Xk } from dXt = [2vθ − 1t Xt ] dt + σ(t) dBt 4: ã1 ← −λ ∇x r̂(XK ); 5: Solve backward 1 6: ã˙ t = −ã⊤ t ∇x [2vbase (Xt , t) − t Xt ] 7: Set LAM (θ) ← 2 2 1 PK 8: k=0 σ(tk ) vθ −vbase + σ(tk )ãk 2 9: Update θ by descending LAM (θ) 10: until convergence 11: return vθ
Mathematically, this restriction means that DRL will generally never be able to target q exactly. Nevertheless, its new target has a simple and intuitive characterization. It is the solution to the following constrained optimization problem: minp KL(p ∥ pbase ) subject to pϕ = q ϕ , (8)
where pϕ , q ϕ denote the pushforwards of p, q under ϕ (see Proposition E.1 for a formal statement and proof). In words, DRL makes the smallest KL change to pbase that aligns its representation-space distribution with that of the target, while leaving any variation invisible to ϕ unchanged; choosing ϕ thus chooses which aspects of q DRL is allowed to correct. Finally, when the learned reward is imperfect, we recover a feature-space test-function bound that augments the standard KL-regularized RL suboptimality bound with an additional term capturing the expected discrepancy between the ideal target in Equation (8) and the target implied by the learned reward. We defer the precise statement and proof to Section E.1.2. The DRL pipeline. These two ideas combine into the DRL algorithm specialized to flow models summarized in Algorithm 1. Stage 1 trains a discriminator in a frozen representation space to distinguish q from pbase and defines the reward as its logit. Stage 2 fine-tunes the base model with KL-regularized RL under that reward. To exploit reward gradients, we instantiate the RL stage with adjoint matching (Domingo-Enrich et al., 2025), a state-of-the-art RL algorithm for flow models. Further implementation details for both stages are deferred to Section H. A short description of adjoint matching is given in Section F.
5
Experiments
We validate DRL by using it to fine-tune four ImageNet-pretrained flow models: SiT (Ma et al., 2024), JiT (Li and He, 2025), REPA (Yu et al., 2024), and RAE (Zheng et al., 2025), covering latent- and pixel-space architectures. Notably, RAE and REPA already use SSL representations during pretraining as the latent space and regularization respectively. Unless noted otherwise, the discriminator is a class-conditional linear projection head (Miyato and Koyama, 2018) on frozen DINOv2-Large (Oquab et al., 2024) features, trained for 10k steps. The RL stage is 3k steps of adjoint matching (Domingo-Enrich et al., 2025). Importantly, this is only a small fraction of the 1M+ steps typically used to pre-train these models. Full description of every experimental setup, evaluation protocol, and hyperparameter is in Section I. Finally, we run adjoint matching with a new local-linear integrator we introduce for the memoryless SDE required for RL; we found it essential for stable training and view it as an important contribution of independent interest but leave its full description in App. Section G. 6
Base
λ=1
λ*
SiT REPA JiT RAE SiT REPA JiT RAE SiT REPA JiT RAE SiT REPA JiT RAE
Base λ = 1
λ*
242 159 148 37.5 88.2 63.9 33.9 6.66 42.6 31.6 39.4 10.2 9.38 6.48 7.16 1.31
54.4 40.3 42.8 20.6 19.3 15.6 11.3 4.49 16.9 12.8 21.0 9.13 2.62 2.14 2.73 1.29
95.5 58.3 97.3 30.2 43.8 29.4 23.7 5.71 20.2 14.3 29.7 9.54 2.62 2.14 3.72 1.38
0.25 0.5 0.75 1 1.25 1.5 FD / Base
Best CFG
Base λ = 1
λ*
47.0 36.9 41.5 25.9 21.6 18.6 10.9 4.83 11.8 9.52 18.2 9.50 1.50 1.21 1.91 1.28
42.4 32.9 30.6 20.2 16.4 13.0 8.10 4.26 11.6 9.32 16.5 9.06 1.78 1.50 1.87 1.19
42.4 32.9 36.1 23.7 18.3 14.6 10.2 4.57 12.2 10.0 17.1 9.29 2.43 1.84 1.87 1.25
0.25 0.5 0.75 1 1.25 1.5 FD / Base
Figure 5 Distribution alignment. Each row shows DRL’s FD normalized by the corresponding Base FD (vertical line);
lower is better. Squares show λ=1 and filled circles show the value optimized over λ. Raw FD values in tables: indicates improvement, indicates degradation over Base. Most markers fall left of the Base line, showing broad alignment gains, with the most significant gains in the no-CFG setting.
DRL Reduces the Distributional Gap Along Semantically Meaningful Directions. We first investigate whether DRL reduces the distributional gap to the data along semantically meaningful directions by computing Fréchet Distance (FD) (Heusel et al., 2017) in four feature spaces: DINOv2 (Oquab et al., 2024), DINOv3 (Siméoni et al., 2025), SigLIP (Zhai et al., 2023), and InceptionV3 (Szegedy et al., 2016). We consider two settings: the theoretically motivated unit λ and a tuned λ⋆ value selected over 1, 5, 10, 20, 40. We consider values larger than 1 because they often improve RL optimization in practice. To stabilize training at larger λ, we found it helpful to add an R1 penalty (Mescheder et al., 2018), which penalizes the discriminator’s input gradient on real samples and smooths the reward landscape (see design choices for more details). Figure 5 plots FD ratios (DRL/Base) for each (model, feature) pair, alongside their raw FD values with and without tuned CFG. Without CFG, DRL substantially improves distribution alignment: the tuned λ⋆ improves all 16 model–feature pairs, and λ = 1 improves all but one. The gains are especially large in DINOv2 and DINOv3, where FD often drops by more than half. With the best CFG setting, the Base models are already much closer to the data, but DRL still further improves alignment in most cases. In particular, λ⋆ improves FD in 14/16 pairs, with consistent gains in DINOv2, DINOv3, and SigLIP. The only clear exceptions are InceptionV3 for SiT and REPA under CFG, where FD increases slightly relative to Base. Across both settings, tuning λ consistently improves over the unit-weight setting, indicating that pushing past the theoretical value pays off in practice. To further investigate the role of λ, Figure 6 plots FD against λ in DINOv2 space alongside the corresponding precision–coverage trajectory. FD bottoms out around λ ∈ [5, 10] and then degrades; coverage drops at λ=20, and by λ=40 (not pictured; see Figures 17 and 18) the trajectory has collapsed off the frontier—the model overfits to the discriminator’s reward signal. We adopt λ=10 as the default for all downstream experiments, since it sits at the knee of this trade-off. Additional metrics and λ ablations for each model– feature pair are in Section J.1. Importantly, KD and FDval follow the same trends despite not being used to tune λ.
RAE
SiT
JiT
REPA
FD vs. λ 1
.75 .5
base
λ=10
Precision vs. Coverage Coverage
FD / base FD ↓
Incep.
SigLIP DINOv3 DINOv2
No CFG
0.8
0.6
.25 1
5 10 20 40
0.6
λ
Precision
0.8
Figure 6 FD vs. λ and Precision–Coverage in DINOv2-L. FD
decreases through λ=5−10 and degrades at high λ as coverage drops.
7
Improvement / σbase
ImageReward ↑ 0.5
0.0
HPSv2 ↑
1.0
Aesthetic ↑ 0.4
0.5
REPA SiT
JiT RAE
0.0
1.0 0.5
0.2
REPA SiT
No CFG
0.0
JiT RAE
Best CFG
REPA
PickScore ↑
REPA SiT
SiT
JiT RAE
JiT
0.0
REPA SiT
JiT RAE
RAE
Figure 7 Held-out preference reward gains. Normalized improvement (rDRL − rbase )/σbase at λ=10 on four rewards, with
and without CFG, where σbase is the base-model reward standard deviation. DRL improves every reward without ever seeing preference data.
DRL Improves Image Quality Without Preference Data. We next examine whether these distributional gains translate into perceptually better images. As proxies for image quality, we use four held-out preference reward models trained on human comparisons: ImageReward (Xu et al., 2023), PickScore (Kirstain et al., 2023), Aesthetics v2.5 (discus0434, 2024), and HPSv2 (Wu et al., 2023). Figure 7 plots the per-model improvement on each reward, (rDRL − rbase )/σbase , both without CFG and with best CFG. Without preference data or access to these rewards during training, DRL improves every reward on every architecture, with the largest gains for SiT and REPA and smaller but consistent gains for JiT and RAE. These improvements are also evident visually. The accompanying samples show fixed-seed outputs at λ values of 1, 10, and 20. As λ grows, shapes sharpen and global structure becomes more coherent, e.g., the spider’s shape, the koala’s face, and the car’s front end all become noticeably more defined. We provide more samples and λ values in App. Figures 30 to 33. DRL Provides a Better Foundation for Preference Alignment. DRL improves distribution match-
DRL λ=1
λ=10
λ=20
JiT-H/16
REPA
SiT-XL/2
ref
RAE
ing, but aligning to genuine subjective preferences still requires fine-tuning with preference-based RL (PRL). We argue that DRL also makes PRL itself more effective. PRL from the base model is asked to do two things with one imperfect scalar reward: repair distributional errors left by the generative model and optimize subjective preference. The result is a reward–drift trade-off: small λPRL leaves structural failures uncorrected; large λPRL exploits the proxy along nuisance directions like oversaturation or excessive brightness (Domingo-Enrich et al., 2025). By handling part of the data repair before preference optimization begins, DRL should ease this trade-off and let PRL focus on genuinely subjective improvements.
Base
We test this by running KL-regularized PRL from either the base model or the DRL checkpoint. We train with ImageReward (Xu et al., 2023) at λPRL ∈ {1, 10, 40} and evaluate on the held-out Figure 8 Effect of λ. Same noise and class label at HPSv2 reward (Wu et al., 2023), tracking five lowλ ∈ {1, 10, 20}. Larger λ produces sharper, more coherlevel statistics that commonly drift under aggresent samples while preserving content and composition. sive PRL: brightness, saturation, contrast, colorfulness, and whiteness; see Section I.3 for details on how these statistics are computed. Figure 9 plots HPSv2 against each statistic. DRL+PRL shifts the reward–drift frontier upward across all five axes, achieving higher HPSv2 than Base+PRL at comparable drift. The pattern holds with other forms of guidance, and on other rewards, including ImageReward itself (Section J.4).
8
DRL+PRL
Base+PRL
Base
ImageNet
SiT
RAE 1
0.23
DRL+PRL
SiT
JiT
λPRL
0.25
0.20
0.25
REPA
REPA
10
0.23
40
JiT
0.20
λPRL
0.22
1
Base+PRL
RAE
0.20
0.24
10
0.22
40 0.5
0.6
Bright.
0.2
0.3
Sat.
0.18
0.21
Contr.
0.15
Color
0.17 0.4
0.5
Whiteness
Figure 9 DRL provides a better starting point for PRL. Left: HPSv2 reward against five low-level statistics under PRL
from the base model (dashed, hollow) or the DRL checkpoint (solid, filled); green dashed lines mark ImageNet reference statistics. Right: matched fixed-seed samples at λPRL ∈ {1, 10, 40}. Across statistics, DRL+PRL gives better reward–drift Pareto fronts than Base+PRL and less brightness/whiteness drift as PRL strength increases. See Section I.3 for details on how the statistics are computed.
Samples from the model further support this conclusion. In Figure 9, Base+PRL at low λPRL leaves structural errors unresolved—the JiT dog has a malformed face, the REPA scorpion a shell-like body, the SiT bus a distorted chassis. Larger λPRL corrects some of these, especially for JiT and REPA, but at the cost of noticeably brighter and whiter images. DRL+PRL avoids the trade-off: samples remain structurally coherent and naturally colored across all λPRL . Per-model grids and guided variants appear in Figures 34 to 37. DRL Performance Cannot Be Reached with Flow Matching Alone. In Section 3 we argued that RL was needed because flow matching cannot faithfully recover many properties of the data distribution. To demonstrate that DRL genuinely learns aspects of the distribution that are hard to learn with flow matching alone, we use a distillation experiment similar to the one in Section 3. Concretely, we take the REPA model post-trained with DRL (λ = 1, R1 = 0) and use it as a teacher by fine-tuning REPA on samples from the teacher. If the features learned by DRL were learnable by the flow matching objective, we should see no gap between the teacher model and the student model. As Figure 10 demonstrates, this is not the case. Despite training for over 900k gradient steps, and seeing more than 50 million samples (over 150× what RL sees and over 40× the size of ImageNet), the student model is not able to faithfully replicate the teacher. This supports the argument from Section 3 that the RL objective matters, rather than only the samples produced by the teacher. We provide further details of the setup of this experiment in Section I.5. Design Choices. We close by ablating DRL’s two main design choices—the discriminator feature space and
the discriminator architecture—on REPA. Full setup is in Section I.
Discriminator design. We compare four discriminators—a linear head and an MLP-2 head on frozen DINOv2 features, a fully fine-tuned DINOv2, and a DINOv2 architecture trained from scratch—at λ=1 with R1 =0, and at λ=10 with R1 ∈ {0, 10−5 , 10−3 , 10−1 }. Figure 11 plots DINOv2-L FD across this sweep, with representative samples along the top row. The full R1 sweep at λ=1 is in Section J.5. 9
pDRL (teacher)
Inception
Base model
SFT from pDRL
DINOv2
DINOv3
SigLIP
160
FD (↓)
6
60
30
50
25
140
5
120
4
100 gap
3
gap
gap
40
20 15
30 0
500k Optimization Steps
gap
80
0
500k Optimization Steps
0
500k Optimization Steps
0
500k Optimization Steps
Figure 10 DRL performance cannot be reached with flow matching alone. FD (↓) over optimization steps for a student
trained with flow matching on samples from a DRL teacher (λ=1, no R1 ), evaluated in four feature spaces. The student improves rapidly over the base model (grey, dotted) but plateaus well above the teacher (amber, dashed) across all four feature spaces.
Pretrained features are essential: the from-scratch discriminator underperforms despite reaching 95% validation accuracy, and even fine-tuning DINOv2 trails frozen features in the theoretically motivated setting (λ=1, R1 =0). At λ=10 without R1 , the model collapses to high-reward regions; a small R1 not only stabilizes training but improves over (λ=1, R1 =0) in both FD and visual quality. Curiously, while R1 helps at λ=10, it degrades performance at λ=1 (Figure 23), suggesting its role here differs from the vanishing-gradient role it plays in standard GANs (Mescheder et al., 2018). Finally, since the MLP and linear heads perform similarly, we prefer the linear head for simplicity.
6
Base
Linear
We highlight three primary threads of related work below and delay a more detailed discussion to App. Section A. RL vs. Imitation Learning. A long line of research has
λ = 10 R1 = 0
λ = 10 R1 = 10−5
MLP-2
Finetune
Scratch
103
base 10
2
λ=1 R1 = 0
Related Work
λ=1 R1 = 0
FD
Feature space. We re-train the discriminator on six other frozen embedders—DINOv2-B, DINOv3 (B/L), SigLIP (B/L), and InceptionV3—at λ=1, R1 =0, and evaluate FD in every embedder’s space. Table 1 reports a representative subset; the full 7×7 table is in Table 5 with KD as an additional metric. All six SSL embedders improve both FD and KD over the base on every evaluation space, with small spread between them; DINOv2-L is consistently strongest, confirming our default. The one exception is InceptionV3, which barely improves over the base on any SSL evaluation space and is outperformed on its own evaluation space by every SSL embedder—we suspect its classification-only training fails to encode perceptual structure as linear directions. Larger λ and nonzero R1 narrow the gap (FID 3.4, FD 107 on DINOv2-B) but do not close it.
λ = 10 R1 = 0
λ = 10 λ = 10 R1 = 10−5 R1 = 10−3
λ = 10 R1 = 0.1
Figure 11 Discriminator ablation. Frozen feature heads work well at λ=1 without R1 . At λ=10, removing R1 leads to collapse while a small R1 restores performance. The "from-scratch" discriminator is the worst.
studied the limitations of supervised imitation learning and the benefits of on-policy methods (Ross and Bagnell, 2010; Ross et al., 2011). We extend this line to flow and score matching, showing that they may suffer from similar pathologies despite not being demonstrationbased, with the continuous-time setting yielding even weaker guarantees than standard DAgger-style bounds. 10
Inverse RL. Our remedy mirrors the structure of inverse RL (Abbeel and Ng, 2004; Ziebart et al., 2008), which
recovers a reward from expert demonstrations and trains a policy against it. Generative adversarial imitation learning (GAIL) (Ho and Ermon, 2016) and adversarial inverse reinforcement learning (AIRL) (Fu et al., 2018) apply this idea to standard RL control; we adapt it to improve generative models. Moreover, while these prior works are largely motivated by closing the train–test gap, our observation that the method only helps when paired with an embedder suggests that its principal benefit may be addressing geometric obstructions, not the mismatch itself. We provide further discussion of this point in Section B. Calibration. Our theory also asks a calibration question: for a property r, does small flow- or score-matching
loss imply Ep [r] ≈ Eq [r]? Related work enforces distribution-level constraints directly in controlled generation and fairness-oriented fine-tuning (Khalifa et al., 2021; Shen et al., 2023). Closest to us is Smith et al. (2025), who calibrate generators by finding the KL-closest model satisfying user-specified moment constraints. DRL solves a related KL-constrained problem: it seeks a distribution close to the original model whose pushforward through a feature map matches the data pushforward. Unlike Smith et al. (2025), we focus on aligning the full distribution in high-dimensional SSL feature spaces, and solve the implicit KL-constrained problem by learning a discriminator between real and generated samples, whose logit is then used for RL.
7
Conclusions
In this paper, we have argued that for flow-based models, RL is valuable not only as a method for steering toward externally specified rewards, but also as a means of evaluating the model on its own samples and exploiting reward geometry rather than ℓ2 distance. DRL turns this view into a method, and with its discriminator-based reward, consistently improves distributional alignment, preference metrics, and performance–drift tradeoffs under preference-based RL.
Table 1 Feature space ablation (REPA).
λ=1, R1 =0, FD (↓). Each row is an evaluation feature space; each column is the feature space used to train the discriminator. Best DRL per row in bold. SSL embedders are L variants. Full FD/KD table in Table 5. Training Base DINOv2 DINOv3 SigLIP Incep.
Eval
Several directions follow: 1) The reliance on frozen SSL Incep. 6.43 2.14 3.02 2.34 5.68 features invites the question of whether comparable repreDINOv2 159 58.3 63.2 87.1 151 sentations can be learned jointly with the model. 2) Our DINOv3 63.7 29.4 18.5 38.8 62.0 14.3 16.5 16.1 29.9 SigLIP 31.2 bounds are worst-case, and quantifying a priori which properties FSM struggles with, as well as how these obstructions translate to practical SDE/ODE samplers, would clarify when an RL stage is helpful; we discuss these limitations further in Section B. 3) More broadly, RL is one way to exploit on-policy evaluation and non-ℓ2 geometry; other sample-based losses like MMD may offer different trade-offs. Together our results point to a view of generative training in which matching objectives and RL are not two methods with different goals, but rather two complementary tools for getting the most out of our data.
Acknowledgments We thank Brian Trippe, Brian Karrer, Ricky T. Q. Chen, and Sebastian Salazar for helpful comments on the manuscript and for insightful discussions about the ideas presented here.
11
References Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In Proceedings of the twenty-first international conference on Machine learning, 2004. Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist, and Olivier Bachem. On-policy distillation of language models: Learning from self-generated mistakes. In International Conference on Learning Representations (ICLR), 2024. Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In Proceedings of the International Conference on Learning Representations (ICLR), 2023. V. I. Arnold. Ordinary Differential Equations. MIT Press, Cambridge, MA, 1978. Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. In The Twelfth International Conference on Learning Representations, 2024. https://openreview.net/ forum?id=YCWjhGrJFD. Nicholas M Boffi, Michael S Albergo, and Eric Vanden-Eijnden. Flow map matching with stochastic interpolants: A mathematical framework for consistency models. TMLR, 2025. Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Fei-Fei Li. ImageNet: A large-scale hierarchical image database. In IEEE conference on computer vision and pattern recognition (CVPR), 2009. discus0434. Aesthetic predictor V2.5: SigLIP-based aesthetic score predictor. https://github.com/discus0434/ aesthetic-predictor-v2-5, 2024. GitHub repository. Carles Domingo-Enrich, Michal Drozdzal, Brian Karrer, and Ricky TQ Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. In International Conference on Learning Representations (ICLR), 2025. Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In Proceedings of the International Conference on Learning Representations (ICLR), 2021. Bradley Efron. Tweedie’s formula and selection bias. Journal of the American Statistical Association, 2011. Jiajun Fan, Shuaike Shen, Chaoran Cheng, Yuxin Chen, Chumeng Liang, and Ge Liu. Online reward-weighted fine-tuning of flow matching with Wasserstein regularization. In The Thirteenth International Conference on Learning Representations, 2025. Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. DPOK: Reinforcement learning for fine-tuning text-to-image diffusion models. Advances in Neural Information Processing Systems, 36:79858–79885, 2023. Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse reinforcement learning. In International Conference on Learning Representations (ICLR), 2018. Igor Vladimirovich Girsanov. On transforming a certain class of stochastic processes by absolutely continuous substitution of measures. Theory of Probability & Its Applications, 1960. Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in Neural Information Processing Systems, 2014. Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by a two time-scale update rule converge to a local Nash equilibrium. In Advances in Neural Information Processing Systems (NeurIPS), 2017. Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in Neural Information Processing Systems, 2016. Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 2020. Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 2005.
12
Jonas Hübotter, Frederike Lübeck, Lejs Behric, Anton Baumann, Marco Bagatella, Daniel Marta, Ido Hakimi, Idan Shenfeld, Thomas Kleine Buening, Carlos Guestrin, and Andreas Krause. Reinforcement learning via self-distillation, 2026. https://arxiv.org/abs/2601.20802. Alexia Jolicoeur-Martineau, Rémi Piché-Taillefer, Rémi Tachet des Combes, and Ioannis Mitliagkas. Adversarial score matching and improved sampling for image generation. arXiv preprint arXiv:2009.05475, 2020. Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems, 2022. Tero Karras, Miika Aittala, Tuomas Kynkäänniemi, Jaakko Lehtinen, Timo Aila, and Samuli Laine. Guiding a diffusion model with a bad version of itself. Advances in Neural Information Processing Systems, 2024. Muhammad Khalifa, Hady Elsahar, and Marc Dymetman. A distributional approach to controlled text generation. In The International Conference on Learning Representations (ICLR), 2021. Yuval Kirstain, Adam Polyak, Uriel Singer, Shahbuland Matiana, Joe Penna, and Omer Levy. Pick-a-Pic: An open dataset of user preferences for text-to-image generation. In Advances in Neural Information Processing Systems (NeurIPS), 2023. Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. In Advances in Neural Information Processing Systems (NeurIPS), 2019. Tuomas Kynkäänniemi, Miika Aittala, Tero Karras, Samuli Laine, Timo Aila, and Jaakko Lehtinen. Applying guidance in a limited interval improves sample and distribution quality in diffusion models. Advances in Neural Information Processing Systems, 37:122458–122483, 2024. Michael Laskey, Jonathan Lee, Roy Fox, Anca Dragan, and Ken Goldberg. DART: Noise injection for robust imitation learning. In Conference on robot learning. PMLR, 2017. Mingxiao Li, Tingyu Qu, Ruicong Yao, Wei Sun, and Marie-Francine Moens. Alleviating exposure bias in diffusion models through sampling with shifted time steps. arXiv preprint arXiv:2305.15583, 2023. Tianhong Li and Kaiming He. Back to basics: Let denoising generative models denoise. arXiv preprint arXiv:2511.13720, 2025. Yangming Li and Mihaela van der Schaar. On error propagation of diffusion models. arXiv preprint arXiv:2308.05021, 2023. Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022. Jie Liu, Gongye Liu, Jiajun Liang, Yangguang Li, Jiaheng Liu, Xintao Wang, Pengfei Wan, Di Zhang, and Wanli Ouyang. Flow-GRPO: Training flow matching models via online RL, 2025a. https://arxiv.org/abs/2505.05470. Jie Liu, Gongye Liu, Jiajun Liang, Ziyang Yuan, Xiaokun Liu, Mingwu Zheng, Xiele Wu, Qiulin Wang, Menghan Xia, Xintao Wang, et al. Improving video generation with human feedback. arXiv preprint arXiv:2501.13918, 2025b. Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003, 2022. Kevin Lu and Thinking Machines Lab. On-policy distillation. Thinking Machines Lab: Connectionism, 2025. doi: 10.64434/tml.20251026. https://thinkingmachines.ai/blog/on-policy-distillation. Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden-Eijnden, and Saining Xie. SiT: Exploring flow and diffusion-based generative models with scalable interpolant transformers. In European Conference on Computer Vision, pages 23–40. Springer, 2024. Lars Mescheder, Andreas Geiger, and Sebastian Nowozin. Which training methods for GANs do actually converge? In International conference on machine learning, pages 3481–3490. PMLR, 2018. Takeru Miyato and Masanori Koyama. cGANs with projection discriminator, 2018. https://arxiv.org/abs/1802.05637. Muhammad Ferjad Naeem, Seong Joon Oh, Youngjung Uh, Yunjey Choi, and Jaejun Yoo. Reliable fidelity and diversity metrics for generative models. In International Conference on Machine Learning (ICML), 2020. Andrew Y Ng and Stuart Russell. Algorithms for inverse reinforcement learning. In International Conference on Machine Learning (ICML), pages 663–670, 2000.
13
Mang Ning, Mingxiao Li, Jianlin Su, Albert Ali Salah, and Itir Onal Ertugrul. Elucidating the exposure bias in diffusion models. arXiv preprint arXiv:2308.15321, 2023a. Mang Ning, Enver Sangineto, Angelo Porrello, Simone Calderara, and Rita Cucchiara. Input perturbation reduces exposure bias in diffusion models. arXiv preprint arXiv:2301.11706, 2023b. Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, Mahmoud Assran, Nicolas Ballas, Wojciech Galuba, Russell Howes, Po-Yao Huang, Shang-Wen Li, Ishan Misra, Michael Rabbat, Vasu Sharma, Gabriel Synnaeve, Hu Xu, Hervé Jegou, Julien Mairal, Patrick Labatut, Armand Joulin, and Piotr Bojanowski. DINOv2: Learning robust visual features without supervision. Transactions on Machine Learning Research (TMLR), 2024. William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pages 4195–4205, 2023. Jan Peters and Stefan Schaal. Reinforcement learning by reward-weighted regression for operational space control. In Proceedings of the 24th international conference on Machine learning, pages 745–750, 2007. Peter Potaptchik, Cheuk-Kit Lee, and Michael S Albergo. Tilt matching for scalable sampling and fine-tuning. arXiv preprint arXiv:2512.21829, 2025. Nived Rajaraman, Lin Yang, Jiantao Jiao, and Kannan Ramchandran. Toward the fundamental limits of imitation learning. Advances in Neural Information Processing Systems, 2020. Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684–10695, 2022. Stéphane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 661–668. JMLR Workshop and Conference Proceedings, 2010. Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pages 627–635. JMLR Workshop and Conference Proceedings, 2011. Axel Sauer, Dominik Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffusion distillation. In European Conference on Computer Vision, pages 87–103. Springer, 2024. Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, Patrick Schramowski, Srivatsa Kundurthy, Katherine Crowson, Ludwig Schmidt, Robert Kaczmarczyk, and Jenia Jitsev. LAION-5B: An open large-scale dataset for training next generation image-text models. In Advances in Neural Information Processing Systems (NeurIPS), 2022. John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, et al. DeepSeekMath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300, 2024. Xudong Shen, Chao Du, Tianyu Pang, Min Lin, Yongkang Wong, and Mohan Kankanhalli. Finetuning text-to-image diffusion models for fairness. arXiv preprint arXiv:2311.07604, 2023. Idan Shenfeld, Mehul Damani, Jonas Hübotter, and Pulkit Agrawal. Self-distillation enables continual learning, 2026. https://arxiv.org/abs/2601.19897. Oriane Siméoni, Huy V Vo, Maximilian Seitzer, Federico Baldassarre, Maxime Oquab, Cijo Jose, Vasil Khalidov, Marc Szafraniec, Seungeun Yi, Michaël Ramamonjisoa, et al. DINOv3. arXiv preprint arXiv:2508.10104, 2025. Henry D Smith, Nathaniel L Diamant, and Brian L Trippe. Calibrating generative models to distributional constraints. arXiv preprint arXiv:2510.10020, 2025. Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020.
14
Gokul Swamy, Sanjiban Choudhury, J Andrew Bagnell, and Steven Wu. Of moments and matching: A game-theoretic framework for closing the imitation gap. In International Conference on Machine Learning, pages 10022–10032. PMLR, 2021. Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2818–2826, 2016. Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 2011. Bram Wallace, Meihua Dang, Rafael Rafailov, Linqi Zhou, Aaron Lou, Senthil Purushwalkam, Stefano Ermon, Caiming Xiong, Shafiq Joty, and Nikhil Naik. Diffusion model alignment using direct preference optimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2024. Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 1992. Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score v2: A solid benchmark for evaluating human preferences of text-to-image synthesis. arXiv preprint arXiv:2306.09341, 2023. Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion GANs. arXiv preprint arXiv:2112.07804, 2021. Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong. ImageReward: Learning and evaluating human preferences for text-to-image generation. In Advances in Neural Information Processing Systems (NeurIPS), 2023. Yanwu Xu, Yang Zhao, Zhisheng Xiao, and Tingbo Hou. UFOGen: You forward once large scale text-to-image generation via diffusion GANs. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2024. Haotian Ye, Kaiwen Zheng, Jiashu Xu, Puheng Li, Huayu Chen, Jiaqi Han, Sheng Liu, Qinsheng Zhang, Hanzi Mao, Zekun Hao, Prithvijit Chattopadhyay, Dinghao Yang, Liang Feng, Maosheng Liao, Junjie Bai, Ming-Yu Liu, James Zou, and Stefano Ermon. Data-regularized reinforcement learning for diffusion models at scale, 2025. https://arxiv.org/abs/2512.04332. Tianwei Yin, Michaël Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Fredo Durand, and William T Freeman. Improved distribution matching distillation for fast image synthesis. Advances in neural information processing systems, 2024. Sihyun Yu, Sangkyung Kwak, Huiwon Jang, Jongheon Jeong, Jonathan Huang, Jinwoo Shin, and Saining Xie. Representation alignment for generation: Training diffusion transformers is easier than you think. arXiv preprint arXiv:2410.06940, 2024. Xiaohua Zhai, Basil Mustafa, Alexander Kolesnikov, and Lucas Beyer. Sigmoid loss for language image pre-training. In Proceedings of the IEEE/CVF international conference on computer vision, pages 11975–11986, 2023. Lvmin Zhang, Anyi Rao, and Maneesh Agrawala. Adding conditional control to text-to-image diffusion models. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 3836–3847, 2023. Boyang Zheng, Nanye Ma, Shengbang Tong, and Saining Xie. Diffusion transformers with representation autoencoders. arXiv preprint arXiv:2510.11690, 2025. Brian D Ziebart. Modeling Purposeful Adaptive Behavior with the Principle of Maximum Causal Entropy. PhD thesis, Machine Learning Department, Carnegie Mellon University, Dec 2010. Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, and Anind K Dey. Maximum entropy inverse reinforcement learning. In Proc. AAAI, pages 1433–1438, 2008.
15
Appendix Appendix Contents A
Detailed Related Work. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 18
B
Limitations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 20
C
Proofs for Section 2: Preliminaries . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 22
D
C.1
Score-velocity relationship . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 22
C.2
SDE and ODE marginal equivalence . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 22
Proofs for Section 3: Understanding the Limitations of Flow and Score Matching . . . . . . . . . . . . . . . . . . . . . 24 D.1 Distribution Shift and Error Accumulation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 24 D.2 Geometry mismatch . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 34 D.3 How RL Helps . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 40
E
Proofs for Section 4: Method: Discriminator-Guided RL . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 42 E.1
Why Representation Spaces Help . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 42
F
Adjoint Matching: Background . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 45
G
Local Linear Integrator for the Memoryless SDE . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 47 G.1 The Memoryless SDE and the Local Linear Integrator . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 47 G.2 Derivation of the Local Linear Integrator . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 50
H
Additional Implementation Details. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 53
I
Experimental Setup. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 55
J
I.1
Distributional alignment setup . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 55
I.2
Better image quality (reward-transfer evaluation) setup . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 56
I.3
Preference-based RL setup. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 56
I.4
Distillation from RL teachers setup . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 57
I.5
Distillation from DRL Teachers . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 58
I.6
Feature-space ablation setup . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 58
I.7
Discriminator architecture and training ablation setup . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 58
Extended Results . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 60 J.1
Alignment: Quantitative λ Sweep . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 62
J.2
Alignment: Full Distribution Metrics (PRDC) . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 64
J.3
Image Quality: Reward Improvement vs. λ. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 66
J.4
Image Quality / Preference RL: Pareto Plots . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 67
J.5
Ablations: Discriminator Full R1 Sweep. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 70
J.6
Ablations: Full Feature-Space Sweep . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 70
J.7
Effect of CFG on Reward Scores and Fréchet Distance . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 71
16
K
Qualitative Samples. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 72 K.1 Base vs. Fine-tuned . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 73 K.2 Effect of DRL Strength . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 77 K.3 RL Fine-tuning Sample Images. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 81
17
Appendix A
Detailed Related Work
Covariate shift, imitation learning, and exposure bias. The distribution mismatch pathology we identify in flow
and score matching (Proposition 3.1) is analogous to the classical covariate-shift problem in imitation learning with behavioral cloning. Behavioral cloning controls prediction error under the expert’s state distribution, but the learned policy is deployed under its own induced state distribution. As a result, errors that are large under the model’s rollout distribution can appear small under the expert’s distribution (Ross and Bagnell, 2010; Ross et al., 2011). This phenomenon has been extensively studied in reinforcement learning and imitation learning (Laskey et al., 2017; Rajaraman et al., 2020; Swamy et al., 2021), and has recently received renewed attention in language modeling through on-policy distillation and related methods (Agarwal et al., 2024; Hübotter et al., 2026; Lu and Lab, 2025; Shenfeld et al., 2026). Independently, the diffusion literature has studied analogous train–test mismatch under the names exposure bias, sampling drift, and error propagation (Ning et al., 2023b; Li and van der Schaar, 2023; Ning et al., 2023a). Because this thread has evolved separately from the RL literature, existing approaches often address the mismatch through ad hoc, model-specific off-distribution corrections, such as perturbing training inputs to mimic sampling-time errors (Ning et al., 2023b), modifying the sampling trajectory via shifted time steps (Li et al., 2023), or modifying the norm of transitions in the sampling process (Ning et al., 2023a). The closest work to an on-policy correction is Li and van der Schaar (2023), who use an MMD regularizer to match forward noising marginals to short model-induced denoising marginals. However, their rollouts are warm-started from the data noising distribution and are very short in their main experiments: only L = 5 reverse steps out of T = 1000. Thus the method addresses local train–test mismatch rather than the full sampling-trajectory occupancy mismatch. We connect these two lines of work in two ways. First, we show that flow and score matching admit worst-case pathologies analogous to those in DAgger-style imitation learning, with even weaker guarantees due to the continuous-time nature of the dynamics. Second, motivated by this connection, we focus on an explicitly on-policy solution, closer in spirit to reinforcement learning and imitation learning than to off-policy perturbation-based corrections. Given our positive results, we believe an interesting direction for future work is to revisit sampling-based approaches such as Li and van der Schaar (2023) with the stronger tooling used in our paper, such as adjoint-based training and the use of semantic feature spaces in the endpoint distributions for the MMD computations. Inverse RL and adversarial imitation. A classical alternative to behavioral cloning is inverse reinforcement
learning (IRL), which infers a reward under which expert behavior is optimal and then optimizes that reward with RL (Ng and Russell, 2000; Abbeel and Ng, 2004). Maximum-entropy and maximum-causal-entropy IRL are especially close in form to our setting: they model expert trajectories as an exponential tilt of a reference trajectory measure by cumulative reward, pr (τ ) ∝ µenv (τ ) exp(Rr (τ )), where µenv is induced by the initial-state distribution and environment dynamics (Ziebart et al., 2008; Ziebart, 2010). KL-regularized posttraining has the analogous endpoint form pr (x) ∝ pbase (x) exp(λr(x)), so r∗ (x) = λ−1 (log q(x) − log pbase (x)) makes q the optimal tilted distribution, up to additive constants. Closely related IRL algorithms include adversarial imitation methods, which use discriminators in related but different ways than DRL. GAIL trains a discriminator between expert and current learner state-action samples, yielding an adversarial objective for matching occupancy measures rather than a fixed reusable reward (Ho and Ermon, 2016). AIRL keeps the adversarial imitation loop but constrains the discriminator logit to decompose into a reward term and a potential-based shaping term, with the goal of recovering a reward that can be re-optimized under changes in dynamics rather than just a policy that matches the expert in the training MDP (Fu et al., 2018). DRL adapts the discriminator-as-reward idea from GANs and adversarial imitation (Goodfellow et al., 2014) to generative post-training, but with a fixed endpoint reference: the discriminator is trained only once on (q, pbase ) through a representation ϕ, and its logit defines a fixed reward used by KL-regularized RL. This avoids the notoriously unstable alternating min-max optimization (Mescheder et al., 2018), and exposes a single hyperparameter λ controlling the strength of the correction. We also differ from prior work in our use of self-supervised representations for ϕ, without which, as we show, the procedure
18
is not practical. Adversarial training in diffusion models. Adversarial training has also been combined with diffusion models
directly, in two main ways. Early work on CIFAR-10 and LSUN Churches (Jolicoeur-Martineau et al., 2020) augmented the matching loss with an adversarial term on the model’s posterior-mean estimate E[X1 | Xt ] sharpening the one-step denoiser used at the final sampling step. More recent approaches (Xiao et al., 2021; Sauer et al., 2024; Xu et al., 2024; Yin et al., 2024) use GANs to essentially amortize the transition kernel p(xs | xt ) for s and t, exploiting the fact that GANs are good implicit models in order to enable few-step sampling and distillation.
In addition to the differences described in the previous section, DRL differs from these lines of work in two ways. First, our goal is to correct an existing flow, not to learn an implicit model or a few-step sampler. Second, while Jolicoeur-Martineau et al. (2020) share our motivation, training the discriminator on the posterior mean E[X1 | Xt ] is not principled: E[X1 | Xt ] is the MSE denoiser, not a sample from the model, so matching its distribution to data does not in general yield sharp, high-quality samples. DRL instead trains the discriminator on actual model samples and corrects the velocity field directly via adjoint matching. Reward-based post-training for diffusion and flow models. RL post-training of diffusion and flow models requires
two ingredients: a reward and an algorithm to optimize it. As discussed throughout the paper, DRL is a contribution to the first; here we discuss its relation to the second. A growing literature studies optimization algorithms for this setting (Black et al., 2024; Fan et al., 2023; Domingo-Enrich et al., 2025; Liu et al., 2025a). While we expect DRL to benefit directly from continued progress in this literature, our analysis in Section 3 predicts that not every optimizer is suited to the role: an effective Stage 2 optimizer needs to be on-policy and to propagate reward information through the model in a way that exploits its landscape—either through ∇x r, as in adjoint matching, or through score-function estimators that estimate ∇θ Epθ [r] directly.
Methods that recast reward optimization as a matching problem against a reward-tilted target, such as Tilt Matching (Potaptchik et al., 2025) and reward-weighted regression (Peters and Schaal, 2007; Fan et al., 2025; Black et al., 2024), take a different route: they use the reward to estimate a matching objective under the target distribution, rather than leveraging the reward landscape directly in the update. From the perspective of our analysis, this leaves them subject to the same matching geometry that we identify as a potential limitation. Some instantiations are additionally off-policy, which may further reintroduce a train–test mismatch. These considerations need not be decisive in every problem instance, but they suggest a plausible partial explanation for the underperformance of RWR reported in Black et al. (2024) and of Tilt Matching reported in Table 1 of Potaptchik et al. (2025), relative to methods that either propagate ∇x r or rely on policy-gradient-style estimators. Further empirical validation of this hypothesis is an interesting direction for future work.
Calibration. Our analysis of surrogate objectives is closely related to calibration: for a statistic r, we ask whether low flow- or score-matching loss is enough to control the calibration error |Epv [r] − Eq [r]|. Distributionlevel constraints of this form have been studied in controlled language generation (Khalifa et al., 2021) and in fairness-oriented text-to-image fine-tuning (Shen et al., 2023). The closest connection is to Smith et al. (2025). Their applications differ significantly from ours — e.g. they rebalance animal-class proportions in a conditional image model so that lions, tigers, and other wildlife categories appear evenly, and balance male and female character frequencies in language-model stories about different professions — but one of their methods, CGM-reward, is closely related to DRL. CGM-reward finds the KL-closest model satisfying moment constraints Epθ [h(x)] = h⋆ , with h typically a class indicator (e.g., the gender of the character in a generated story); the max-entropy dual yields the exponential tilt pα⋆ (x) ∝ pθbase (x) exp(α⋆⊤ h(x)). DRL targets a more general object, the full pushforward pϕ = q ϕ under KL regularization (Proposition E.1), but in the linear case the two coincide in form: our class-conditional discriminator ry (x) = wy⊤ ϕ(x) + by gives a tilt linear in ϕ. Nevertheless, even in this special case the two generally estimate different targets, since CGM-reward solves an empirical max-entropy dual against user-specified moments while DRL recovers the tilt from a discriminator’s logit. Moreover, as shown in design choices, the linear head is not essential: MLP discriminators yield similar gains.
19
B
Limitations
While our results provide evidence that DRL is an effective post-training method, several caveats remain. These include practical limitations inherited from RL-based fine-tuning, limitations of the representation space used to define the reward, and limitations of the simplified theoretical setting used to motivate the method. Classifier-free guidance. One limitation of DRL, and of flow-based RL methods more generally, is that they do
not provide a clean way to incorporate classifier-free guidance (CFG) during training. CFG typically disrupts the structure of the learned flow used to pass between ODE and SDE samplers. We experimented with CFG during training, but it often led to instability; similar difficulties have been reported by Domingo-Enrich et al. (2025).1 We also observed that the improvements of DRL over the base model are smaller when both are sampled with their best CFG scale, as shown in figure 5. We attribute these issues to the fact that current RL algorithms do not yet handle CFG cleanly, and consider developing principled CFG-compatible RL training an important direction for future work. Choice of λ. As shown in section 5, the KL regularization weight λ plays a significant role in performance, and
a small number of values typically need to be explored. While we found λ=10 to be a reliable default across our experiments, this may not hold in all settings, and the need to tune λ adds to the overall computational cost—a burden shared by RL methods more broadly. An interesting direction for future work is to make the method more robust to the choice of λ, and more generally to amortize the cost of tuning this hyperparameter, for example with a ControlNet-like mechanism (Zhang et al., 2023). Reverse KL behavior. As discussed in the main text, DRL optimizes a reverse KL objective and therefore
inherits both its advantages and its drawbacks. In particular, while DRL is robust at low λ values, at large λ the model can become mode-seeking, as visible in the PRDC curves (section J.2). We believe this partly explains why R1 gradient regularization proved beneficial at higher λ (design choices). Although we did not find this to be a practical issue at the λ values we recommend, it is worth keeping in mind. Additional compute. As a post-training method, DRL requires additional compute beyond base-model training (see section I for details). While this cost is only a small fraction of pretraining, it is nonetheless an extra expense that may not be justified in all scenarios. Role of the representation. Our experiments suggest that DRL is most effective when the discriminator is
trained in a pretrained representation space. This is a strength, because it lets the reward focus on semantic discrepancies, but it is also a limitation: the method can only correct distributional differences that are visible to the chosen embedder, so it requires a suitable representation in the first place. Moreover, although our analysis emphasizes both train–test mismatch and reward-geometry mismatch, the empirical importance of the embedder suggests that a large part of DRL’s benefit may come from changing the geometry in which the reward is estimated, rather than from on-policy optimization alone. Scope of the theory. Our theoretical analysis is intentionally carried out in a simplified probability-flow ODE
setting. We believe this captures the core geometric obstructions behind propositions 3.1 and 3.2, but it is not a complete model of every practical training setup. In particular, while the reward certificate in proposition 3.2 should have close SDE analogues under suitable drift control, the no-certificate construction in proposition 3.1 uses the ODE structure directly. It is therefore not clear how far that worst-case construction extends to stochastic samplers, especially because path-KL objectives can yield Pinsker-type control at the trajectory level. This is nevertheless consistent with our intent: the theory is meant to identify failure modes and motivate the method, not to claim that the worst case occurs generically. Furthermore, while a construction like the ODE one is possibly too extreme, as discussed in Section A many works have documented significant train–test gaps in practice for SDE samplers as well. A related gap is the distinction between training and inference dynamics. In our experiments, RL training uses the memoryless SDE, while evaluation uses the ODE sampler. The memoryless-schedule equivalence guarantees agreement at the optimum, but during training the SDE and ODE rollout distributions need not coincide exactly. While this means that SDE-level Pinsker bounds do not by themselves certify the ODE samples used at inference, we did not observe large discrepancies between the two in practice. Moreover, as argued in the paper, we hypothesize that the main driver of DRL’s improvements is correcting geometric 1 https://github.com/microsoft/soc-fine-tuning-sd
20
obstructions, which should affect both SDE and ODE samplers similarly through the reward gradient. This does not affect our explanation of the limitations of flow and score matching.
21
C
Proofs for Section 2: Preliminaries
This appendix collects proofs for the statements in section 2.
C.1
Score-velocity relationship
Proposition C.1 (Score-velocity equivalence). Let X1 ∼ q and X0 ∼ N (0, I) be independent, and define the interpolation Xt = α(t)X1 + β(t)X0 with marginal density qt . The velocity field vt (x) := E[α̇(t)X1 + β̇(t)X0 | Xt = x] and the score st (x) := ∇x log qt (x) are related by: ! α̇(t) β̇(t) α̇(t) vt (x) = x + β(t)2 − st (x). α(t) α(t) β(t) Proof. To compute the velocity, we need E[X1 | Xt = x] and E[X0 | Xt = x].
From the interpolation, we can express the noise as X0 = (Xt − α(t)X1 )/β(t). Therefore: E[X0 | Xt = x] =
x − α(t)E[X1 | Xt = x] . β(t)
By Tweedie’s identity (Efron, 2011), the posterior mean of X1 given Xt = x is: E[X1 | Xt = x] =
x β(t)2 x + β(t)2 st (x) = + st (x). α(t) α(t) α(t)
Substituting: E[X0 | Xt = x] =
x − α(t)
2
β(t) x α(t) + α(t) st (x)
β(t)
=
−β(t)2 st (x) = −β(t)st (x). β(t)
Now we compute the velocity: vt (x) = α̇(t)E[X1 | Xt = x] + β̇(t)E[X0 | Xt = x] x β(t)2 = α̇(t) + st (x) + β̇(t) (−β(t)st (x)) α(t) α(t) α̇(t) α̇(t)β(t)2 = x+ − β̇(t)β(t) st (x) α(t) α(t) ! α̇(t) β̇(t) α̇(t) 2 x + β(t) − st (x). = α(t) α(t) β(t) This is exactly (C.1).
C.2
SDE and ODE marginal equivalence
Proposition C.2 (Marginal equivalence). Let vt and st = ∇ log qt be the true velocity and score fields for marginals qt . Then the SDE dXt = vt (Xt ) + 21 σ(t)2 st (Xt ) dt + σ(t) dWt and the probability flow ODE dXt /dt = vt (Xt ) produce identical marginal distributions qt at all times t, for any noise schedule σ(t) ≥ 0.
22
Proof. The Fokker–Planck equation describes how the density ρt of a diffusion process evolves. For an SDE of the form dXt = b(Xt , t)dt + σ(t)dWt , the density satisfies ∂t ρt = −∇ · (b ρt ) +
σ(t)2 ∆ρt . 2
For our SDE with drift b = vt + 12 σ(t)2 st , this becomes ∂t ρt = −∇ ·
σ(t)2 ∆ρt . vt + 21 σ(t)2 st ρt + 2
We verify that ρt = qt satisfies this equation. Using st = ∇ log qt , we have st qt = ∇qt . Therefore: −∇ ·
2 1 2 σ(t) st qt
+
σ(t)2 σ(t)2 σ(t)2 ∆qt = − ∇ · (∇qt ) + ∆qt 2 2 2 σ(t)2 σ(t)2 =− ∆qt + ∆qt = 0. 2 2
The diffusion and score terms cancel exactly, leaving ∂t qt = −∇ · (vt qt ), which is the continuity equation for the probability flow ODE. Since both the ODE and SDE satisfy the same continuity equation with the same initial condition q0 , they have identical marginals qt at all times.
23
D Proofs for Section 3: Understanding the Limitations of Flow and Score Matching This appendix collects proofs for the statements in section 3. The material follows the order of the main discussion: we first prove the no-certificate result for the standard probability-flow ODE, then collect the reward-geometry results used in proposition 3.2 into a single section, and finally prove the RL reward-regret bound (proposition D.7).
D.1
Distribution Shift and Error Accumulation
D.1.1
Proof of proposition 3.1: no reward certificate for the standard probability-flow sampler
We now prove the no-certificate claim by giving two complementary counterexamples for the probability-flow ODE. The first is a one-dimensional Gaussian velocity example with a closed-form calculation. It is useful because every quantity can be written down explicitly, so the qt versus pt mismatch is visible in a single formula and provides a simple illustration of what goes wrong. The second proof gives a general construction in the coordinate induced by the true target flow and applies equally to the velocity and score parametrizations. Throughout we use the setup in Section 2 and the following standard assumptions and notation. We let (α, β) be a C 1 interpolation schedule (1) with α(0) = 0,
β(0) = 1,
α(1) = 1,
β(1) = 0,
and α(t), β(t) > 0 for t ∈ (0, 1). For a target endpoint law q, write qt for the interpolation marginal, sq for its score, and vq for its probability-flow velocity. For score-parametrized fields we also write ! α̇(t) β̇(t) 2 γ(t) := β(t) − , (9) α(t) β(t) so that vq (x, t) =
α̇(t) x + γ(t)sq (x, t) α(t)
by (C.1). Proposition D.1 (Velocity case: Gaussian variance inflation). For every ε > 0 and every δ ∈ (0, 1/4), there exist a target distribution q, a bounded reward r : R → [0, 1], and a velocity field vbad such that sup EX∼qt |vbad (X, t) − vq (X, t)|2 ≤ ε, t∈[0,1]
but if pbad denotes the endpoint law of the standard probability-flow ODE sampler X0 ∼ N (0, 1),
Ẋt = vbad (Xt , t), then
Eq [r] − Epbad [r] ≥ 1 − 2δ. Intuition. The marginals qt have analytical form given by: qt = N (0, ρτ (t)2 ),
ρτ (t)2 := α(t)2 τ 2 + β(t)2 .
The idea is to define a perturbation given by, ∆v(x, t) =
√
ε
x . ρτ (t)
Under qt , the normalized coordinate X/ρτ (t) is standard Gaussian, so the flow-matching error is exactly ε at every time. However, as time progresses this small error accumulates and becomes larger as |x| grows. This ends up increasing the variance of the rollout. Therefore, by choosing the target endpoint distribution to be sufficiently concentrated, we can make a reward on the high-likelihood regions of q1 to be missed by most of the distribution. 24
4
1.00
3
0.50
St
2 1
x
x
0.25 0.00
shell branches
thin overlap with qt
reward at t = 1 variance inflation
0.75
−0.25
0
−0.50
−1
−0.75
low-reward interval at t = 1
left shell branch sweeps the rollout inward
−2
−1.00 0
1
0
1
t
t
qt band true paths
pbad,t envelope rollout paths
reward at t = 1
qt band true paths
shell branches St rollout paths
low-reward interval
Figure 12 Appendix sketches of the two constructions. Left: velocity case, corresponding to Proposition D.1 and its
proof below. The true trajectories contract with qt , while the learned rollout trajectories and the one-sigma envelope of pbad,t show variance inflation under the same linear perturbation. Right: a one-dimensional Gaussian instance of Proposition D.3, drawn directly in x-space. In the general proof the perturbation is a shrinking radial shell in the target-flow coordinate; in one dimension that shell becomes two moving intervals, shown in orange here, and the left branch sweeps representative rollout trajectories into the low-reward interval at t = 1. In both panels, the shaded regions indicate typical mass under the relevant marginals, the dashed dark paths denote representative true trajectories, and the solid orange paths denote representative learned rollout trajectories. The obstruction is the same: training measures error under the data marginals qt , while sampling evaluates the learned field along the rollout marginals pt .
Remark D.2. This is not the same statement in the main text but it immediately disproves the existence of a guarantee of the form “with small flow-matching error, we obtain a small reward gap”. Proof. Fix ε > 0 and δ ∈ (0, 1/4). Let τ ∈ (0, 1) be chosen later and set q := N (0, τ 2 ). Under the interpolation (1), X1 ∼ q,
Xt = α(t)X1 + β(t)X0 , the marginal is
qt = N (0, ρτ (t)2 ),
X0 ∼ N (0, 1),
ρτ (t)2 := α(t)2 τ 2 + β(t)2 .
The corresponding probability-flow velocity is vq (x, t) =
ρ̇τ (t) x. ρτ (t)
Indeed, Ẋt = α̇(t)X1 + β̇(t)X0 . Since (Xt , Ẋt ) are jointly centered Gaussian, the conditional expectation is linear: vq (x, t) = E[Ẋt | Xt = x] = 25
Cov(Ẋt , Xt ) x. Var(Xt )
Using the independence of X1 and X0 , 1 d ρτ (t)2 = ρτ (t)ρ̇τ (t), 2 dt
Cov(Ẋt , Xt ) = α(t)α̇(t)τ 2 + β(t)β̇(t) = while Var(Xt ) = ρτ (t)2 . Therefore
ρ̇τ (t) x. ρτ (t)
vq (x, t) = Now define
vbad (x, t) := vq (x, t) + ∆v(x, t),
∆v(x, t) :=
√
ε
x . ρτ (t)
If X ∼ qt , then E[X 2 ] = ρτ (t)2 , so Eqt |vbad (X, t) − vq (X, t)|2 = ε Eqt
X2 = ε. ρτ (t)2
Thus the training error is exactly ε uniformly in time. The learned ODE is
Ẋt =
√ ρ̇τ (t) ε + Xt . ρτ (t) ρτ (t)
Solving this scalar linear ODE gives Xt = ρτ (t) exp
√
Z t ε
du 0 ρτ (u)
X0 .
Hence the rollout marginal is 2
pbad,t = N (0, ρτ (t) k(t)),
√ k(t) := exp 2 ε
In particular, Epbad,t |∆v(X, t)|2 = ε Epbad,t
Z t
du 0 ρτ (u)
.
X2 = εk(t), ρτ (t)2
so the same perturbation is measured as ε under qt but as εk(t) under the rollout law pbad,t . We now show that k(1) can be made arbitrarily large by taking τ small. Because β is C 1 with β(1) = 0, there exist M > 0 and t0 < 1 such that for all t ∈ [t0 , 1].
β(t) ≤ M (1 − t) Let A := supt∈[0,1] |α(t)| < ∞. Then ρτ (t) = Therefore
Z 1 0
for t ∈ [t0 , 1].
p p α(t)2 τ 2 + β(t)2 ≤ A2 τ 2 + M 2 (1 − t)2
dt ≥ ρτ (t)
Z 1
1 dt p = sinh−1 2 2 2 2 M A τ + M (1 − t) t0
which diverges as τ ↓ 0. Hence k(1) → ∞ as τ ↓ 0. Choose K > 0 such that, for Z ∼ N (0, 1),
Pr(|Z| ≤ K) ≥ 1 − δ. Define r(x) := 1{|x| ≤ Kτ }. 26
M (1 − t0 ) Aτ
,
Under the target endpoint law q = N (0, τ 2 ), Eq [r] = Pr(|Z| ≤ K) ≥ 1 − δ. Under the bad endpoint law pbad = N (0, τ 2 k(1)), K Epbad [r] = Pr |Z| ≤ p k(1)
!
r ≤
2 K p . π k(1)
Here we used that, for Z ∼ N (0, 1) and any a ≥ 0, Z a
r 1 −z2 /2 2 √ e dz ≤ a. Pr(|Z| ≤ a) = π 2π −a p p Since k(1) → ∞, choose τ small enough that 2/π K/ k(1) ≤ δ. Then Eq [r] − Epbad [r] ≥ 1 − 2δ. For the general proof, we use the standard notion of a flow map; see, for example, (Arnold, 1978; Boffi et al., 2025). Intuitively, the flow map is the function that, given an initial time a, a final time t, and a point y, returns the point reached by evolving y from time a to time t. Formally, given 0 < a < 1 and a point y ∈ Rd , let X a,y denote the solution of Ẋt = vq (Xt , t), Xa = y. The associated flow map is defined by
Φa,t (y) := Xta,y ,
that is, Φa,t (y) is the point reached at time t by the exact probability-flow ODE initialized at y at time a. Proposition D.3 (Radial-shell no-certificate construction). Let q be the endpoint target law, let (qt )t∈[0,1] be the interpolation marginals, let vq and sq denote the corresponding population velocity and score, and let r : Rd → [0, 1] be a bounded reward. Fix 0 < a < b < 1. Assume that the exact probability-flow ODE Ẋt = vq (Xt , t) generates flow maps Φa,t : Rd → Rd that are C 1 diffeomorphisms for every t ∈ [a, 1]. Assume further that there exist x⋆ ∈ Rd , ρ⋆ > 0, and η ∈ [0, 1] such that r(x) ≤ η
whenever ∥x − x⋆ ∥ ≤ ρ⋆ .
Then for every ε > 0 and every δ ∈ (0, 1), the following two statements hold. Velocity parametrization. There exists a field vbad such that
vbad (x, t) = vq (x, t)
for all t ∈ / [a, b],
and sup EX∼qt ∥vbad (X, t) − vq (X, t)∥2 ≤ ε.
t∈[0,1]
Let pbad denote the endpoint law of the probability-flow ODE sampler (2) run with vbad in place of vq . Then Epbad [r] ≤ η + δ. Score parametrization. If inf t∈[a,b] |γ(t)| > 0, then there exists a field sbad such that
sbad (x, t) = sq (x, t)
for all t ∈ / [a, b],
and sup EX∼qt ∥sbad (X, t) − sq (X, t)∥2 ≤ ε.
t∈[0,1]
27
Define the induced velocity field vbad (x, t) =
α̇(t) x + γ(t)sbad (x, t), α(t)
and let pbad denote the endpoint law of the probability-flow ODE sampler (2) driven by this vbad . Then Epbad [r] ≤ η + δ. Intuition. The goal is to hide a harmful perturbation inside a small score-matching or flow-matching error. We do this with a thin shell that moves inward over time. Away from the shell, the learned field agrees with the true score/flow, so if the shell is thin enough then the training error stays small. At the same time, the shell is arranged to catch trajectories and carry them into a low-reward region by time 1. Figure 12 (right) gives a sketch in one dimension, where the shell becomes two moving intervals. We apply the construction only on a subinterval [a, b] ⊂ (0, 1) to avoid the endpoints. However, designing the shell directly in x-space is awkward because the true dynamics already move trajectories around on their own, so the shell’s effect gets tangled with the motion of the original dynamics. We avoid this by working in a coordinate in which the true dynamics are trivial: for a path Xt , define the inverse-flow coordinate Yt := Φ−1 a,t (Xt ),
which records the starting point at time a of the trajectory passing through Xt . By construction, Yt is constant along the true ODE, so any motion of a trajectory in Y -space comes entirely from the perturbation. We then translate the motion in this space back to x-space and check that the training error is small while still carrying trajectories into the low-reward region. Proof. The proof is relatively straightforward but requires some bookkeeping. We organize it in steps for clarity. Under the assumptions of the proposition, the laws qt for t ∈ [a, 1] admit densities. With a slight abuse of notation, we use the same symbol for the law and its density when no confusion arises. Step 1: Work in the target-flow coordinate and rewrite the reward there. For any path Xt on [a, 1], define
Yt := Φ−1 a,t (Xt ). This coordinate records where Xt came from at time a under the true target flow. If Xt follows the true dynamics then Xt = Φa,t (Ya ), so Yt = Ya is constant. Because the shell construction will be carried out in Y -space, it is convenient to rewrite the endpoint reward in the same coordinate. Set µ := qa ,
B(z, ρ) := {y ∈ Rd : ∥y − z∥ ≤ ρ},
Let
re(y) := r(Φa,1 (y)).
z := Φ−1 a,1 (x⋆ ),
and consider the open ball {x ∈ Rd : ∥x − x⋆ ∥ < ρ⋆ } in endpoint space. Because Φa,1 is a diffeomorphism, its preimage is an open neighborhood of z, so there exists ρ > 0 such that B(z, ρ) ⊂ Φ−1 x ∈ Rd : ∥x − x⋆ ∥ < ρ⋆ . a,1 Therefore re(y) ≤ η
for all y ∈ B(z, ρ).
Thus it is enough to construct a bad rollout whose terminal target-flow coordinate Y1 lies in the ball B(z, ρ). In the remaining steps we will first define the bad dynamics directly in this coordinate, where the true dynamics are frozen, and then recover the corresponding path in x-space by applying the flow map.
28
Step 2: Build the shrinking shell in the target-flow coordinate. We first choose a large ball that contains almost
all of the mass at time a. Because
as R → ∞,
µ(B(z, R)) → 1 we can choose R > ρ such that
µ(B(z, R)) ≥ 1 − δ.
Our goal is to build a thin shell that starts just outside B(z, R) and then moves inward until it lies inside the low-reward ball B(z, ρ). Let
R + 1 − ρ/2 ℓ(t) := R + 1 − c(t − a) for t ∈ [a, b]. T Here T is the length of the active interval, c is the inward speed, and ℓ(t) is the radius of the shell at time t. Then ℓ(a) = R + 1, ℓ(b) = ρ/2. T := b − a,
c :=
So the shell starts strictly outside B(z, R) and ends strictly inside B(z, ρ). Fix w > 0, to be chosen later, with
w < min
1 ρ , 2 4
.
This will be the thickness scale of the shell. The bound w ≤ 1/2 is only used later, when we keep the shell inside a fixed radial region for the training-error estimate. Define a piecewise-linear function ψw : R → [0, c] that specifies how strongly the perturbation acts as a function of signed distance from the shell: 0, s ≤ −w, s c 1+ , −w ≤ s ≤ 0, w ψw (s) :=
c, 0 ≤ s ≤ w, s c 2− , w ≤ s ≤ 2w, w 0, s ≥ 2w.
This function is zero outside [−w, 2w], ramps up linearly on [−w, 0], is flat at height c on [0, w], and ramps down linearly on [w, 2w]. For y ̸= z, define the outward radial direction n(y) :=
y−z , ∥y − z∥
and set n(z) := 0. Now define the perturbation in the target-flow coordinate by uw (y, t) := −1{t ∈ [a, b]} ψw (∥y − z∥ − ℓ(t)) n(y). The minus sign means that the perturbation pushes points inward. It is supported on the moving shell St := {y ∈ Rd : ∥y − z∥ − ℓ(t) ∈ [−w, 2w]}. So at time t, the perturbation is active only on radii in the interval [ℓ(t)−w, ℓ(t)+2w]. Because ℓ(t) ∈ [ρ/2, R+1] and w < min{1/2, ρ/4}, every point y ∈ St satisfies ∥y − z∥ ∈ [ℓ(t) − w, ℓ(t) + 2w] ⊂ [ρ/4, R + 2]. 29
Thus the shell St always lies inside the fixed radial region {ρ/4 ≤ ∥y − z∥ ≤ R + 2}. Step 3: Define and analyze the bad dynamics in Y -space. Under the true dynamics, the target-flow coordinate
does not move: Yt is constant. We now define the bad dynamics in this same coordinate by letting Yt solve on [a, 1]
Ẏt = uw (Yt , t)
with initial law Ya ∼ µ. Because the shell support stays in the radial region {ρ/4 ≤ ∥y − z∥ ≤ R + 2}, it stays a positive distance from z. Therefore the field uw (·, t) is globally Lipschitz in y for each t and piecewise continuous in t. Hence this ODE has a unique absolutely continuous solution on [a, 1]. The mechanism has two phases: first the shell catches a trajectory while the trajectory stays fixed in Y -space, and then the shell carries the trajectory inward until it reaches B(z, ρ). Let ϱt := ∥Yt − z∥,
St := ϱt − ℓ(t).
Here ϱt is the distance to the shell center and St records the trajectory’s position relative to the moving shell. The sign of St tells us where the trajectory sits: St < −w means it is strictly inside the shell, while St ∈ [−w, 0] means it has been caught by the shell’s inner side.
Consider t ∈ [a, b]. Because uw (y, t) always points in the radial direction n(y), it changes only the distance to z. For ϱt > 0, ⊤ Yt − z Ẏt = n(Yt )⊤ uw (Yt , t) = −ψw (St ). ϱ̇t = ∥Yt − z∥
Since ℓ̇(t) = −c, we also have
Ṡt = ϱ̇t − ℓ̇(t) = −ψw (St ) + c.
If Ya ∈ B(z, ρ), then whenever ϱt > 0,
ϱ̇t = −ψw (St ) ≤ 0,
so the trajectory stays in B(z, ρ). It therefore remains to consider the case ρ ≤ ϱa ≤ R. In this case Sa = ϱa − ℓ(a) = ϱa − (R + 1) ∈ [ρ − R − 1, −1].
Since w < 1 the trajectory starts strictly inside the shell, with Sa < −w.
While St ≤ −w, the function ψw vanishes and
Ṡt = ϱ̇t − ℓ̇(t) = 0 − (−c) = c. So before capture, the shell simply moves inward until it catches the trajectory. The latest catch occurs for the smallest radius that still needs to be captured, namely ϱa = ρ. Then the catch time is at most −Sa − w R+1−ρ−w R + 1 − ρ/2 = < = b − a, c c c where we used w < ρ/2. Hence every trajectory with Ya ∈ B(z, R) \ B(z, ρ) reaches the inner side of the shell before time b. Once St ∈ [−w, 0], the definition of ψw gives
St Ṡt = −c 1 + w
c + c = − St . w
Therefore St remains in [−w, 0] and moves toward 0. In particular, Sb ≤ 0. 30
Since ℓ(b) = ρ/2, we conclude that ϱb = ℓ(b) + Sb ≤ ℓ(b) = ρ/2 < ρ. So every trajectory that starts in B(z, R) lies in B(z, ρ) by time b. For t > b, the perturbation is switched off, so uw (Yt , t) = 0 and the target-flow coordinate remains fixed. Hence Y1 ∈ B(z, ρ) on the event {Ya ∈ B(z, R)}.
At this point the geometry is complete: every trajectory that starts in B(z, R) has been moved into B(z, ρ) by time 1. By Step 1, re(Y1 ) ≤ η on {Ya ∈ B(z, R)}. Since Ya ∼ µ and µ(B(z, R)) ≥ 1 − δ, we obtain
E[e r(Y1 )] ≤ η · µ(B(z, R)) + 1 · µ(B(z, R)c ) ≤ η + δ.
This is the bad behavior we want in the target-flow coordinate. It remains to realize these same dynamics as a perturbation of the probability-flow ODE in x-space. Step 4: Push the construction to x-space. Let Xt denote the exact probability-flow ODE started from
X0 ∼ N (0, I).
Ẋt = vq (Xt , t), Then Xa ∼ qa = µ. Set
Ya := Xa ,
let Yt evolve by Ẏt = uw (Yt , t)
on [a, 1],
( Xt ,
t ∈ [0, a],
and define the bad rollout by X̄t :=
Φa,t (Yt ),
This is well defined because
t ∈ [a, 1].
Φa,a (Ya ) = Ya = Xa = X̄a . Moreover, because X̄1 = Φa,1 (Y1 ), the definition of re and the Step 3 bound give E[r(X̄1 )] = E[r(Φa,1 (Y1 ))] = E[e r(Y1 )] ≤ η + δ. This proves the endpoint reward bound for the constructed path. It remains to identify this path with the standard probability-flow ODE driven by a suitable field. For t ∈ [a, 1], differentiating the second branch gives X̄˙ t = ∂t Φa,t (Yt ) + DΦa,t (Yt )Ẏt = vq (Φa,t (Yt ), t) + DΦa,t (Yt )uw (Yt , t) Here DΦa,t (y) denotes the Jacobian matrix of Φa,t with respect to the spatial variable y. Because Φa,t is a diffeomorphism, the second term can be written as a function of (x, t) alone. Define ( ∆v(x, t) :=
−1 DΦa,t (Φ−1 a,t (x)) uw (Φa,t (x), t), t ∈ [a, b],
t∈ / [a, b].
0,
and set vbad (x, t) := vq (x, t) + ∆v(x, t).
31
By definition, ∆v(x, t) = 0 for t < a, so vbad = vq on [0, a). Since X̄t = Xt there, X̄ agrees with the exact ODE on [0, a). On the other hand, for t ∈ [a, 1] the computation above gives X̄˙ t = vbad (X̄t , t). Hence X̄ is a solution of Ẋt = vbad (Xt , t) started from X̄0 ∼ N (0, I). Therefore the endpoint law of X̄1 is exactly pbad , and so Epbad [r] ≤ η + δ. This proves the endpoint reward bound for the velocity construction. It remains to show that the training error under qt is small. Step 5: The velocity error is small under qt . We now show that we can choose w small enough to ensure
sup EX∼qt ∥vbad (X, t) − vq (X, t)∥2 = sup EX∼qt ∥∆v(X, t)∥2 ≤ ε.
t∈[0,1]
t∈[0,1]
If t ∈ / [a, b], then ∆v(·, t) = 0, so there is nothing to prove. It therefore suffices to consider t ∈ [a, b]. In that case, −1 ∆v(x, t) = DΦa,t (Φ−1 a,t (x)) uw (Φa,t (x), t). Now fix t ∈ [a, b], let X ∼ qt , and set
Y := Φ−1 a,t (X),
Because qt = (Φa,t )# µ, we have Y ∼ µ = qa . Therefore Z 2 2 EX∼qt ∥∆v(X, t)∥ = ∥DΦa,t (y)uw (y, t)∥ qa (y) dy. The key observation is that uw (·, t) is supported on a thin radial shell. Indeed, by definition, uw (y, t) ̸= 0
=⇒
∥y − z∥ ∈ [ℓ(t) − w, ℓ(t) + 2w].
This interval has length 3w. Thus the whole expectation is concentrated on a thin shell in Y -space, and it remains to control the weighted qa -mass of such shells. Since the shell support always stays in the radial range [ρ/4, R + 2], it is enough to control thin radial shells inside the fixed radial region AR := y ∈ Rd : ρ/4 ≤ ∥y − z∥ ≤ R + 2 . Because a < 1, we have β(a) > 0, so qa is a Gaussian-smoothed version of q. In particular, qa has bounded density: Z 1 ∥x − α(a)u∥2 1 qa (x) = exp − q(u) du ≤ . 2β(a)2 (2πβ(a)2 )d/2 (2πβ(a)2 )d/2 Because Φa,t is a C 1 flow, its Jacobian is uniformly bounded on the compact set [a, b] × AR . Also, on the bounded radial region AR , the volume of a radial shell is proportional to its thickness. Combining these two facts with the density bound above, there exists a finite constant CR such that for every interval I ⊂ [ρ/4, R+2] and every t ∈ [a, b], Z {y: ∥y−z∥∈I}
∥DΦa,t (y)∥2op qa (y) dy ≤ CR |I|.
Indeed, if SI := {y : ∥y − z∥ ∈ I}, then bounded density and bounded Jacobian on [a, b] × AR give Z ′ ∥DΦa,t (y)∥2op qa (y) dy ≤ CR Vol(SI ), SI
32
′′ while the volume of a radial shell inside the bounded annulus AR satisfies Vol(SI ) ≤ CR |I|. Absorbing the constants gives the claim. It := [ℓ(t) − w, ℓ(t) + 2w].
This interval has length 3w and is contained in [ρ/4, R + 2]. Since uw (·, t) is supported where ∥y − z∥ ∈ It , and |uw (y, t)| ≤ c, we obtain EX∼qt ∥∆v(X, t)∥2 ≤
Z
2
2
∥DΦa,t (y)∥op |uw (y, t)| qa (y) dy Z 2 ≤ c2 ∥DΦa,t (y)∥op qa (y) dy {y: ∥y−z∥∈It }
≤ c2 CR |It | = 3c2 CR w. Hence choosing
w ≤ min
ensures
1 ρ ε , , 2 4 3c2 CR
sup EX∼qt ∥vbad (X, t) − vq (X, t)∥2 ≤ ε.
t∈[0,1]
Step 6: Score parametrization. For the score statement, keep the same reward, the same center z, and the same
shell radii. If needed, choose a smaller width w; this changes the perturbation but not the reward bound. Set γ− := inf |γ(t)| > 0. t∈[a,b]
Define
∆v(x, t) , γ(t) ∆s(x, t) := 0,
Set
t ∈ [a, b], t∈ / [a, b].
sbad (x, t) := sq (x, t) + ∆s(x, t), and define vbad (x, t) =
α̇(t) x + γ(t)sbad (x, t). α(t)
vq (x, t) =
α̇(t) x + γ(t)sq (x, t), α(t)
Because
we obtain vbad (x, t) = vq (x, t) + ∆v(x, t). So the score-parametrized model induces exactly the same bad ODE rollout as above, and therefore the same endpoint reward bound. It remains to check the DSM error. If t ∈ / [a, b], the error is zero. If t ∈ [a, b], then 1 EX∼qt ∥sbad (X, t) − sq (X, t)∥2 = EX∼qt ∥∆s(X, t)∥2 ≤ 2 EX∼qt ∥∆v(X, t)∥2 . γ− Using the bound from Step 5 gives 3c2 CR EX∼qt ∥sbad (X, t) − sq (X, t)∥2 ≤ w. 2 γ− 33
Thus choosing w ≤ min ensures
2 1 ρ εγ− , , 2 2 4 3c CR
sup EX∼qt ∥sbad (X, t) − sq (X, t)∥2 ≤ ε.
t∈[0,1]
This proves the score statement.
D.2
Geometry mismatch
D.2.1
Reward geometry under stronger rollout control
This subsection collects the geometry results used in the main text into one place. We first prove the global certificate stated in proposition 3.2, then derive from the same coupling argument a weaker on-policy variant under the rollout marginals pt . The proposition also records a matching tightness example, and we close with the indicator-reward counterexample that explains why this kind of geometric control applies only to Lipschitz rewards. We use only standard ODE stability ideas and the following elementary differential form of Grönwall’s inequality; see, for example, Arnold (1978) for background. Lemma D.4 (Grönwall inequality). Let m : [0, T ] → [0, ∞) be absolutely continuous, let b : [0, T ] → [0, ∞) be integrable, and let L ≥ 0. If m′ (t) ≤ b(t) + Lm(t) then Lt
m(t) ≤ e m(0) +
Z t
for almost every t ∈ [0, T ],
eL(t−s) b(s) ds
0
for every t ∈ [0, T ].
Proof. The idea is to remove the linear growth term first. Set z(t) := m(t)e−Lt . Then z is absolutely continuous, and for almost every t, z ′ (t) = e−Lt m′ (t) − Lm(t) ≤ e−Lt b(t). Integrating from 0 to t gives z(t) − z(0) ≤
Z t
e−Ls b(s) ds.
0
Multiplying by eLt yields m(t) ≤ eLt m(0) +
Z t
eL(t−s) b(s) ds.
0
Proposition D.5 (Reward certificates under stronger rollout control). Consider the coupled ODEs Ẋt = v(Xt , t),
Ẋt∗ = v ∗ (Xt∗ , t),
X0 = X0∗ ∼ p0 ,
and write pt and p∗t for the laws of Xt and Xt∗ respectively. Assume that v ∗ (·, t) is Lv -Lipschitz in x, uniformly over t ∈ [0, 1]. 1. Global control. If
∥v − v ∗ ∥∞ := sup ∥v(x, t) − v ∗ (x, t)∥ ≤ ε, t,x
34
then
ε (eLv − 1), ∥X1 − X1∗ ∥ ≤ eLv (1−s) ε ds = Lv ε, 0 Z 1
Lv > 0,
a.s.
Lv = 0,
and for every Lr -Lipschitz reward r, |Ep1 [r] − Ep∗1 [r]| ≤ Lr
Z 1
eLv (1−s) ε ds.
0
2. On-policy control. Define the rollout error function e(t) := EXt ∼pt ∥v(Xt , t) − v ∗ (Xt , t)∥ . If e ∈ L1 ([0, 1]), then
E∥X1 − X1∗ ∥ ≤
Z 1
eLv (1−s) e(s) ds.
0
and for every Lr -Lipschitz reward r, |Ep1 [r] − E [r]| ≤ Lr
Z 1
p∗ 1
eLv (1−s) e(s) ds.
0
In particular, if supt∈[0,1] e(t) ≤ ε, then the same endpoint bound as in part 1 holds at t = 1. 3. Tightness. The exponential weighting factor eLv (t−s) and its dependence on Lv , Lr , and the error function cannot be improved in general. More precisely, there exists an explicit one-dimensional population CFM problem whose population minimizer is v ∗ (x, t) = Lv x and such that, for every nonnegative e ∈ L1 ([0, 1]), the perturbation ∆v(x, t) := e(t), v(x, t) = v ∗ (x, t) + ∆v(x, t) satisfies EXt ∼pt ∥v(Xt , t) − v ∗ (Xt , t)∥ = e(t) and E∥X1 − X1∗ ∥ =
Z 1
eLv (1−s) e(s) ds.
0
Thus part 2 is attained with equality for every admissible error function. For the linear reward r(x) = Lr x, the reward bound in part 2 is also attained with equality. In particular, taking e(t) ≡ ε gives ∥v−v ∗ ∥∞ = ε and attains the bound in part 1. Proof. The two certificates come from the same coupling argument. The only difference is how the field mismatch term is controlled: in part 1 it is bounded pathwise, while in part 2 it is bounded only after taking expectation. Set
∆t := Xt − Xt∗ .
Because both trajectories solve ODEs with the same initial condition, ∆ is absolutely continuous and ˙ t = v(Xt , t) − v ∗ (Xt∗ , t) ∆
for almost every t.
We now use this in two slightly different ways. Part 1: global control. Set
y(t) := ∥∆t ∥.
Since ∆ is absolutely continuous, so is y, and for almost every t, ˙ t ∥ ≤ ∥v(Xt , t) − v ∗ (Xt , t)∥ + ∥v ∗ (Xt , t) − v ∗ (X ∗ , t)∥ ≤ ε + Lv y(t). y ′ (t) ≤ ∥∆ t 35
The first term is controlled by the uniform bound ∥v − v ∗ ∥∞ ≤ ε, evaluated at the point Xt . The second term is controlled by the Lv -Lipschitz continuity of v ∗ (·, t) in its spatial argument, applied to the pair (Xt , Xt∗ ). In other words, we split the error into direct field mismatch of size at most ε, and the amplification of the current gap by the dynamics of v ∗ . Because y(0) = 0, applying lemma D.4 with b(t) ≡ ε gives ε Z t (eLv t − 1), Lv > 0, y(t) ≤ eLv (t−s) ε ds = Lv εt, 0 Lv = 0. Evaluating this at t = 1 gives ∥X1 − X1∗ ∥ ≤
Z 1
a.s.
eLv (1−s) ε ds
0
Taking expectation preserves this bound. Now let r be Lr -Lipschitz. Since |r(X1 ) − r(X1∗ )| ≤ Lr ∥X1 − X1∗ ∥ we obtain |Ep1 [r] − E
p∗ 1
[r]| ≤ Lr E∥X1 − X1∗ ∥ ≤ Lr
a.s., Z 1
eLv (1−s) ε ds.
0
Part 2: on-policy control. Part 1 relied on a pathwise bound for the field mismatch term at each time t. Under
on-policy control we no longer have that: the quantity ∥v(Xt , t) − v ∗ (Xt , t)∥ is controlled only after averaging over the rollout law. We therefore rewrite the trajectory gap in integral form and then take expectation, which puts the assumption on the error function e(t) = EXt ∼pt ∥v(Xt , t) − v ∗ (Xt , t)∥ directly into the estimate. Starting from the integral form, Z t ∆t = v(Xs , s) − v ∗ (Xs∗ , s) ds 0
Z t = 0
v(Xs , s) − v ∗ (Xs , s) ds +
Z t 0
v ∗ (Xs , s) − v ∗ (Xs∗ , s) ds.
Taking norms and using the triangle inequality gives Z t Z t ∥∆t ∥ ≤ ∥v(Xs , s) − v ∗ (Xs , s)∥ ds + ∥v ∗ (Xs , s) − v ∗ (Xs∗ , s)∥ ds. 0
0
Now use that v (·, s) is Lv -Lipschitz in x, so ∗
∥v ∗ (Xs , s) − v ∗ (Xs∗ , s)∥ ≤ Lv ∥Xs − Xs∗ ∥ = Lv ∥∆s ∥. Substituting this back yields ∥∆t ∥ ≤ Now define
Z t 0
∥v(Xs , s) − v ∗ (Xs , s)∥ ds + Lv
Z t 0
(∗)
∥∆s ∥ ds.
m(t) := E∥Xt − Xt∗ ∥ = E∥∆t ∥.
Taking expectation in (∗) and using Fubini’s theorem yields Z t Z t Z t Z t ∗ m(t) ≤ E∥v(Xs , s) − v (Xs , s)∥ ds + Lv m(s) ds = e(s) ds + Lv m(s) ds. 0
0
0
To bring this into differential form, set Z t M (t) :=
Z t e(s) ds + Lv
0
m(s) ds. 0
36
0
Then m(t) ≤ M (t), M (0) = 0, and for almost every t, M ′ (t) = e(t) + Lv m(t) ≤ e(t) + Lv M (t). Applying lemma D.4 to M with b = e gives M (t) ≤
Z t
eLv (t−s) e(s) ds.
0
Evaluating at t = 1 and using m ≤ M , we conclude that Z 1 m(1) ≤ eLv (1−s) e(s) ds. 0
Therefore E∥X1 − X1∗ ∥ ≤ If r is Lr -Lipschitz, then so
Z 1
eLv (1−s) e(s) ds.
0
|r(X1 ) − r(X1∗ )| ≤ Lr ∥X1 − X1∗ ∥
|Ep1 [r] − Ep∗1 [r]| ≤ Lr E∥X1 − X1∗ ∥ ≤ Lr If supt e(t) ≤ ε, then
Z 1 e
Lv (1−s)
0
e(s) ds ≤
Z 1
Z 1
a.s., eLv (1−s) e(s) ds.
0
eLv (1−s) ε ds,
0
which recovers the same endpoint constant as in part 1. Part 3: tightness. Work in one dimension. Fix any C 1 function
h πi θ : [0, 1] → 0, 2 satisfying θ(0) = 0,
θ(1) =
π θ(t) ∈ 0, 2
π , 2
Let
U1 ∼ N (0, e2Lv ),
U0 ∼ N (0, 1),
with U0 and U1 independent, and define
where
for t ∈ (0, 1).
X̄t = α(t)U1 + β(t)U0 ,
Yt = α̇(t)U1 + β̇(t)U0 ,
α(t) := eLv (t−1) sin θ(t),
β(t) := eLv t cos θ(t).
Then α(0) = 0, β(0) = 1, α(1) = 1, and β(1) = 0, with α(t), β(t) > 0 for t ∈ (0, 1), so this is a valid interpolation from N (0, 1) to N (0, e2Lv ). The population CFM minimizer is v ∗ (x, t) = E[Yt | X̄t = x]. Since (X̄t , Yt ) is centered and jointly Gaussian, v ∗ (x, t) =
Cov(Yt , X̄t ) x. Var(X̄t )
A direct calculation from the definitions of α and β gives Var(X̄t ) = α(t)2 e2Lv + β(t)2 = e2Lv t , 37
Cov(Yt , X̄t ) = Lv e2Lv t .
Hence v ∗ (x, t) = Lv x. Now fix any nonnegative function e ∈ L1 ([0, 1]) and define v(x, t) := v ∗ (x, t) + ∆v(x, t) = Lv x + e(t).
∆v(x, t) := e(t),
Let W0 ∼ N (0, 1) and run the coupled ODEs Ẋt∗ = v ∗ (Xt∗ , t),
Ẋt = v(Xt , t),
Xt∗ = eLv t W0 ,
Xt = eLv t W0 +
X0∗ = X0 = W0 .
A direct integration gives Z t
eLv (t−s) e(s) ds.
0
Therefore X1 − X1∗ =
Z 1
eLv (1−s) e(s) ds.
0
Since ∆v is independent of x, the on-policy error function coincides with the prescribed function: EXt ∼pt |v(Xt , t) − v ∗ (Xt , t)| = e(t). Therefore E|X1 − X1∗ | =
Z 1
eLv (1−s) e(s) ds,
0
so part 2 is attained with equality. For the linear reward r(x) = Lr x, |Ep1 [r] − Ep∗1 [r]| = Lr |E[X1 − X1∗ ]| = Lr
Z 1
eLv (1−s) e(s) ds.
0
Thus the reward bound in part 2 is also attained with equality. If e(t) ≡ ε, then ∥v − v ∗ ∥∞ = ε and
ε (eLv − 1), Lv > 0, ∗ Lv (1−s) E|X1 − X1 | = e ε ds = Lv ε, 0 Lv = 0. Z 1
so part 1 is also attained with equality. The exact reward-tightness statement above uses the linear test function r(x) = Lr x, which is Lipschitz but unbounded. If one insists on bounded rewards while keeping this same amplification-tight construction, the endpoint Gaussians can be rescaled to make range truncation negligible. Indeed, replace the endpoint Gaussians by U0 ∼ N (0, σ 2 ) and U1 ∼ N (0, σ 2 e2Lv ). The calculation of the CFM minimizer is unchanged, so v ∗ (x, t) = Lv x, and the perturbation ∆v(x, t) = e(t) still shifts every endpoint by Z 1 A := eLv (1−s) e(s) ds. 0
Taking σ small makes X1∗ arbitrarily concentrated, so a clipped affine reward of slope Lr can be placed so
that, with probability arbitrarily close to one, the shift by A either remains inside the linear region, giving gap Lr A, or crosses from the zero plateau to the one plateau, giving gap 1. Hence for every ζ > 0 there is a bounded Lr -Lipschitz reward such that Ep1 [r] − Ep∗1 [r] ≥ min{1, Lr A} − ζ.
Thus the bounded-reward version of this amplification-tight construction is essentially tight as well.
38
Remark: A simpler bounded-reward tightness statement. The tightness construction in part 3 is engineered to
make the full amplification kernel eLv (1−s) sharp, and for that reason it uses a somewhat contrived setup. This is useful for completeness, but it is more than is needed to show the main geometric point in the paper, namely that the factor Lr ε itself cannot be removed. The next result isolates that point already for the standard linear Gaussian bridge with p0 = q = N (0, 1).
Theorem D.6 (Bounded-reward tightness of the Lr ε factor). There is a universal constant c > 0 such that the following holds. Fix ε > 0 and L > 0 with εL < 1, and consider the one-dimensional CFM instance with standard linear interpolation i.i.d.
Xt = tX1 + (1 − t)X0 ,
X0 , X1 ∼ N (0, 1),
whose source and target laws are both N (0, 1) and whose population velocity v ∗ is 1-Lipschitz in space. Then there is a velocity field v satisfying sup t∈[0,1], x∈R
|v(x, t) − v ∗ (x, t)| ≤ ε
and a reward r : R → [0, 1] with Lip(r) ≤ L such that the endpoint law p of the ODE driven by v satisfies Ep [r] − Eq [r] ≥ c Lε,
q = N (0, 1).
The explicit nonsmooth construction below gives c = 1/2. Intuition. The idea is to build a reward together with a matching velocity perturbation that pushes every generated sample uphill of the reward. The reward is a periodic triangular wave: peaks of height 1 spaced 2/L apart, falling linearly to 0 at the midpoints. This makes the reward bounded in [0, 1], L-Lipschitz, and ensures that every point is within distance 1/L of a peak. The bad velocity field is v ∗ plus an admissible nudge of size ε pointing each sample toward its nearest peak. Over unit time, this moves each sample up a slope of magnitude L until it has traveled distance ε or reached a peak, so the reward gain is of order Lε on average. Since every sample moves locally uphill, there is no cancellation between samples moving in different directions. The gain is therefore insensitive to how spread out the samples are, which is why the standard bridge with p0 = q = N (0, 1) already realizes an Ω(Lε) gap. Proof. Write
p ρt := t2 + (1 − t)2 . √ Then qt = N (0, ρ2t ), with ρ0 = ρ1 = 1 and 1/ 2 ≤ ρt ≤ 1.
First, the reference flow is a pure rescaling. The population CFM velocity is v ∗ (x, t) = E[X1 − X0 | Xt = x]. Since (X0 , X1 , Xt ) is centered and jointly Gaussian, as earlier we can compute this explicitly as: v ∗ (x, t) = Since ∂x v ∗ (x, t) =
2t − 1 ρ̇t Cov(X1 − X0 , Xt ) x= 2 x = x. Var(Xt ) 2t − 2t + 1 ρt
2t − 1 , 2t2 − 2t + 1
2t − 1 ≤ 1 for all t ∈ [0, 1], 2t2 − 2t + 1
the field v ∗ (·, t) is 1-Lipschitz uniformly in t. Integrating Ẋt = (ρ̇t /ρt )Xt gives Xt = ρt X0 . Thus, in the rescaled coordinate Xt , Yt := ρt the reference dynamics are frozen, Ẏt = 0, and the reference endpoint is X1∗ = Y0 ∼ N (0, 1) = q.
We now construct a bounded reward and an uphill perturbation to get the desired bound. Set T := 2/L. For a phase θ ∈ [0, T ], to be chosen later, place reward peaks at the equally spaced points {nT − θ : n ∈ Z}. 39
Let dθ (y) be the distance from y to the nearest peak, and define the triangular-wave reward rθ (y) := 1 − L dθ (y).
Every point is at most half a period, T /2 = 1/L, from a nearest peak, so 0 ≤ rθ ≤ 1. Since dθ is 1-Lipschitz, rθ is L-Lipschitz. Let bθ (y) ∈ {−1, 0, +1} be the unit direction of steepest ascent of rθ : it equals +1 if a nearest peak lies to the right of y, −1 if a nearest peak lies to the left, and 0 at peaks and troughs. Define x vθ (x, t) := v ∗ (x, t) + ρt ε bθ . ρt Because ρt ≤ 1 and |bθ | ≤ 1,
∗
sup |vθ (x, t) − v (x, t)| = sup ρt ε bθ t,x
t,x
x ρt
≤ ε.
Substituting this field into the dynamics of Yt = Xt /ρt , the rescaling terms cancel and Ẏt = εbθ (Yt ). Thus, in the rescaled coordinate, every trajectory moves at speed ε toward a nearest reward peak and stops on arrival. Fix Y0 = y away from the null set of peaks and troughs. Over the unit time interval the trajectory moves a distance min{ε, dθ (y)} toward the nearest peak, so dθ decreases by exactly that amount and rθ (Y1 ) − rθ (Y0 ) = L min{ε, dθ (Y0 )}.
Couple the reference and perturbed flows through the same start Y0 ∼ N (0, 1). The reference endpoint is X1∗ = Y0 , while the perturbed endpoint is X1 = Y1 because ρ1 = 1. Therefore, for the fixed reward rθ , G(θ) := Ep [rθ ] − Eq [rθ ] = E rθ (Y1 ) − rθ (Y0 ) = L EZ∼N (0,1) min{ε, dθ (Z)} . It remains only to choose the phase. Average G over θ ∼ Unif[0, T ]. For each fixed z, as θ ranges over one period, the location of z relative to the nearest peak is uniform over one period. Hence dθ (z) is uniform on [0, T /2] = [0, 1/L] with density L. Since ε < 1/L, ! Z 1/L Z ε Z 1/L Lε2 . Eθ min{ε, dθ (z)} = L min{ε, u} du = L u du + ε du = ε − 2 0 0 ε This quantity is independent of z, so Lε2 Eθ G(θ) = L ε − 2
Lε 1 = Lε 1 − ≥ Lε, 2 2
where the last inequality uses Lε < 1. Therefore some phase θ∗ satisfies G(θ∗ ) ≥ Lε/2 by simple properties of the expectation. Taking r = rθ∗ and v = vθ∗ proves the theorem with c = 1/2.
D.3
How RL Helps
D.3.1
RL reward regret bound
Proposition D.7 (RL reward regret bound). Let r : X → [0, 1], let Pbase be the base trajectory distribution, define dP∗ exp(λr(X1 )) (X0:1 ) = , Z := EPbase [exp(λr(X1 ))], dPbase Z and let P be any trajectory distribution with final-sample marginal p. Then the final-sample marginal of P∗ is the tilted target p∗ defined in Section 2, and q q Ep∗ [r] − Ep [r] ≤ 12 KL(P ∥ P∗ ) = λ2 LRL (P) − LRL (P∗ ) , (10) where LRL (P) = λ1 KL(P ∥ Pbase ) − Ep [r]. 40
Proof. Here Pbase , P, and P∗ are distributions over full sampled trajectories X0:1 ; formally, they are path measures. Because the Radon–Nikodym derivative of P∗ depends only on the terminal state X1 , the final-sample marginal of P∗ is exactly exp(λr(x)) pbase (x) p∗ (x) = , Z which is exactly the tilted target. Since TV(p, p∗ ) =
sup f : 0≤f ≤1
the assumption r ∈ [0, 1] gives
|Ep [f ] − Ep∗ [f ]| ,
Ep∗ [r] − Ep [r] ≤ TV(p, p∗ ) .
With the convention TV(p, p∗ ) = 21 ∥p − p∗ ∥1 , Pinsker’s inequality yields TV(p, p∗ ) ≤
q
1 ∗ 2 KL(p ∥ p ),
and the data processing inequality for the endpoint map X0:1 7→ X1 gives KL(p ∥ p∗ ) ≤ KL(P ∥ P∗ ) . Therefore Ep∗ [r] − Ep [r] ≤
q
1 ∗ 2 KL(P ∥ P ).
The first inequality in (10) therefore follows by data processing for the endpoint map X0:1 7→ X1 , which gives KL(p ∥ p∗ ) ≤ KL(P ∥ P∗ ). It remains to identify this trajectory-level KL with the RL objective gap. By the definition of P∗ , dP KL(P ∥ P∗ ) = EP log ∗ dP dP dP∗ = EP log − EP log dPbase dPbase = KL(P ∥ Pbase ) − EP [λr(X1 ) − log Z] = KL(P ∥ Pbase ) − λEp [r] + log Z. Rearranging gives
1 log Z KL(P ∥ P∗ ) − . λ λ Since log Z/λ is constant, evaluating the same identity at P∗ and subtracting gives KL(P ∥ P∗ ) = λ LRL (P) − LRL (P∗ ) , LRL (P) =
which is the second equality in (10).
41
E
Proofs for Section 4: Method: Discriminator-Guided RL
E.1
Why Representation Spaces Help
E.1.1
Feature-space correction
For readability, we work throughout this subsection in the standard density setting. Assume every distribution we mention admits a well-defined continuous density, and that the corresponding conditionals p(x | z) for z = ϕ(x) are well defined. We write pϕ (z) for the density on Z induced by passing x ∼ p through ϕ, and assume that q ϕ (z) > 0 only where pϕbase (z) > 0. Proposition E.1 (Feature-space correction). Let ϕ : X → Z be an encoder, let h : Z → R be any scalar function, and let λ > 0. Assume Zh,λ := Ex∼pbase exp λh(ϕ(x)) < ∞. Define the tilted density by ph,λ (x) :=
exp(λh(ϕ(x))) pbase (x). Zh,λ
(11)
exp(λh(z)) ϕ pbase (z), Zh,λ
(12)
Then: (i) the feature marginal is tilted according to pϕh,λ (z) = where equivalently Zh,λ = Ez∼pϕ [exp(λh(z))] ; base
(ii) the conditional distribution inside each feature cell is unchanged: for pϕh,λ -a.e. z.
ph,λ (x | z) = pbase (x | z)
(13)
(iii) if h∗ (z) = log
q ϕ (z) pϕbase (z)
(14)
,
then pϕh∗ ,1 = q ϕ ; (iv) ph∗ ,1 is the unique solution of min KL(p ∥ pbase )
subject to
p
pϕ = q ϕ .
Proof. Let z = ϕ(x). For part (i), let g : Z → R be any bounded measurable test function. Then Ez∼pϕ [g(z)] = Ex∼ph,λ [g(ϕ(x))] h,λ
1 Ex∼pbase [g(ϕ(x)) exp(λh(ϕ(x)))] Zh,λ Z 1 = g(z) exp(λh(z)) pϕbase (z) dz. Zh,λ Z =
Thus the pushed-forward law under ϕ has density pϕh,λ (z) =
exp(λh(z)) ϕ pbase (z), Zh,λ
42
(15)
since the right-hand side gives the same integral against every bounded measurable test function g. Hence (12). Taking g ≡ 1 in the display above gives Z 1 1= exp(λh(z)) pϕbase (z) dz, Zh,λ Z so
Zh,λ = Ez∼pϕ [exp(λh(z))] . base
This completes part (i). For part (ii), fix a measurable set A ⊆ X , and let g : Z → R be any bounded measurable test function. On the one hand, conditioning under ph,λ gives Z Ex∼ph,λ [1A (x) g(ϕ(x))] = g(z) ph,λ (A | z) pϕh,λ (z) dz. Z
On the other hand, using the definition of ph,λ and then conditioning under pbase , 1 Ex∼pbase [1A (x) g(ϕ(x)) exp(λh(ϕ(x)))] Zh,λ Z 1 = g(z) exp(λh(z)) pbase (A | z) pϕbase (z) dz Zh,λ Z Z = g(z) pbase (A | z) pϕh,λ (z) dz,
Ex∼ph,λ [1A (x) g(ϕ(x))] =
Z
where the last step uses part (i). Since these two expressions agree for every bounded measurable test function g, their integrands must agree dz-a.e.: ph,λ (A | z) pϕh,λ (z) = pbase (A | z) pϕh,λ (z). Therefore ph,λ (A | z) = pbase (A | z)
for pϕh,λ -a.e. z.
Since this holds for every measurable A, the conditional distributions agree. In our density setting, we record this as ph,λ (x | z) = pbase (x | z),
which proves part (ii). For part (iii), let
h∗ (z) = log Then
Z Zh∗ ,1 =
Therefore, by part (i),
q ϕ (z) pϕbase (z)
exp(h∗ (z)) pϕbase (z) dz =
. Z
q ϕ (z) dz = 1.
pϕh∗ ,1 (z) = exp(h∗ (z)) pϕbase (z) = q ϕ (z),
proving part (iii). For part (iv), let p be any distribution satisfying pϕ = q ϕ . The KL chain rule gives KL(p ∥ pbase ) = KL pϕ ∥ pϕbase + Ez∼pϕ [KL(p(· | z) ∥ pbase (· | z))] . Under the constraint pϕ = q ϕ , this becomes KL(p ∥ pbase ) = KL q ϕ ∥ pϕbase + Ez∼qϕ [KL(p(· | z) ∥ pbase (· | z))] ≥ KL q ϕ ∥ pϕbase . 43
(16)
The first term is fixed by the constraint, and the second term is always nonnegative. Equality holds if and only if p(· | z) = pbase (· | z) for q ϕ -a.e. z.
Hence the unique minimizer is the distribution with feature marginal q ϕ and within-cell conditional pbase (· | z). Parts (ii) and (iii) show that ph∗ ,1 has exactly this marginal and these conditionals, so it attains the lower bound. Uniqueness follows because the marginal q ϕ together with the conditionals pbase (· | z) determines the distribution p uniquely. Remark. We stated this subsection in the density setting because it is the clearest way to express the main
idea: RL reweights feature cells by the feature-space density ratio and leaves the within-cell conditional unchanged. A more general measure-theoretic version replaces ordinary ratios such as q ϕ (z)/pϕbase (z) with the Radon–Nikodym derivative dq ϕ /dpϕbase , replaces p(x | z) with regular conditional distributions, and writes feature-space integrals against pϕbase (dz) rather than dz. The proof is otherwise the same. E.1.2
Feature-Space Test-Function Bound
Theorem E.2 (Feature-space test-function bound). Let r̂(x) = h(ϕ(x)) be a feature-based reward with implied target q̂λ ∝ exp(λr̂) pbase , and suppose pθ is ∆-suboptimal for the KL-regularized objective with trade-off parameter λ. Then for any g : Z → R with ∥g∥∞ ≤ C, √ Epθ [g(ϕ(x))] − Eq [g(ϕ(x))] ≤ 2λ∆} + Eq̂λ [g(ϕ(x))] − Eq [g(ϕ(x))] . (17) |C {z | {z } RL suboptimality
feature-space ratio error
Proof. By the triangle inequality, Epθ [g(ϕ(x))] − Eq [g(ϕ(x))] ≤ Epθ [g ◦ ϕ] − Eq̂λ [g ◦ ϕ] + Eq̂λ [g ◦ ϕ] − Eq [g ◦ ϕ] .
(18)
First term (RL suboptimality). Since ∥g∥∞ ≤ C, the function g ◦ ϕ takes values in [−C, C]. Also, because
q̂λ (x) ∝ exp(λr̂(x)) pbase (x),
LRL (p) =
1 1 1 KL(p ∥ pbase ) − Ep [r̂] = KL(p ∥ q̂λ ) − log Ẑλ λ λ λ
(19)
for a constant Ẑλ independent of p. Therefore the unique minimizer of the RL loss is q̂λ , and the ∆suboptimality assumption implies KL(pθ ∥ q̂λ ) ≤ λ∆. By Pinsker’s inequality, r √ λ∆ Epθ [g ◦ ϕ] − Eq̂λ [g ◦ ϕ] ≤ 2C TV(pθ , q̂λ ) ≤ 2C = C 2λ∆. (20) 2 Second term (feature-space ratio error). Since g ◦ ϕ depends on x only through ϕ(x),
Eq̂λ [g(ϕ(x))] = Eq̂ϕ [g] λ
and Eq [g(ϕ(x))] = Eqϕ [g],
so Eq̂λ [g ◦ ϕ] − Eq [g ◦ ϕ] = Eq̂ϕ [g] − Eqϕ [g] . Combining gives (17). λ
44
(21)
F
Adjoint Matching: Background
This appendix contains a self-contained review of adjoint matching (Domingo-Enrich et al., 2025), the method we use to optimize the RL objective LRL (4). We first provide a general description of the method in the SOC setting and then specialize to flow models. Adjoint matching solves LRL (4) under the lens of stochastic optimal control (SOC), a framework for fine-tuning the dynamics of an SDE through a learned drift correction. Concretely, given a base SDE with drift b and noise σ(t), SOC adds a learned control uθ (x, t) to the drift, dXt = b(Xt , t) + σ(t) uθ (Xt , t) dt + σ(t) dBt , X0 ∼ N (0, I), (22) Stochastic optimal control formulation.
and seeks the uθ minimizing a quadratic control cost minus the terminal reward, Z 1 2 1 min E 2 ∥uθ (Xt , t)∥ dt − λ r(X1 ) . uθ
(23)
0
In our setting b(x, t) = vt (x) + 12 σ(t)2 st (x) is the base drift of the sampling SDE (2)—the dynamics that generate pbase —and uθ is the deviation we are learning to fine-tune the model. By Girsanov’s theorem, the path-KL between (22) and the σ(t)-base SDE is exactly the expected control energy, Z 1 KL(Pθ ∥ Pbase ) = 12 E ∥uθ (Xt , t)∥2 dt , 0
so (23) is precisely the path-KL-regularized RL objective LRL (4). Optimizing uθ thus tunes the trajectory drift to push X1 toward high reward while staying close to the base process in path-KL. For flow models we parametrize uθ in terms of a learned velocity field; we defer the specifics to the specialization paragraph below. The memoryless schedule.
pbase (x) exp(λr(x)). Let
Solving (23) does not in general produce the desired endpoint tilt p∗ (x) ∝
Z 1 2 1 V (x, t) := min E 2 ∥uθ (Xs , s)∥ ds − λ r(X1 ) Xt = x uθ
t
denote the optimal cost-to-go from state x at time t. Domingo-Enrich et al. (2025) show that the SOC optimum picks up an X0 -dependent factor through the initial-time value V (·, 0), P∗ (X0 , X1 ) ∝ Pbase (X0 , X1 ) exp λr(X1 ) + V (X0 , 0) , so once we marginalize over X0 , the resulting tilt on X1 is no longer just exp(λr). The bias is largest when X0 and X1 are strongly coupled—exactly the case in flow matching, where the ODE makes X1 a deterministic function of X0 . To remove the bias, Domingo-Enrich et al. (2025) propose choosing σ during training so that the base process makes X0 and X1 independent. Under this memoryless schedule, V (X0 , 0) becomes a global constant and the endpoint marginal recovers the desired tilt p∗ (X1 ) ∝ pbase (X1 ) exp(λr(X1 )). For the linear interpolation, they show that the unique memoryless schedule is r 2(1 − t) σml (t) = , (24) t large at t → 0 (which destroys dependence on X0 ) and small at t → 1 (which keeps the endpoint clean). Importantly, Domingo-Enrich et al. (2025) further show that the memoryless schedule is the unique training schedule for which the fine-tuned optimum can be sampled at inference under any other noise schedule—in particular, the standard deterministic ODE (σ(t) = 0). Training with any non-memoryless schedule locks the fine-tuned model to that specific schedule at inference. A standard way of optimizing (23) is to build an estimator of its gradient in θ by sampling a trajectory {Xt } under uθ , computing the integrand, and backpropagating through Adjoint methods and adjoint matching.
45
the discretized SDE (22). Domingo-Enrich et al. (2025) show that the resulting gradient agrees in expectation with the gradient of the matching objective Z 1 2 (25) uθ (Xt , t) + σ(t)⊤ a(t) dt , LBasic−AM (θ) = 12 E 0
where {Xt } and {a(t)} are stop-gradient and a(t) is the adjoint state a(t) := ∇Xt
Z 1 1 2
t
∥uθ (Xs , s)∥2 ds − λ r(X1 ) ,
computed by a single backward solve of the ODE ⊤ 2 ȧ(t) = − ∇x (b + σ uθ )(Xt , t) a(t) − 21 ∇x uθ (Xt , t) ,
a(1) = −λ ∇r(X1 ).
(26)
In practice, Domingo-Enrich et al. (2025) find this empirically underperforms. They instead propose replacing a(t) with the lean adjoint ã(t), obtained by dropping the control-dependent terms of (26), ⊤ ˙ ã(t) = − ∇x bbase (Xt , t) ã(t),
ã(1) = −λ ∇r(X1 ).
(27)
The adjoint matching loss is the same regression as (25) but with ã in place of a, Z 1
LAM (θ) = 21 E
⊤
uθ (Xt , t) + σ(t) ã(t)
2
dt .
(28)
0
Domingo-Enrich et al. (2025) show that LAM has the same minimizer u∗ as LBasic−AM . Crucially, the lean adjoint is still anchored to the reward through its terminal condition ã(1) = −λ ∇r(X1 ), so the regression target remains informed by the reward landscape; dropping the control-dependent terms only reduces variance and yields significantly better empirical performance. Specialization to flow matching. For the linear interpolation, the base drift of (2) under σml from (24) is
bbase (x, t) = 2 vbase (x, t) − x/t. We parametrize the control uθ so that the controlled drift bbase + σml uθ has the same shape as bbase but with a learned velocity vθ in place of vbase . Solving for uθ yields uθ (x, t) =
2 vθ (x, t) − vbase (x, t) . σml (t)
(29)
The reason for this parametrization is practical: fine-tuning then reduces to learning a new velocity field vθ , which can be plugged directly into the standard deterministic ODE for sampling at inference time, with no need for the memoryless SDE. Substituting (29) into (22) gives the corresponding memoryless training SDE, r Xt 2(1 − t) dXt = 2 vθ (Xt , t) − dt + dBt . (30) t t Substituting (29) and bbase into (27) and (28) gives the lean adjoint and adjoint matching loss for flow models, h i⊤ ˙ ã(t) = − ∇x 2vbase (Xt , t) − Xt /t ã(t), ã(1) = −λ ∇r(X1 ), "Z # 2 1 2 FM LAM (θ) = 21 E vθ − vbase + σml (t) ã(t) dt . σml (t) 0
(31) (32)
The discretization of (30) and (31) introduces numerical stiffness at t → 0, which motivates the local linear integrator (section G). Implementation choices are collected in section H.
46
G
Local Linear Integrator for the Memoryless SDE JiT
RAE
600
REPA
SiT
600
kxt k2
500
80
80
60
60
500
400 400
300
40
40
200 0.0
0.5
1.0
0.0
t
0.5
1.0
0.0
t ODE Euler (target)
0.5
1.0
0.0
0.5
t
SDE Euler (with offset)
1.0
t
SDE Local linear (ours)
Figure 13 Latent norm trajectories under Euler ODE, Euler SDE, and local linear SDE sampling. Estimate of expected latent
norm E[∥Xt ∥] against time for 1000 trajectories sampled under the Euler ODE integrator, Euler SDE integrator with offsets, and the local linear SDE integrator. A perfect SDE integrator should match the ODE norm distribution. The panels show that the Euler SDE integrator contains an early time blow-up in norm, while our local linear SDE integrator stays close to the ODE reference across models. Red arrows mark the early-time Euler blow-up in each panel, while blue arrows point to the corresponding local linear trajectory at t = 0, where no such blow-up occurs.
This appendix describes the training-time integrator used in our experiments, explains the numerical issues that motivate it, and presents empirical results that validate its use. Although our experiments focus on adjoint matching, the same problems arise whenever an RL algorithm samples the memoryless SDE, and therefore apply to any path-based RL method on flow models, including non-adjoint-matching approaches such as Flow-GRPO (Liu et al., 2025a). Given its superior performance and simplicity, we believe this integrator should become the standard for all methods that require sampling from the memoryless SDE for flow models.
G.1
The Memoryless SDE and the Local Linear Integrator
Recall from section F that solving the KL-regularized RL problem for flow models requires sampling from the memoryless SDE during training, r 1 2(1 − t) dXt = 2 vθ (Xt , t) − Xt dt + dBt . (33) t t Using this memoryless SDE during training is what guarantees that the optimal terminal marginal is the tilted distribution p∗ (x) ∝ pbase (x) exp(r(x)), while still allowing us to switch back to the standard ODE at test time. The difficulty is that equation (33) is numerically stiff near t = 0: both the linear drift coefficient −1/t and p the diffusion coefficient 2(1 − t)/t diverge as t → 0+ . A naive explicit discretization therefore struggles on the first few steps. Euler–Maruyama with offsets. Domingo-Enrich et al. (2025) handle this p p stiffness by replacing the singular terms 1/t and 2(1 − t)/t with the regularized versions 1/(t + δ) and 2(1 − t + ε)/(t + δ), and then applying Euler–Maruyama to the modified SDE. With the practical choice δ = ε = ∆t, the update is Euler–Maruyama with δ = ε = ∆t (Domingo-Enrich et al., 2025):
xEuler k+1 = xk +
s xk 2(1 − tk + ∆t)∆t 2 vθ (xk , tk ) − ∆t + ξk , tk + ∆t tk + ∆t
47
ξk ∼ N (0, I).
(34)
Local linear
Euler
The shift δ moves the singularity away from t = 0, and ε provides slack near t = 1 in the adjoint matching loss, which DomingoEnrich et al. (2025) found to improve training speed.2
Figure 14 Matched-noise JiT samples at K=50. Top: Euler–Maruyama. Bottom:
local linear. Under the same Brownian motion, Euler with δ=ε=∆t already produces visibly degraded samples while local linear with the same offsets remains stable.
a pronounced early-time blow-up.
In practice, we observed that this discretization still produced lowquality samples for some models. The failure mode matches the intuition above: even with the offsets, an explicit first-order step cannot keep up with the rapid contraction of the stiff linear part at early times, so the noise increment dominates and trajectories blow up in norm, taking the model far from the training distribution. Figure 13 illustrates the effect across the four ImageNet models used in the paper. It plots the latent norm of 1000 trajectories under the Euler ODE, under Euler SDE with offsets, and under the local linear SDE integrator introduced next. A marginal-preserving SDE integrator should produce the same norm distribution as the ODE, but the Euler–Maruyama SDE discretization instead exhibits
The local linear integrator. We address the same stiffness with a different splitting. Our key observation is that
if we freeze the velocity vθ (Xt , t) to its left-endpoint value on each interval [tk , tk+1 ], the remaining dynamics form a linear SDE with time-dependent coefficients, which can be integrated in closed form via an integrating factor. Writing uk := tk + δ for the shifted time and specializing to δ = ε = ∆t as in the Euler–Maruyama baseline, this yields Local linear update with δ = ε = ∆t (ours):
xLL k+1 = Φk xk + Ωk vθ (xk , tk ) + Φk =
uk , uk+1
Ωk = uk+1 (1 − Φ2k ),
p
Vk ξk ,
(35)
ξk ∼ N (0, I),
Vk = (1 + 2∆t)(1 − Φ2k ) −
2 uk+1 − Φ2k uk . 3
A full derivation, including the general case δ, ε ≥ 0, is given in section G.2.
Empirical improvements. Replacing Euler–Maruyama with the local linear integrator improves both sampling and training. The clearest case is JiT, which was the most challenging model in our experiments: figure 14 shows that at K = 50 steps the local linear integrator already produces coherent images, while Euler–Maruyama with offsets produces visibly degraded samples under the same Brownian motion.
This can also be verified quantitatively, and the pattern persists across step counts and models. Figure 16 reports
−0.2
ImageReward
Compared to Euler–Maruyama, the local linear integrator has the same computational cost, namely one network evaluation per step, but it handles the stiff linear part of the SDE exactly. The contraction factor Φk = uk /uk+1 stays in (0, 1) throughout and is smallest near t = 0 (where the stiffness is largest), so it damps the early-time dynamics that would otherwise drive Euler– Maruyama trajectories to blow up. This is why local linear trajectories stay close to the ODE reference in figure 13. Moreover, no term has any problematic singularity.
−0.4 −0.6
SDE Euler (with offset) K=50 SDE Euler (with offset) K=100 SDE Local linear (ours) K=50 SDE Local linear (ours) K=100
−0.8 −1.0 0
500
1000 1500 2000 Training step
2500
Figure 15 JiT training reward during adjoint matching.
Local linear with δ=ε=∆t at K=50 closely tracks the K=100 run, while Euler with the same offsets and step count barely improves the reward. Both local linear runs outperform the Euler run at K=100, showing that the integrator matters for both optimization quality and speed.
2 The precise role of ε is orthogonal to the present discussion, but we keep it in the expressions for completeness.
48
3000
Fréchet distances under the Inception-v3, DINOv2-B, and DINOv3-L feature spaces for all four ImageNet models. Across model families, local linear with δ = ε = ∆t Pareto-dominates Euler–Maruyama with the same offsets, with the clearest gains at low step counts. The exception is RAE under Inception-v3, where both integrators produce high-quality samples and Euler-Maruyama with offsets slightly outperforms local linear. These improvements carry over to training as well. Figure 15 shows the training reward for JiT on ImageNet using the setup of the main reward-finetuning experiments. At K = 50, local linear matches the reward achieved by its K = 100 equivalent. On the other hand, Euler–Maruyama with offsets at K = 50 barely improves over the baseline, and even at K = 100 it still underperforms the local linear runs. This shows that the choice of integrator changes the effective optimization problem that RL sees, and that a better integrator can lead to faster training and better final performance. The 2× speedup is significant as sampling is the most expensive part of training. Inception-v3 150
DINOv2-B
SDE Euler (with offset)
JiT FD
SDE Local linear (ours)
100
DINOv3-L
1000 750
100
500 50 250 40
60
80
100
50 40
60
80
100
SiT FD
400 20
40
60
80
100
40
60
80
100
40
60
80
100
40
60
80
100
110 100
300
90 10
200 40
60
80
100
80 40
60
80
100
REPA FD
30
90 300
80
20
70
200 10
60
RAE FD
40
60
80
100
40
5
35.0
4
32.5
60
80
100
8
30.0
7
3 27.5 40
60
80
100
40
60
K
K
80
100
K
Figure 16 Cross-model endpoint FD for Euler vs. local linear integrators. Rows correspond to models and columns to the
three evaluation feature spaces. Local linear with δ=ε=∆t consistently matches or improves upon Euler with the same offsets, especially at low K. All metrics were computed with 50, 000 samples. Extension to the backward adjoint. The same idea applies to the backward adjoint ODE that appears in
adjoint matching. Defining u(t) = t + δ (so uk = u(tk ) matches the shifted time used above) and writing Φk = uk /uk+1 for brevity, recall that adjoint matching requires integrating ⊤ ⊤ at ȧt = − ∇x 2 vbase (Xt , t) − Xt /u(t) at = −2 ∇x vbase (Xt , t) at + , u(t)
(36)
where the Jacobian of the stiff linear drift −x/u(t) contributes the scalar −1/u(t) times the identity. (We write equation (36) in terms of u(t) rather than t, as in equation (31), given the practical choice of shifting the time to avoid the singularity at t = 0.) Domingo-Enrich et al. (2025) use a standard backward Euler step to discretize equation (36): 49
Standard backward Euler adjoint step (Domingo-Enrich et al., 2025):
ãEuler = ãk+1 + ∆t k
⊤ ãk+1 2 ∇x vbase (Xk+1 , tk+1 ) ãk+1 − . uk+1
(37)
Using a similar splitting idea, which we derive in section G.2 below, one can instead integrate the stiff scalar term exactly and replace this step by: Local linear backward adjoint step (ours):
⊤ ãLL = Φ ã + 2∆t ∇ v (X , t ) ã , k k+1 x base k+1 k+1 k+1 k
Φk =
uk . uk+1
(38)
In practice we did not see gains as large as for the forward integrator, but this step performed at least as well as the Euler discretization, so we kept it.
G.2
Derivation of the Local Linear Integrator
Both derivations rest on the following simple lemma (for further background in the context of ODEs see Arnold (1978)). Lemma G.1. Let f, g, h : [s, T ] → R be continuous and deterministic. The unique solution of the SDE dYt = f (t) Yt dt + g(t) dt + h(t) dBt , is RT
YT = e s f (ζ) dζ ys +
Z T
Ys = ys , Z T
RT
e τ f (ζ) dζ g(τ ) dτ +
s
The stochastic integral is a centered Gaussian with variance recovered as a special case.
RT s
e
RT
e τ f (ζ) dζ h(τ ) dBτ .
s RT 2 τ f (ζ) dζ
h(τ )2 dτ . The ODE case h ≡ 0 is
Rt Proof. Let µ(t) = exp − s f (ζ) dζ , so that µ(s) = 1 and µ̇(t) = −f (t) µ(t). Since µ is deterministic and of bounded variation, Itô’s product rule applied to µ(t) Yt gives d[µ(t)Yt ] = µ̇(t) Yt dt + µ(t) dYt = −f (t) µ(t) Yt dt + µ(t) f (t)Yt dt + g(t) dt + h(t) dBt . The two f (t) Yt terms cancel, leaving d[µ(t) Yt ] = µ(t) g(t) dt + µ(t) h(t) dBt . Integrating from s to T and using µ(s) Ys = ys , Z T µ(T ) YT = ys +
Z T µ(τ ) g(τ ) dτ +
s
µ(τ ) h(τ ) dBτ . s
Dividing through by µ(T ) gives YT =
1 ys + µ(T )
Z T
RT
s
µ(τ ) g(τ ) dτ + µ(T )
Z T s
µ(τ ) h(τ ) dBτ , µ(T )
RT and the identities 1/µ(T ) = exp s f (ζ) dζ and µ(τ )/µ(T ) = exp τ f (ζ) dζ yield the claimed formula. Gaussianity and the variance formula for the stochastic integral follow from Itô’s isometry applied to the deterministic integrand µ(τ ) h(τ )/µ(T ).
Throughout this section we keep the offsets δ, ε ≥ 0 and continue to use u(t) = t + δ,
uk = u(tk ), 50
Φk =
uk . uk+1
Forward SDE
The forward SDE
s Xt 2(1 − t + ε) dXt = 2 vθ (Xt , t) − dt + dBt u(t) u(t)
(39)
matches lemma G.1 with f (t) = −1/u(t), smooth forcing g(t) = 2 vθ (Xt , t), and noise coefficient h(t) = p 2(1 − t + ε)/u(t). The exponential factor in the lemma simplifies because −1/u has a logarithmic antiderivative: Rt u(s) . e s f (ζ) dζ = e−[log u(t)−log u(s)] = u(t) Applying lemma G.1 over a single step [tk , tk+1 ] gives the exact representation s Z tk+1 Z tk+1 2 1 2(1 − τ + ε) Xtk+1 = Φk Xtk + u(τ ) vθ (Xτ , τ ) dτ + u(τ ) dBτ . uk+1 tk uk+1 tk u(τ )
(40)
We make a single approximation—freezing the velocity at its left-endpoint value, vθ (Xτ , τ ) ≈ vθ (xk , tk ) =: vk — and integrate the rest exactly. The stochastic-integral clause of lemma G.1 turns the noise integral into a centered Gaussian. The velocity term collapses to Ωk vk with Z tk+1 u2 − u2k 2 u(τ ) dτ = k+1 = uk+1 1 − Φ2k . Ωk = uk+1 tk uk+1 √ The noise term is Vk ξk with ξk ∼ N (0, I) and Z tk+1 2 Vk = 2 u(τ ) (1 − τ + ε) dτ. uk+1 tk
(41)
Using the algebraic identity 1 − τ + ε = (1 + δ + ε) − u(τ ), the integrand splits into elementary terms: Z tk+1 Z tk+1 u3 − u3k u2 − u2k , . u(τ )2 dτ = k+1 u(τ ) dτ = k+1 2 3 tk tk Substituting and simplifying with u3k /u2k+1 = Φ2k uk , Vk = (1 + δ + ε) 1 − Φ2k
−
2 uk+1 − Φ2k uk . 3
(42)
Setting δ = ε = ∆t recovers the boxed forward update equation (35). Backward adjoint ODE
We now derive the boxed local linear backward adjoint step equation (38). Recall the split form of the adjoint ODE equation (36), ⊤ at . ȧt = −2 ∇x vbase (Xt , t) at + u(t) To handle the stiff scalar term exactly, apply lemma G.1 to this ODE—now with f (t) = 1/u(t) and smooth forcing g(t) = −2[∇x vbase (Xt , t)]⊤ at . The exponential factor is again a ratio of u’s, with the opposite sign from the forward case: Rt u(t) e s f (ζ) dζ = . u(s) Applying the lemma over [tk , tk+1 ] to express atk+1 in terms of atk and rearranging for the backward direction gives the exact identity Z tk+1 ⊤ uk ãk = Φk ãk+1 + 2 ∇x vbase (Xτ , τ ) aτ dτ. u(τ ) tk 51
The integrand depends on the unknown aτ , so we must approximate. Following the same spirit as in the forward case, we freeze the integrand at a single point and use the resulting constant times ∆t. We freeze at τ = tk+1 , the only point where aτ is known: this replaces aτ by ãk+1 , the Jacobian by ∇x vbase (Xk+1 , tk+1 ), and the kernel uk /u(τ ) by uk /uk+1 = Φk . The integral collapses to 2Φk ∆t times the frozen value, yielding ⊤ ãk = Φk ãk+1 + 2∆t ∇x vbase (Xk+1 , tk+1 ) ãk+1 , which is equation (38).
52
H
Additional Implementation Details
This appendix gives the method-level implementation recipe used in our DRL experiments. The end-to-end procedure is summarized in algorithm 2, which combines the buffered training pattern used in practice for both stages with the local linear integrator for the memoryless SDE and its companion adjoint (section G). We provide further details below. Experiment-level configuration (base models, default hyperparameters, sampling, evaluation metrics, baseline and ablation setups) is collected in section I. The shifted-time offsets in Stage 2 are written as δ, ε ≥ 0 throughout the algorithm; in practice we use the choice δ = ε = ∆t from section G, but stating the algorithm in terms of δ, ε keeps the dependence on the offsets explicit. Buffered discriminator training. To amortize the cost of generating model samples, we do not draw a fresh
mini-batch of fakes for every gradient step. Instead, we periodically refresh a buffer BD of size ND holding paired real samples xreal and base-model samples xfake , and then take ED epochs of gradient steps over that buffer with the standard logistic loss plus an R1 gradient penalty. The R1 penalty is computed by backpropagating the discriminator output through the frozen encoder ϕ to the image space. In principle, the flow-space could be used as well but we found this to work well. Given that it is also more memory efficient as we avoid backpropagating through the flow space, we decided to use it instead of the flow-space penalty. Buffered adjoint matching. We use the same buffering pattern for the RL stage. Periodically, we sample NT
trajectories from the memoryless training SDE under the current vθ using the local linear forward integrator (35), compute their adjoints by solving the local linear backward adjoint (38) starting from ãK = −λ ∇x r̂(XK ), and store both quantities in a trajectory buffer BT as stop-gradient tensors. We then take ET epochs of gradient steps over BT , with gradients flowing only through the current velocity evaluations vθ (Xk , tk ). This is technically off-policy (the buffer drifts as vθ updates), but in practice we observed no degradation versus a fully on-policy implementation, and Girsanov-based importance weighting gave no measurable improvement either. Training versus inference. The memoryless SDE used during post-training is integrated with the local linear
scheme described in section G. After fine-tuning, we sample from the updated model with the standard deterministic ODE; differences between the two samplers were small but non-negligible.
Other implementation details. A few smaller choices, all following or adapted from the original adjoint-matching
implementation:
1. Final denoising step. The forward loop in algorithm 2 runs K − 1 local linear SDE steps and replaces the final step from tK−1 to tK with a single deterministic denoising call to the network. 2. Stratified timestep subsampling for the loss. Rather than summing the AM loss over all K timesteps, each gradient step uses a fresh subset S ⊂ {0, . . . , K − 1} of size ⌈f K⌉, with f = 0.4 by default. The subset is split 50/50 between early (t < 0.6) and late (t ≥ 0.6) timesteps so that bothP ends of the trajectory are always represented. The loss line in Stage 2 then becomes LAM = (1/(2|S|)) k∈S ∥ · ∥2 . 2
(i)
3. Adaptive outlier clipping for the AM loss. The per-sample loss ℓk = σ2k (vθ − vbase ) + σk ãk has a heavy right tail, dominated by a few trajectories whose adjoints blow up and destabilize training. We maintain q an EMA-smoothed threshold τn across optimizer steps n on the per-sample norms τn = ρ τn−1 + (1 − ρ) min qα
q
(i)
ℓk ,
(i) ℓk , c τn−1 ,
where qα is the empirical α-quantile over the current batch, ρ is the EMA decay, andqthe spike cap (i)
c prevents a single outlier batch from rapidly inflating the threshold. Samples with ℓk ≥ τn (or non-finite) are masked out of the gradient. We use α = 0.9, ρ = 0.9 for SiT-XL/2, REPA SiT-XL/2, and JiT-H/16, and α = 0.75, ρ = 0.995 for RAE DiTDH-XL, with spike cap c = 5 throughout. The spike cap c was particularly important for RAE, whose loss exhibited noticeably more frequent spikes than the other models.
53
Algorithm 2 Discriminator-Guided RL with Buffered Training and the Local Linear Integrator Stage 1: Buffered Feature-Space Discriminator Training Require: Target distribution q, base model pbase , frozen encoder ϕ, buffer size ND , epochs per refresh ED , R1 weight γ, learning
rate ηD 1: Initialize discriminator parameters ψ 2: repeat 3: Refresh buffer BD : 4: for i = 1, . . . , ND do (i) (i) 5: Sample xreal ∼ q and xfake ∼ pbase (i) (i) 6: Append xreal , xfake to BD 7: end for 8: 9: Train discriminator on BD : 10: for epoch n = 1, . . . , ED do 11: for minibatch (xreal , xfake ) ⊂ BD do 2 12: ℓdisc (ψ) ← − log Dψ (ϕ(xreal )) − log 1 − Dψ (ϕ(xfake )) + γ2 ∇x Dψ (ϕ(xreal )) 13: ψ ← ψ − ηD ∇ψ ℓdisc (ψ) 14: end for 15: end for 16: until convergence 17: Define reward r̂(x) ← logit Dψ (ϕ(x))
Stage 2: Buffered Adjoint Matching with the Local Linear Integrator
Require: Reward r̂, base velocity vbase , KL weight λ, uniform time grid 0 = t0 < · · · < tK = 1 with ∆t = 1/K, shifted-time
offsets δ, ε ≥ 0 with uk := tk + δ, trajectory buffer size NT , epochs per refresh ET , learning rate ηT 1: Initialize vθ ← vbase 2: repeat 3: Refresh trajectory buffer BT : 4: for i = 1, . . . , NT do (i) 5: Sample X0 ∼ N (0, I) 6: Local linear forward pass (with current vθ ): 7: for k = 0, . . . , K − 1 do 8: Φk ← uk /uk+1 , Ωk ← uk+1 1 − Φ2k , Vk ← (1 + δ + ε) 1 − Φ2k − 23 uk+1 − Φ2k uk
9:
10: 11: 12: 13: 14: 15: 16: 17: 18: 19: 20: 21: 22: 23: 24:
(i)
ξk ∼ N (0, I) (i)
(i)
(i)
Xk+1 ← Φk Xk + Ωk vθ (Xk , tk ) +
end for
√
(i)
Vk ξk
(i) (i) Terminal adjoint ãK ← −λ ∇x r̂ XK Local linear backward adjoint pass (with vbase ): for k = K − 1, . . . , 0 do ⊤ (i) (i) (i) (i) ãk ← Φk ãk+1 + 2∆t ∇x vbase Xk+1 , tk+1 ãk+1
end for
(i) (i) K Store {Xk }K k=0 , {ãk }k=0 in BT as stop-gradient quantities
end for
Train velocity on BT : for epoch n = 1, . . . , ET do for minibatch {Xk }, {ãk } ⊂ BT do
p 2(1 − tk + ε)/uk 2 1 PK−1 2 LAM (θ) ← 2K vθ (Xk , tk ) − vbase (Xk , tk ) + σk ãk k=0 σ
Set σk ←
k
25: θ ← θ − ηT ∇θ LAM (θ) 26: end for 27: end for 28: until convergence 29: return vθ
54
I
Experimental Setup
We collect all experiment-level configuration here: shared setup (base models, sampling, evaluation metrics) followed by per-experiment subsections covering our main distributional alignment experiments (section 5), the baselines we compare against, and the ablations discussed in design choices. The general method recipe is given separately in section H. Base models. We evaluate four pretrained ImageNet (Deng et al., 2009) 256 × 256, class-conditional generators
spanning latent and pixel-space flow/diffusion architectures:
• SiT-XL/2 (Ma et al., 2024) (0.72B params total: 675M flow model, 50M decoder): A generative model built on the DiT backbone (Peebles and Xie, 2023) using the stochastic interpolant framework (Albergo and Vanden-Eijnden, 2023). We use the XL/2 variant, operating in the standard SD VAE latent space (4 × 32 × 32).
• JiT-H/16 (Li and He, 2025) (0.95B params total: 953M flow model, 0 decoder): A plain Vision Transformer (Dosovitskiy et al., 2021) that performs flow matching directly in pixel space on raw image patches, without a VAE tokenizer. We use the H/16 variant operating directly on 3×256×256 pixel inputs. • REPA SiT-XL/2 (Yu et al., 2024) (0.72B params total: 675M flow model, 50M decoder): A flow-based diffusion transformer regularized by aligning intermediate noisy representations with clean semantic features from a frozen encoder. We use the variant built on the SiT-XL/2 backbone (same SD VAE latent space, 4 × 32 × 32), trained with REPA alignment against DINOv2-ViT-B features.
• RAE DiTDH-XL (Zheng et al., 2025) (1.25B params total: 839M flow model, 415M ViT-XL decoder): A diffusion transformer that replaces the VAE latent space with a Representation Autoencoder—a frozen vision encoder paired with a lightweight decoder—and trains via flow matching in this semantic latent space. We use the DiTDH-XL stage-2 model with a frozen DINOv2-ViT-B encoder and ViT-XL decoder, operating in a 768 × 16 × 16 semantic latent space.
Sampling. For all models, unless otherwise indicated, we use a Heun 50-step sampler with a linear schedule,
except for RAE where we follow the original paper and use their suggested scheduler (Zheng et al., 2025). This corresponds to 99 NFEs since the final step reduces to a regular Euler step. Evaluation metrics and feature spaces. For our main results, we report metrics in the DINOv2-Large, DINOv3-
Large, SigLIP-Large, and Inception-v3 feature spaces. The feature-space ablation also reports the corresponding base-size variants (DINOv2-B, DINOv3-B, SigLIP-B); patterns are largely similar to their large-size counterparts. All distributional metrics are computed on 50 000 generated samples with class labels balanced across the 1000 ImageNet classes, against an equally sized reference set so that the generated and reference sets are matched in size. As is commonly done in the literature, FD is reported against the (µ, Σ) feature statistics of the full ImageNet training set; FDval , KD, Precision, Recall, Density, and Coverage are reported against 50 000 ImageNet validation images. Precision/Recall are computed with k=3 nearest neighbors and Density/Coverage with k=5, matching the defaults of Kynkäänniemi et al. (2019) and Naeem et al. (2020) respectively. All manifold metrics use Euclidean distance directly on the embedding vectors. Whenever we report “best CFG” numbers, the configuration is selected per checkpoint by sweeping the CFG scale from 1 to 3 in steps of 0.25 together with the application interval {(0.0, 1.0), (0.1, 1.0), (0.3, 1.0)} (the time range t ∈ [low, high] over which CFG is applied; outside it the guidance scale falls back to 1.0) (Kynkäänniemi et al., 2024). For the autoguidance experiments (section J) we also sweep scales from 0 to 1 in the same step. Best-CFG configuration.
I.1
Distributional alignment setup
DRL is run as two stages—a discriminator on frozen-encoder features, followed by KL-regularized RL via adjoint matching—both using buffered training (algorithm 2). The same checkpoints back the distributional alignment results and the image-quality transfer results. Discriminator architecture. Unless stated otherwise, the discriminator operates on frozen DINOv2-Large
features (CLS token; the encoder is never updated during DRL). The head is a class-conditional projection 55
discriminator (Miyato and Koyama, 2018) of the form D(x, y) = φ(h(x)) + ⟨h(x), e(y)⟩ + b(y), where h is a linear map (no hidden layers) with hidden dim 512 and e(y) is a learned class embedding. The reward used by the RL stage is the discriminator logit, r̂(x) = logit D(ϕ(x)), as constructed in section 4. Discriminator training. We optimize the standard logistic GAN loss with the R1 gradient penalty above. Each
refresh of the discriminator buffer holds 20 mini-batches of paired (real, fake) images (in the ambient space of the model), and we take 10 epochs of gradient steps over that buffer before refreshing. We optimize with Adam at learning rate 10−4 at an effective batch size of 512, for up to 10 000 optimizer steps, evaluating on a held-out validation set every 50 steps and early-stopping with patience 15. RL training (adjoint matching). We optimize the RL objective with adjoint matching (Domingo-Enrich et al., 2025), sampling trajectories from the memoryless training SDE using the local linear integrator (section G) at K = 100 uniform steps with η = 1 and offsets δ = ε = ∆t. CFG is disabled during training. We use Adam at learning rate 10−5 (no weight decay) with a 200-step linear warmup, run for 3 000 optimizer steps, and sum the loss across timesteps. In addition to the per-trajectory adaptive clipping and stratified timestep subsampling described in section H, trajectories and their adjoints are pre-computed and stored stop-gradient in a buffer; we then take 5 epochs of velocity gradient steps over that buffer per refresh. We use an effective batch size of 576 for all models, with a buffer of 2304 trajectories for SiT-XL/2 and REPA SiT-XL/2 and a buffer of 3072 trajectories for JiT-H/16 and RAE DiTDH-XL. We observed that RAE was noticeably more unstable to train than the other models, so for RAE we increased the quantile clipping (q = 0.75 with EMA decay 0.995, vs q = 0.9 and EMA decay 0.9 elsewhere) and reduced the loss to its mean rather than its sum across timesteps to compensate for the larger latent space (768 × 16 × 16 vs 4 × 32 × 32). We also use a slightly higher learning rate of 5 × 10−5 (vs 10−5 elsewhere), which we found to work better in practice for RAE, together with gradient clipping of 1.0 and a spike cap of 5 on the clipped loss-threshold EMA. Compute. Each base model requires a one-off discriminator training run plus one adjoint-matching run per λ. All numbers below are measured on a cluster of NVIDIA H200 GPUs (8 per node). The discriminator stage takes roughly 3–6 hours of wall-clock time on a single 8×H200 node, i.e. approximately ∼24–48 GPU-hours per base model. The per-λ adjoint-matching cost is approximately ∼190 GPU-hours for SiT-XL/2 and REPA SiT-XL/2 (∼110 in optimizer steps and ∼80 in buffer-fill sampling), ∼500 GPU-hours for JiT-H/16 (∼290 in optimizer steps, ∼210 in buffer-fill sampling, with the higher cost reflecting pixel-space ViT-H/16 forward passes), and ∼420 GPU-hours for RAE DiTDH-XL (∼190 in optimizer steps, ∼230 in buffer-fill sampling, with the buffer-fill share being largest because sampling in the 768 × 16 × 16 semantic latent is more expensive than in the SD VAE latent).
For comparison, the base models we post-train are themselves the result of substantial pretraining budgets: SiT-XL/2 is reported to be trained for approximately 18,300 TPU v4 chip-hours (Ma et al., 2024), and REPA SiT-XL/2 for 1,646 H100 GPU-hours (Yu et al., 2024) (to our knowledge no GPU-hour figures have been reported for JiT or RAE). DRL post-training is therefore a small fraction of the cost of producing the base model in either case. We also note that we did not spend significant effort tuning DRL for training-time efficiency: we believe it is straightforward to roughly halve the per-λ cost by trading off the configuration knobs we held fixed—fewer discretization steps, larger trajectory buffers, more epochs per buffer refresh, and so on—without changing the qualitative results.
I.2
Better image quality (reward-transfer evaluation) setup
The image-quality transfer results reuse the distributional alignment checkpoints (section I.1) without any additional training, and only the evaluation protocol is new. We evaluate four held-out preference reward models—ImageReward (Xu et al., 2023), PickScore (Kirstain et al., 2023), Aesthetics v2.5 (discus0434, 2024), and HPSv2 (Wu et al., 2023)—over 50 000 generated images per checkpoint (balanced across the 1 000 ImageNet classes). For the text-conditioned rewards, the ImageNet class label is plugged into the prompt template “a photo of a {class}”.
I.3
Preference-based RL setup
For the preference alignment results, we post-train each base model with KL-regularized RL using ImageReward (Xu et al., 2023) as the reward (computed on the rendered image with the same “a photo of a {class}” 56
prompt template as in section I.2), and compare two starting points: (i) the base model (Base+PRL) and (ii) our DRL λ=10 checkpoint (DRL+PRL). For each starting point we sweep the PRL reward–KL trade-off λPRL ∈ {1, 10, 40}. For DRL+PRL the DRL checkpoint serves as both the initial weights and the KL reference, so the KL is anchored to the distributionally aligned model rather than the base. RL training (adjoint matching). PRL uses the same adjoint matching setup as the distributional alignment
stage (section I.1—same batch size, buffer sizes, and RAE-specific overrides), trained for 5 000 optimizer steps instead of 3 000. We chose this larger budget after observing that the Base+PRL runs at 3 000 steps were close to but not yet fully converged. Compute. Same per-λ figures as section I.1, scaled proportionally to the longer 5 000-step budget (a factor of
5/3). Evaluation for the Pareto / over-optimization analysis. Beyond the same distributional and reward metrics from
section I.1, the analysis behind figure 9 (and the appendix Pareto figures) tracks five low-level image statistics that reward-hacking can push to extreme values: mean brightness, saturation, contrast, colorfulness, and whiteness. Each statistic is computed per image on RGB pixels in [0, 1] and then averaged across all images: • Brightness: mean grayscale luminance Y = 0.299 R + 0.587 G + 0.114 B over all pixels (BT.601 weights). • Saturation: mean of (maxc − minc )/ maxc across pixels (channel-wise max/min over RGB; clamped to avoid division by zero). • Contrast: pixel-wise standard deviation of the grayscale luminance Y . q q 2 + σ 2 and µ = µ2rg + µ2yb on the • Colorfulness: the Hasler–Süsstrunk metric σ + 0.3 µ, with σ = σrg yb opponent channels rg = R − G and yb = 12 (R + G) − B.
• Whiteness: mean of minc (R, G, B) across pixels (high when pixels are close to white). Statistics are computed on 3 000 generated images per checkpoint (three images per ImageNet class, sampled with fixed noise seeds shared across checkpoints so trajectories are directly comparable). The ImageNet reference values are computed on 3 000 ImageNet validation images (three per class, resized and 256 × 256 center-cropped).
I.4
Distillation from RL teachers setup
The distillation-gap experiment (section 3) on Stable Diffusion 1.5 has two stages: (i) train an RL teacher with adjoint matching, and (ii) distill its samples into a fresh student with standard flow / score matching. RL teacher (adjoint matching). We use the official adjoint-matching codebase released by Domingo-Enrich
et al. (2025) (https://github.com/microsoft/soc-fine-tuning-sd). Starting from Stable Diffusion 1.5, we run buffered adjoint matching with the codebase’s multi_prompt_buffer.yaml settings, and only deviate as noted below. Training is on an ∼8 000-prompt subset of the 10 000 prompts shipped in the codebase, themselves extracted from the ReFL training data of Xu et al. (2023), with ImageReward as the reward and reward multiplier 300 (vs 100 in the default config). Sampling uses 50-step memoryless DDIM (η = 1). Optimization uses Adam at learning rate 3×10−6 , β = (0.9, 0.95), fp32; per-rank batch 6 with 11 accumulation steps on 8 H200 GPUs (effective batch 528). Buffered training uses buffer_size=100 trajectories with passes_per_buffer=10. We train for 10 epochs over the prompt set and keep the checkpoint with the best validation ImageReward as the teacher. These are essentially the default settings provided in multi_prompt_buffer.yaml. Distillation (SFT student). The student is a fresh Stable Diffusion 1.5 UNet trained with the standard DDPM
ε-prediction loss on samples drawn from the RL teacher. We use the same prompt set used to train the RL teacher above. Teacher samples are generated without classifier-free guidance using 50-step DDIM with η = 0 on the standard SD 1.5 noise schedule, which is also the schedule the student trains on (βstart = 8.5 × 10−4 , βend = 1.2 × 10−2 , scaled-linear). Training is buffered: every cycle we generate a fresh batch of 4 096 teacher images and take 10 passes over the buffer, with per-rank batch 32 and 2 accumulation steps on 8 H200 GPUs (effective batch 512). Optimization uses Adam at learning rate 10−5 with 100-step linear warmup, gradient clipping at 1.0, and bfloat16 mixed precision. We evaluate ImageReward on a held-out validation split every
57
1 000 optimizer steps and early-stop when no improvement is observed for 5 consecutive eval rounds. The run plotted in figure 4 stopped after ∼ 9 000 optimizer steps under this criterion—roughly 50× more (prompt, teacher-image) pairs seen by the gradient than the RL teacher saw during its own training. The picture is unchanged if we evaluate on training prompts rather than the held-out validation split. We report and train at CFG=1 (no guidance) as the “clean” comparison: the RL teacher is never fine-tuned along the non-conditional branch (adjoint matching only works well without CFG). This also avoids issues given that the CFG distribution of the teacher would not be the CFG distribution of the student if we were to fine-tune the student on the CFG teacher samples. This is also the setting we discuss in the main text.
I.5
Distillation from DRL Teachers
For the DRL distillation experiment, we use the same DRL model (λ=1, R1 =0, linear-conditional discriminator on DINOv2-L) used throughout the main paper as the frozen teacher. At a high level, the setup mirrors the distillation experiment on Stable Diffusion 1.5 in section 3: we sample from the post-DRL model to fill a buffer, train the student on those samples with the standard flow matching velocity loss, and repeat. Unlike the Stable Diffusion 1.5 experiments, we do not apply early stopping to drive home the point that the gap is not solvable with more compute. Hyperparameters. The student is initialized from the base REPA SiT-XL/2 pretrained checkpoint and
trained with the standard flow matching velocity loss on teacher-generated samples (no REPA projection loss). The teacher generates samples with Heun-50. We use AdamW with lr = 10−4 (constant schedule, 100-step warmup), no weight decay, gradient clipping at 1.0, and EMA decay of 0.9999. Training runs in full precision (float32), as we occasionally encountered instabilities with bf16 accumulation. The effective batch size is 576. Teacher-generated samples are stored in a buffer of 57,600 samples; the student trains on this buffer for 10 passes (1,000 optimizer steps) before the buffer is discarded and refilled with fresh teacher generations. Evaluation settings follow those described in the main paper. The results reported in the main text are at CFG=1 and with EMA. Non-EMA results are similar, albeit slightly worse. Additionally, these hyperparameters correspond almost exactly to the original REPA training configuration, with the only differences being a larger batch size and the absence of the REPA projection loss. We tried other learning rates and buffer-refresh schedules, but found similar or worse results. Training samples. We train for 900k optimizer steps. Since the buffer is refreshed every 1,000 steps, this
amounts to 900 buffer fills, each producing 57,600 fresh samples. The student therefore sees a total of ∼52M teacher samples over the course of training—over 150× the number of samples consumed during DRL post-training (∼346k) and ∼40× the size of ImageNet (∼1.28M images). Compute. The total cost of the 900k-step run is approximately 4,050 GPU-hours on H200 GPUs. Buffer
filling dominates: each fill requires a 50-step Heun solve for 57,600 samples and takes ∼400 seconds, so 900 fills account for ∼2,385 GPU-hours (∼59% of the budget). The remaining ∼1,710 GPU-hours cover the optimization steps.
I.6
Feature-space ablation setup
We base our feature-space ablation on REPA SiT-XL/2. We train separate linear conditional discriminators on features from seven embedders—DINOv2 (base, large), DINOv3 (base, large), SigLIP (base, large), and Inception-v3—and use each as the reward for adjoint matching. Only the embedder varies across configs; all other discriminator and adjoint-matching settings follow the distributional alignment setup (section I.1). Every fine-tuned model is then evaluated across all seven feature spaces; full results are in table 5.
I.7
Discriminator architecture and training ablation setup
The discriminator ablation on REPA SiT-XL/2 with λ = 10 follows the distributional alignment setup (section I.1) with the following differences and additions. Discriminator. We explore three additional configurations on top of the default linear projection head on frozen
DINOv2-Large: (i) MLP head —two GELU residual blocks of hidden dim 512 on top of the frozen DINOv2Large CLS token, followed by the same linear map used in the default; (ii) fine-tuning the embedder —unfreeze 58
DINOv2-Large and train it jointly with the linear head; (iii) training the embedder from scratch—initialize the same DINOv2-L architecture randomly and train it jointly with the linear head. In all three configurations we keep the class-conditional projection wrapper D(x, y) = φ(h(x)) + ⟨h(x), e(y)⟩ + b(y) from the alignment setup. Training. Everything else matches the alignment setup, except for the embedder learning rate when the
embedder is updated. We use 10−5 for fine-tuning—we initially tried 10−4 but training was too unstable and the resulting discriminators were noticeably worse. We nevertheless kept 10−4 for the from-scratch variant. At these rates we observed no training instabilities and the discriminator reached the same convergence behavior as the frozen-embedder default.
59
J
Extended Results
This section reports the extended quantitative results that complement the main-text experiments. Subsections are ordered to match the main-text experiment flow: alignment (distributional alignment), image quality (image-quality transfer), preference RL (preference alignment), and ablations (design choices). Setups for each experiment, including the ablation training configurations, are collected in section I. A note on autoguidance. Standard classifier-free guidance (CFG) replaces the conditional velocity field v(x, y) with a combination of the conditional and unconditional fields, vcfg (x, y) = v(x, ∅) + w (v(x, y) − v(x, ∅)), where w is a positive scalar, usually greater than 1, and ∅ denotes the null class. Autoguidance (Karras et al., 2024) proposes replacing v(x, ∅) with the velocity field of a weaker version of the same model, typically an earlier checkpoint from pretraining. In practice, it has been shown to improve sample quality and diversity.
To our knowledge, autoguidance has only been applied to pretrained models. As a curiosity, we explored applying it in the RL setting, guiding the RL fine-tuned conditional model with the corresponding non-RL conditional version, vag (x, y) = vbase (x, y) + w (vRL (x, y) − vbase (x, y)). (43)
We tried this in two settings: using the non-DRL model to guide the DRL model, and, after preference alignment, using the non-preference-aligned model to guide the preference-aligned model. Generally, we found the results to be mixed. In the DRL setting, when λ was set to a high value (e.g., 40), a value of w lower than 1 usually improved performance over both the base and the DRL model. However, autoguidance generally underperformed standard CFG with an interval. In the preference-alignment setting, w > 1 generally led to the highest reward values across all models, but also to the worst image fidelity: with sufficiently large w, images became very bright and white, which is presumably why they scored so highly under the proxy reward despite being poorly aligned with true human preferences. For these reasons, we do not feature autoguidance in the main text. Nevertheless, we provide Pareto fronts with autoguidance in figures 20 to 22 and the corresponding sample grids in sections K.2 and K.3, so that interested readers can see the effect in more detail and so that the observations may prove useful for future work on the topic. Alignment. figures 17 and 18 show how all seven distributional metrics vary with λ at CFG = 1 (no guidance).
For each (model, embedder, λ) combination, the plotted value corresponds to the interval achieving the lowest FD, with all other metrics taken from the same entry. The dashed gray line marks the base model value, and the dotted line (in the model’s color) marks the theoretically motivated λ=1 (R1 =0) checkpoint. tables 2 and 3 extend the Fréchet Distance comparison from figure 5 to the full set of distributional metrics: FD, FDval , KD, Precision, Recall, Density, and Coverage. For each (model, embedder) pair, we select the DRL configuration (λ, interval) that achieves the lowest FD and report all seven metrics from that same generation run, along with the chosen configuration. This ensures internal consistency across all numbers in each row.
The takeaways are largely unchanged from the main text: DRL consistently improves over the base model across the board. Two patterns become clearer when seeing the full sweep of metrics. First, λ ∈ {5, 10} generally performs best across most metrics, with the exception of FID, which tends to prefer smaller λ. Second, as discussed in the main paper, larger λ generally brings improvements over the theoretically motivated baseline. Additionally, λ=40 is never selected as the best configuration and λ=20 is selected only sparingly, which further validates our choice of λ=10 as the default for the reward-improvement and preference-RL experiments. Finally, while we select the best configuration according to its FD value (computed against the ImageNet training distribution), the improvements persist on FDval and KD, neither of which we select for. Together, these suggest genuine distributional alignment rather than overfitting to the optimized signal. Image quality. Figure 19 extends the reward-improvement comparison from figure 7, which fixes λ=10, to the
full λ sweep. Each panel plots normalized reward improvement (DRL−Base)/σbase against λ for one held-out reward. The dotted horizontal line per model marks the theoretically motivated λ=1 (R1 =0) checkpoint for reference. Improvements are largely monotone in λ up to λ=20, with the strongest gains in ImageReward, HPSv2, and PickScore as reported in the main text. For λ=40, however, improvements often dip—consistent
60
with the alignment results in figures 17 and 18, where λ=40 is rarely the best configuration. Table 4 reports the raw reward values underlying figure 7: λ=10. Preference RL: Pareto plots. figure 9 in the main text shows HPSv2 vs. low-level image statistics at CFG = 1 (no guidance). figures 20 to 22 extend this to a full 3 × 3 grid: three held-out rewards (HPSv2, ImageReward, Aesthetic v2.5) × three guidance settings (no CFG, CFG = 2, autoguidance = 2; see section J for context on autoguidance). We do not sweep over the guidance value because the results are easier to read this way; the plots look essentially identical for other CFG values.
The results are essentially unchanged from the main text: the patterns we showed for HPSv2 also hold for ImageReward and Aesthetic v2.5, with DRL achieving the best Pareto fronts in nearly every plot. Moreover, this conclusion is robust across the three guidance settings. Furthermore, as discussed in section J, autoguidance pushes reward values higher at the cost of substantially worse image-statistic distortion—for example, on REPA the maximum brightness rises from ∼0.6 under base+CFG to ∼0.7 under autoguidance, an effect also visible in the sample grids of section K.3, where the images are noticeably brighter. This further supports our claim that reward proxies are usually exploitable and negatively entangled with image fidelity. Even so, DRL retains the best Pareto front in the autoguidance setting. Ablations. Figure 23 reports DINOv2-L FD across the full R1 sweep at both λ=1 and λ=10, complementing
the path-based view in Figure 11. The two regimes show qualitatively different behavior: at λ=1 all R1 >0 values hurt — the cleanest baseline is at R1 =0 — while at λ=10 a small R1 rescues frozen-feature heads from collapse. This supports the claim in design choices that R1 here functions as a stabilizer for aggressive λ rather than as a vanishing-gradient remedy as in the standard GAN literature (Mescheder et al., 2018). Table 5 reports the full 7 × 7 feature-space ablation summarized by table 1 in the main text: FD and KD when training a discriminator on each of seven embedders (columns) and evaluating on each of seven embedders (rows). The training setup is identical across embedders (section I.6); only the embedder varies. We see that even with the additional embedders the conclusions are unchanged. CFG effect. Finally, to provide additional insight into how our method interacts with CFG, Figure 24 and
Figure 25 break down the per-reward and per-embedder FD performance of DRL as a function of the CFG scale. The reward plot is constructed by running CFG and using the same setup for estimating rewards described earlier. The FD plot is constructed by taking the minimum over the three CFG interval values described in section I at each CFG scale. We didn’t notice meaningful differences across the three intervals for the rewards. Several observations are in order. For the reward values, we see that DRL has a strong positive effect that is maintained at all CFG levels. For some (model, reward) pairs (e.g., SiT on ImageReward) DRL with no CFG already attains higher reward than the base model at any CFG value, and the trend is consistent across model classes. The only exception is RAE, where a small λ slightly decreases the reward; this disappears at larger λ. For FD the picture is more mixed, with two evident patterns. First, on Inception and SigLIP smaller λ generally yields better performance, while on DINOv2 and DINOv3 larger λ is generally more beneficial. Second, CFG becomes much less effective the larger λ gets. We believe these are artifacts of the current training setup, in which we only finetune the conditional branch (a constraint imposed by adjoint matching). Finding ways to make the model less dependent on λ and trainable jointly with both the CFG and non-CFG branches is an interesting direction for future work.
61
J.1
Alignment: Quantitative λ Sweep Base
FD ↓
75
FDval ↓
30
KD ×1000 ↓
.76
Precision ↑
80
50
20
60 40
10
10
12
5
.78
7.5
9
4
.76
3
.74
5
.68
27
.848
24
.84
16 12.5
14
10
12
6
7.5
4 2
3 1
5 102040
1
λ
.725
.56
1
5 102040
1
160
160
KD ×1000 ↓
1
DRL (λ=1, R1 =0)
.5
.75
.5
.4
.5
DINOv3-L SigLIP-L
30
30
60
.76
20
20
30
.72
.2
15
18
7.5
.8
.6
10
12
5
.7
.45
.6
.3
λ
.8
1
5 102040
λ
1
5 102040
Density ↑
.75
Coverage ↑
.6
1
.75
.9
.6
.8
.4
1.2
.3
1 .8
.84 .78 .72 .66
1.5 .9
1.2
2.5
6 5 102040
90
λ
.45
.84
40
1
.25
.5
25
5
.3
.6
40
5 102040
.45
.35 .7
1
DRL (λ sweep, R1 =10−5 )
1
.6
40
30
Recall ↑
.6
5 102040
λ
160
80
50
1
Base
120
60
5 102040
λ
90
.88
1.08
λ
Precision ↑
.92
1.14
.7
80
75
.96 1.2
240
80
80
5 102040
λ
.84
1.1
.44
.6
.88
1.15
.48
.75
FDval ↓
240
1.2
Inception
DINOv2-L
FD ↓
.52
.8 .76
.92
SiT-XL/2 240
.8
2
λ
.88 .84
.6
.64
5 102040
.76
.9 .85
.775
1
Coverage ↑
.8
.64
3
4.5
.88 .84
.8
.68
.824
6
Density ↑
.9
.7
.832
21
Recall ↑
.78
.72
6
15
DRL (λ sweep, R1 =10−5 )
DRL (λ=1, R1 =0)
.75
.72
25
Inception
SigLIP-L
DINOv3-L
DINOv2-L
RAE DiTDH-XL
1
5 102040
λ
λ
.84
.9
.78 1
5 102040
λ
1
5 102040
λ
1
5 102040
λ
Figure 17 Quantitative lambda sweep (RAE + SiT). Each subplot shows one metric for one embedder. Dashed gray: base
model at CFG = 1. Dotted (model color): theoretically motivated λ=1 (R1 =0) checkpoint. Solid line with markers: DRL values at λ ∈ {1, 5, 10, 20, 40} from the R1 =10−5 sweep. Rows: embedders (DINOv2-L, DINOv3-L, SigLIP-L, Inception). Columns: metrics (FD↓, FDval ↓, KD↓, Precision↑, Recall↑, Density↑, Coverage↑).
62
FD ↓ 120
150
KD ×1000 ↓
100
80
50
32
32
24
30
24
16
40 32
20
16
10
40
100
60
8
6
6
4.5
4
3 1
5 102040
Precision ↑
Recall ↑ .75
.6
.6
.56
.45
.6
.75
.475
.85
.8
.7
.45
.8
.72
.425
.75
.84
.3
1.2
.81
.28
.78
.26
3
.72
2
.66
λ
.6 1
λ
5 102040
1
DINOv3-L
32
Inception
100
50
60
FDval ↓ 150
100
SigLIP-L
DINOv2-L
FD ↓
150
KD ×1000 ↓
.8
100
.7
50
.6
60
40
40
20
16
5
λ
1
.9
.48
.8 1
5 102040
.87 1
λ
Base
DRL (λ=1, R1 =0)
Precision ↑
Recall ↑
5 102040
1
λ
DRL (λ sweep, R1 =10−5 )
Density ↑
Coverage ↑ .8
.75
.7
.4
.5
.6
.8
.6 .5
.4
1.1 1
25
.6
.3
.9
5 102040
λ
1
.35
.7
50
16
25
λ
.88
.8
.4
1.05
.8
.76
.32
.9
.72
.7
.8
.6
1.25
.9
1
.84
.5
.7
3
5 102040
1.2
.24
6
1
.48
.84
.4 .75
.6
0 5 102040
.93
.56
.7
24
6 1
1.2
.8
75
12
10
.81 .78
50
18
15
.84
.6
32
24
.87
1
75
20
.64
1.1
λ
50
Coverage ↑
.7
.64
5 102040
λ
REPA SiT-XL/2
150
.8
.64
.78
5 102040
Density ↑
.6
4
1 1
DRL (λ sweep, R1 =10−5 )
.66
.65
24
7.5
DRL (λ=1, R1 =0)
80
32
24
Base
.54
40
SigLIP-L Inception
FDval ↓
160 120
80
DINOv3-L
DINOv2-L
JiT-H/16
1
5 102040
1
5 102040
λ
1
λ
Figure 18 Quantitative lambda sweep (JiT + REPA). Same layout as figure 17.
63
5 102040
λ
.78 1
5 102040
λ
1
5 102040
λ
J.2
Alignment: Full Distribution Metrics (PRDC)
Table 2 Distribution metrics — No CFG. For each (model, embedder) pair, all rows are evaluated at cfg = 1; the DRL (λ∗ ) row uses the (λ, time interval) with the lowest FD. All seven metrics come from a single sweep entry. For each (model, embedder), three rows compare Base, DRL at the theoretically motivated λ=1 (R1 =0), and DRL with λ tuned over {1, 5, 10, 20, 40}. Bold marks the best value per metric within each block. FD is computed against the ImageNet training set; FDval against the validation set. KD is reported as value ± standard error of the subset-mean estimator (100 random subsets). The right-most column shows the (λ, R1 ) of the run each row is drawn from. FD↓
FDval ↓
KD↓
Prec↑ Rec↑ Dens↑ Cov↑
(λ, R1 )
Base 37.47 DINOv2-L DRL (λ=1, R1 =0) 30.17 ∗ DRL (λ ) 20.60
49.27 42.47
29.75 ± 0.33 22.00 ± 0.29 9.31 ± 0.20
0.669 0.787 0.692 0.832 0.689 0.781 0.737 0.849 0.734 0.759 0.849 0.882
(−−, −−) (1, 0) (10, 10−5 )
Base DINOv3-L DRL (λ=1, R1 =0) DRL (λ∗ )
6.66 5.71
8.04 7.14
4.49
6.10
5.20 ± 0.03 4.21 ± 0.03 2.75 ± 0.02
0.731 0.697 0.802 0.849 0.742 0.698 0.821 0.859 0.769 0.684 0.886 0.877
(−−, −−) (1, 0) (10, 10−5 )
SigLIP-L
Base 10.19 DRL (λ=1, R1 =0) 9.53 ∗ DRL (λ ) 9.13
11.88 11.31
0.822 0.541 1.118 0.910 0.830 0.535 1.143 0.915 0.844 0.518 1.200 0.919
(−−, −−) (1, 0) (10, 10−5 )
Incep.
Base DRL (λ=1, R1 =0) DRL (λ∗ )
0.719 0.643 1.066 0.956 0.736 0.632 1.114 0.957 0.728 0.640 1.092 0.957
(−−, −−) (1, 0) (1, 10−5 )
276.55 ± 1.28 0.469 0.621 0.399 0.451 81.43 ± 0.51 0.564 0.609 0.551 0.659 22.62 ± 0.18 0.733 0.510 0.920 0.782
(−−, −−) (1, 0) (10, 10−5 )
117.44 ± 0.51 0.470 0.314 0.803 0.453 46.23 ± 0.23 0.581 0.347 0.754 0.647 10.90 ± 0.05 0.781 0.278 1.019 0.749
(−−, −−) (1, 0) (20, 10−5 )
108.37 ± 0.52 0.712 0.395 0.761 0.700 40.37 ± 0.24 0.776 0.417 0.977 0.823 30.20 ± 0.21 0.817 0.382 1.128 0.843
(−−, −−) (1, 0) (5, 10−5 )
5.27 ± 0.08 0.85 ± 0.02 0.85 ± 0.02
0.586 0.675 0.719 0.843 0.713 0.597 1.068 0.937 0.713 0.597 1.068 0.937
(−−, −−) (1, 0) (1, 0)
153.22 ± 0.71 0.515 0.648 0.463 0.596 89.20 ± 0.52 0.568 0.659 0.545 0.690 21.31 ± 0.20 0.709 0.622 0.847 0.825
(−−, −−) (1, 0) (20, 10−5 )
33.84 ± 0.13 21.87 ± 0.09 8.13 ± 0.03
(−−, −−) (1, 0) (20, 10−5 )
Model Embedder Method
RAE
SiT
JiT
REPA
1.31 1.38 1.29
33.88
11.15 2.39
0.43 ± 0.02 0.82 ± 0.03 0.61 ± 0.03
2.72 2.52
Base 241.93 DINOv2-L DRL (λ=1, R1 =0) 95.51 DRL (λ∗ ) 54.37
249.99 108.56
Base 88.22 DINOv3-L DRL (λ=1, R1 =0) 43.76 ∗ DRL (λ ) 19.25
88.53 44.88
SigLIP-L
Base 42.59 DRL (λ=1, R1 =0) 20.17 ∗ DRL (λ ) 16.90
42.91 22.11
9.38
8.13
Incep.
Base DRL (λ=1, R1 =0) DRL (λ∗ )
2.62 2.62
3.63 3.63
70.97
21.19
18.94
Base 148.03 DINOv2-L DRL (λ=1, R1 =0) 97.33 ∗ DRL (λ ) 42.84
158.85 109.00
Base 33.89 DINOv3-L DRL (λ=1, R1 =0) 23.66 ∗ DRL (λ ) 11.32
35.11 25.08
SigLIP-L
Base 39.41 DRL (λ=1, R1 =0) 29.72 DRL (λ∗ ) 20.99
40.94 31.36
Incep.
Base DRL (λ=1, R1 =0) DRL (λ∗ )
6.47 3.73
7.16 3.72 2.73
20.64 ± 0.13 19.30 ± 0.15 19.68 ± 0.17
57.29
13.24
23.04
168.91 72.52
Base 63.87 DINOv3-L DRL (λ=1, R1 =0) 29.43 ∗ DRL (λ ) 15.58
64.46 30.81
SigLIP-L
Base 31.56 DRL (λ=1, R1 =0) 14.34 ∗ DRL (λ ) 12.77
32.13 16.56
6.48
5.81
Incep.
Base DRL (λ=1, R1 =0) DRL (λ∗ )
2.14 2.14
3.66 3.66
56.12
0.754 0.470
0.862
0.810
100.22 ± 0.38 0.772 0.297 0.921 0.779 71.41 ± 0.28 0.802 0.315 1.030 0.833 45.97 ± 0.20 0.856 0.302 1.270 0.875
(−−, −−) (1, 0) (20, 10−5 )
2.81 ± 0.05 0.79 ± 0.02 0.61 ± 0.01
0.588 0.664 0.714 0.860 0.652 0.645 0.871 0.905 0.699 0.616 1.001 0.930
(−−, −−) (1, 0) (5, 10−5 )
163.05 ± 0.95 0.517 0.668 0.465 0.606 45.83 ± 0.33 0.644 0.653 0.693 0.773 23.77 ± 0.22 0.721 0.615 0.877 0.829
(−−, −−) (1, 0) (5, 10−5 )
77.62 ± 0.38 28.16 ± 0.16 10.88 ± 0.06
0.540 0.385 0.848 0.597 0.660 0.412 0.855 0.763 0.791 0.386 1.066 0.827
(−−, −−) (1, 0) (10, 10−5 )
79.18 ± 0.45 28.71 ± 0.21 23.97 ± 0.21
0.729 0.456 0.809 0.769 0.795 0.470 1.033 0.864 0.823 0.449 1.138 0.874
(−−, −−) (1, 0) (5, 10−5 )
0.600 0.687 0.750 0.872 0.719 0.611 1.072 0.943 0.719 0.611 1.072 0.943
(−−, −−) (1, 0) (1, 0)
3.28
Base 159.46 DINOv2-L DRL (λ=1, R1 =0) 58.33 DRL (λ∗ ) 40.33
0.620 0.428 0.749 0.636 0.651 0.464 0.735 0.711
17.37
14.88
2.94 ± 0.06 1.16 ± 0.03 1.16 ± 0.03
64
Table 3 Distribution metrics — Best CFG. For each (model, embedder) pair and each method row, the (CFG scale,
time interval) combination minimising FD is selected; all seven metrics are taken from that single run. For each (model, embedder), three rows compare Base, DRL at the theoretically motivated λ=1 (R1 =0), and DRL with λ tuned over {1, 5, 10, 20, 40}. Bold marks the best value per metric within each block. FD is computed against the ImageNet training set; FDval against the validation set. KD is reported as value ± standard error of the subset-mean estimator (100 random subsets). The right-most column shows the configuration (λ, R1 , cfg, [interval]) of the run each row is drawn from. Model Embedder Method
RAE
SiT
FD↓ FDval ↓
Base 25.94 DINOv2-L DRL (λ=1, R1 =0) 23.66 ∗ DRL (λ ) 20.24
40.78 37.74
Base 4.83 DINOv3-L DRL (λ=1, R1 =0) 4.57 ∗ DRL (λ ) 4.26
6.66 6.44
SigLIP-L
Base 9.50 DRL (λ=1, R1 =0) 9.40 DRL (λ∗ ) 9.06
11.81 11.77
Incep.
Base 1.28 DRL (λ=1, R1 =0) 1.25 ∗ DRL (λ ) 1.19
2.32 2.43
Base 46.96 DINOv2-L DRL (λ=1, R1 =0) 42.36 ∗ DRL (λ ) 42.36
62.68
Base 21.63 DINOv3-L DRL (λ=1, R1 =0) 18.32 DRL (λ∗ ) 16.40
24.10 20.61
34.68
5.97
11.03
2.31 58.94 58.94
18.54 13.48
SigLIP-L
Base 11.81 DRL (λ=1, R1 =0) 12.17 ∗ DRL (λ ) 11.58
14.95 13.64
Incep.
Base 1.50 DRL (λ=1, R1 =0) 2.43 ∗ DRL (λ ) 1.78
4.17 3.26
2.90
Base 41.55 DINOv2-L DRL (λ=1, R1 =0) 36.10 DRL (λ∗ ) 30.64
54.96 49.96
Base 10.89 DINOv3-L DRL (λ=1, R1 =0) 10.19 ∗ DRL (λ ) 8.10
13.16 12.43
Base 18.21 DRL (λ=1, R1 =0) 17.11 DRL (λ∗ ) 16.47
19.93 18.94
Base 1.91 DRL (λ=1, R1 =0) 1.87 ∗ DRL (λ ) 1.87
3.16
Base 36.87 DINOv2-L DRL (λ=1, R1 =0) 32.87 ∗ DRL (λ ) 32.87
52.75
Base 18.60 DINOv3-L DRL (λ=1, R1 =0) 14.58 ∗ REPA DRL (λ ) 13.01
20.97 16.80
JiT
SigLIP-L
Incep.
45.59
10.14
18.43 3.05 3.05 49.11 49.11
15.00 11.18
SigLIP-L
Base 9.52 DRL (λ=1, R1 =0) 10.02 DRL (λ∗ ) 9.32
12.77 11.34
Incep.
Base 1.21 DRL (λ=1, R1 =0) 2.14 ∗ DRL (λ ) 1.50
3.66 3.05
2.77
KD↓
Prec↑ Rec↑ Dens↑ Cov↑
(λ, R1 , cfg, [int.])
16.59 ± 0.22 0.764 0.677 0.924 0.887 15.22 ± 0.22 0.743 0.723 0.870 0.884 10.51 ± 0.18 0.762 0.719 0.915 0.897
(−−, −−, 1.5, [0, 1]) (1, 0, 1.25, [0, 1]) (5, 10−5 , 1.25, [0, 1])
3.37 ± 0.02 2.91 ± 0.02 2.71 ± 0.02
(−−, −−, 1.5, [0, 1]) (1, 0, 1.5, [0, 1]) (5, 10−5 , 1.25, [0, 1])
0.810 0.608 0.992 0.895 0.820 0.605 1.013 0.895 0.800 0.651 0.954 0.895
22.86 ± 0.20 0.855 0.494 1.252 0.925 23.36 ± 0.23 0.860 0.488 1.280 0.927 19.18 ± 0.17 0.839 0.524 1.185 0.917 0.35 ± 0.02 0.49 ± 0.02 0.38 ± 0.02
(−−, −−, 1.25, [0, 1]) (1, 0, 1.25, [0, 1]) (10, 10−5 , 1.75, [.1, 1])
0.714 0.651 1.043 0.953 0.716 0.645 1.052 0.953 0.714 0.652 1.042 0.952
(−−, −−, 2, [.1, 1]) (1, 0, 2.5, [.1, 1]) (1, 10−5 , 2.25, [.1, 1])
28.17 ± 0.25 0.801 0.436 1.171 0.862 23.46 ± 0.25 0.796 0.465 1.148 0.860 23.46 ± 0.25 0.796 0.465 1.148 0.860
(−−, −−, 2.75, [.1, 1]) (1, 0, 2.75, [.3, 1]) (1, 0, 2.75, [.3, 1])
16.70 ± 0.08 0.835 0.156 1.337 0.800 15.16 ± 0.08 0.837 0.220 1.303 0.840 10.99 ± 0.05 0.851 0.234 1.354 0.831
(−−, −−, 3.5, [0, 1]) (1, 0, 3, [.1, 1]) (10, 10−5 , 2.25, [.3, 1])
25.30 ± 0.21 0.842 0.400 1.241 0.898 27.46 ± 0.29 0.842 0.405 1.238 0.887 24.76 ± 0.26 0.845 0.401 1.253 0.896 0.70 ± 0.02 1.77 ± 0.04 1.08 ± 0.03
0.713 0.626 1.065 0.941 0.742 0.582 1.169 0.947 0.734 0.610 1.136 0.946
(−−, −−, 3.5, [.3, 1]) (1, 0, 2.25, [.3, 1]) (1, 10−5 , 2.75, [.3, 1]) (−−, −−, 2, [.3, 1]) (1, 0, 1.25, [.3, 1]) (1, 10−5 , 1.5, [.3, 1])
31.69 ± 0.25 0.737 0.576 0.896 0.851 24.87 ± 0.23 0.753 0.585 0.943 0.867 15.39 ± 0.17 0.788 0.574 1.054 0.884
(−−, −−, 4, [.1, 1]) (1, 0, 3.5, [.1, 1]) (10, 10−5 , 2.75, [.1, 1])
8.56 ± 0.04 7.74 ± 0.03 5.94 ± 0.03
0.813 0.376 1.038 0.841 0.816 0.397 1.039 0.847 0.818 0.456 1.029 0.862
(−−, −−, 2.75, [0, 1]) (1, 0, 2.25, [0, 1]) (20, 10−5 , 3.25, [.1, 1])
46.80 ± 0.19 0.837 0.348 1.168 0.888 42.95 ± 0.20 0.846 0.348 1.206 0.896 42.03 ± 0.21 0.854 0.339 1.253 0.899
(−−, −−, 3.5, [.1, 1]) (1, 0, 3, [.1, 1]) (5, 10−5 , 3, [.1, 1])
0.74 ± 0.02 0.74 ± 0.02 0.74 ± 0.02
0.684 0.629 0.950 0.928 0.698 0.627 0.976 0.932 0.698 0.627 0.976 0.932
(−−, −−, 2.25, [.1, 1]) (1, 0, 2.5, [.3, 1]) (1, 0, 2.5, [.3, 1])
22.95 ± 0.22 0.809 0.511 1.111 0.884 20.44 ± 0.21 0.797 0.567 1.058 0.882 20.44 ± 0.21 0.797 0.567 1.058 0.882
(−−, −−, 2.75, [.1, 1]) (1, 0, 2.75, [.3, 1]) (1, 0, 2.75, [.3, 1])
14.02 ± 0.08 0.831 0.238 1.233 0.818 11.71 ± 0.06 0.841 0.324 1.235 0.866 9.02 ± 0.05 0.852 0.354 1.243 0.862
(−−, −−, 3.25, [0, 1]) (1, 0, 2.75, [.1, 1]) (10, 10−5 , 2.25, [.3, 1])
20.65 ± 0.20 0.830 0.488 1.141 0.910 23.03 ± 0.27 0.832 0.475 1.172 0.899 20.71 ± 0.23 0.833 0.482 1.165 0.907 0.71 ± 0.03 1.16 ± 0.03 0.94 ± 0.03
65
0.702 0.642 1.004 0.939 0.719 0.611 1.072 0.943 0.718 0.621 1.067 0.945
(−−, −−, 3.75, [.3, 1]) (1, 0, 2.25, [.3, 1]) (1, 10−5 , 3, [.3, 1]) (−−, −−, 2.25, [.3, 1]) (1, 0, 1, [0, 1]) (1, 10−5 , 1.5, [.3, 1])
J.3
Image Quality: Reward Improvement vs. λ
Figure 19 Reward improvement vs. λ. Per-model normalized improvement on each of four held-out preference rewards.
Solid lines with markers: DRL at λ ∈ {1, 5, 10, 20, 40} from the R1 =10−5 sweep. Dotted lines (in the model’s color): theoretically motivated λ=1 (R1 =0) checkpoint. Dashed gray line at zero: Base model. Top row: CFG = 1 (no guidance). Bottom row: best CFG. Same units as figure 7. Base
ImageReward
Best CFG (DRL−Base)/σbase
No CFG (DRL−Base)/σbase
1.25
DRL (λ sweep, R1 =10−5 )
DRL (λ=1, R1 =0)
HPSv2
1.25
SiT
REPA
Aesthetic
1.25
JiT
1.00
1.00
1.00
0.75
0.75
0.75
0.75
0.50
0.50
0.50
0.50
0.25
0.25
0.25
0.25
0.00
0.00
0.00
0.00
0.4
0.4
0.4
0.4
0.3
0.3
0.3
0.3
0.2
0.2
0.2
0.2
0.1
0.1
0.1
0.1
0.0
0.0
0.0
0.0
−0.1
−0.1 1
5
10
20
40
−0.1 1
λ
5
10
20
40
PickScore
1.25
1.00
RAE
−0.1 1
5
λ
10
20
40
λ
1
5
10
20
40
λ
Table 4 Numerical values for Figure 7. DRL is pinned to the headline λ=10; the guidance scale is tuned per (model, metric) cell. Bold indicates the better value within each Base/DRL pair.
No guidance Base
DRL
Guidance
Model
Metric
Base
DRL
SiT-XL/2
ImageReward↑ -1.252 -0.486 -0.584 -0.309 HPSv2↑ 0.197 0.224 0.233 0.242 Aesthetic↑ 3.149 3.362 3.473 3.533 PickScore↑ 19.140 20.086 20.153 20.383
REPA SiT-XL/2
ImageReward↑ -1.187 -0.562 -0.622 -0.415 HPSv2↑ 0.200 0.223 0.230 0.237 Aesthetic↑ 3.185 3.386 3.456 3.508 PickScore↑ 19.240 20.044 20.094 20.292
JiT-H/16
ImageReward↑ -1.159 -0.804 -0.723 -0.551 HPSv2↑ 0.196 0.210 0.226 0.231 Aesthetic↑ 3.078 3.192 3.321 3.366 PickScore↑ 19.239 19.710 19.904 20.079
ImageReward↑ -0.933 -0.764 -0.583 -0.488 HPSv2↑ 0.217 0.222 0.227 0.230 RAE DiTDH-XL Aesthetic↑ 3.292 3.338 3.398 3.425 PickScore↑ 19.784 19.990 20.187 20.297
66
J.4
Image Quality / Preference RL: Pareto Plots DRL+PRL (ours) Base+PRL
ImageNet λPRL =1
λPRL =10
λPRL =40
SiT
0.25
REPA 0.20
JiT
no CFG
0.20 0.25
0.22
RAE
0.20 0.24
0.22 0.5
0.6
0.2
0.3
0.5
0.6
0.25 0.30 0.35
0.18
0.20
0.200
0.225
0.22
0.14
0.16
0.16
0.18
0.4
0.5
0.4
0.5
SiT
0.26 0.24
REPA
0.24
0.24
JiT
CFG=2
0.26
0.22
RAE
0.25 0.24
SiT
0.23 0.20
0.25
REPA
AG=2
JiT
0.20 0.250
0.225
0.225
RAE
0.200 0.250 0.225 0.5
0.6
0.7
Brightness
0.2
0.3
0.15
Saturation
0.20
Contrast
0.12
0.14
0.16
Colorfulness
0.4
0.6
Whiteness
Figure 20 HPSv2 vs. image statistics across guidance settings. Each panel-group is a 4 (model) × 5 (image statistic) grid;
rows correspond to no CFG (CFG=1), CFG=2, and autoguidance=2 (see section J for context on autoguidance). Markers: base (⋆), Base+PRL (dashed, hollow), DRL+PRL (solid, filled) at λ ∈ {1, 10, 40}. Vertical dashed green line: ImageNet reference value per statistic. The autoguidance group uses each method’s pre-PRL checkpoint as the negative signal at scale 2 (Base for Base+PRL, DRL for DRL+PRL); the base reference still uses standard CFG at the same scale. DRL+PRL consistently attains better Pareto fronts across guidance settings.
67
SiT
DRL+PRL (ours) Base+PRL
ImageNet λPRL =1
λPRL =10
λPRL =40
0
REPA
0 −1 0
JiT
no CFG
−1
−1
RAE
0.0 −0.5 0.5
0.6
0.2
0.3
0.5
0.6
0.25 0.30 0.35
0.18
0.20
0.200
0.225
0.22
0.14
0.16
0.16
0.18
0.4
0.5
0.4
0.5
SiT
0.5 0.0
REPA
0.5 0.0 −0.5 0.0 −0.5
RAE
JiT
CFG=2
−0.5
0.0
SiT
−0.5 0.20
0
REPA
0 −1 0
JiT
AG=2
−1
−1
RAE
0.5 0.0
−0.5 0.5
0.6
0.7
Brightness
0.2
0.3
0.15
Saturation
0.20
Contrast
0.12
0.14
0.16
Colorfulness
0.4
0.6
Whiteness
Figure 21 ImageReward vs. image statistics across guidance settings. Mirrors figure 20 using ImageReward as the held-out
reward. DRL+PRL consistently attains better Pareto fronts across guidance settings.
68
SiT
DRL+PRL (ours) Base+PRL
ImageNet λPRL =1
λPRL =10
λPRL =40
3.50
REPA
3.50 3.25 3.4
JiT
no CFG
3.25
3.2
RAE
3.6 3.4 0.5
0.6
0.2
0.3
0.5
0.6
0.25 0.30 0.35
0.18
0.20
0.200
0.225
0.22
0.14
0.16
0.16
0.18
0.4
0.5
0.4
0.5
SiT
3.75 3.50
REPA
3.50 3.5
JiT
CFG=2
3.75
3.4
RAE
3.3 3.6 3.4 0.20
SiT
3.75 3.50
REPA
3.75 3.50 3.25 3.50
JiT
AG=2
3.25
3.25
RAE
3.75 3.50
0.5
0.6
0.7
Brightness
0.2
0.3
0.15
Saturation
0.20
Contrast
0.12
0.14
0.16
Colorfulness
0.4
0.6
Whiteness
Figure 22 Aesthetic v2.5 vs. image statistics across guidance settings. Mirrors figure 20 using Aesthetic v2.5 as the held-out
reward. DRL+PRL consistently attains better Pareto fronts across guidance settings.
69
J.5
Ablations: Discriminator Full R1 Sweep
Figure 23 Discriminator ablation: full R1 sweep. DINOv2-L FD vs. R1 for four head/training configurations, at λ=1 (left)
and λ=10 (right). Top: fixed-seed samples at four representative configurations. Dashed gray line: base model FD. λ=1 R1 = 0
Base
Linear
λ = 10 R1 = 10−5
λ = 10 R1 = 0
MLP-2
Finetune
Scratch
λ=1
λ = 10
FD
103
base 102 10−5
0
J.6
R1
10−3
0.1
10−5
0
R1
10−3
0.1
Ablations: Full Feature-Space Sweep
Table 5 Feature space ablation results. Fréchet Distance (FD, ↓) and Kernel Distance (KD×103 , ↓) when training the
discriminator in different feature spaces (columns) and evaluating in different feature spaces (rows). Best fine-tuned result per row in bold. Training a discriminator on a particular embedding consistently improves downstream alignment metrics based on that embedder; DINOv2-L gives the best cross-feature transfer. Training feature space
KD×103
FD
Eval. space Base DINOv2-B DINOv2-L DINOv3-B DINOv3-L SigLIP-B SigLIP-L Incep. Inception DINOv2-B DINOv2-L DINOv3-B DINOv3-L SigLIP-B SigLIP-L
6.43 130 159 22.4 63.7 7.54 31.2
2.18 50.8 71.0 10.8 31.1 3.85 16.1
Inception DINOv2-B DINOv2-L DINOv3-B DINOv3-L SigLIP-B SigLIP-L
2.89 192 165 36.3 78.3 18.5 78.8
0.66 57.0 60.5 16.3 30.0 8.42 34.4
2.14 44.3 58.3
9.81 29.4 3.41 14.3
1.16 47.5 45.8
14.5 28.2 6.84 28.7
2.19 54.3 74.7 10.7 31.0 3.80 16.0 0.62 63.8 64.7 15.7 29.2 8.47 33.9
70
3.02 47.3 63.2 9.53 18.5
4.28 16.5 1.69 52.9 53.1 13.9 14.3
10.8 38.7
2.57 72.9 96.1 14.0 44.0 4.11 18.0
2.34 64.5 87.1 12.4 38.8 3.88 16.1
5.68 123 151 21.5 62.0 7.19 29.9
0.38 94.5 89.1 21.9 49.0 8.90 38.5
0.36
2.30 181 156 34.9 75.8 17.5 74.8
80.3 78.9 19.1 41.1 8.28 33.2
J.7
Effect of CFG on Reward Scores and Fréchet Distance Base
DRL λ = 1
SiT
HPSv2 ↑
3.50
REPA
DRL λ = 10
Aesthetic ↑
DRL λ=1 (R1=0)
PickScore ↑
−0.5
20.0
3.25
−1.0
19.5
0.225
3.4
−0.5
20.0
3.2
−1.0
19.5
0.200 0.225
3.3
−0.5
20.0
3.2 0.200
3.1
−1.0
19.5
3.4
−0.50
20.25
−0.75
20.00
0.23
RAE
DRL λ = 20
ImageReward ↑
0.225 0.200
JiT
DRL λ = 5
0.22
3.3 2
4
2
CFG scale
4
2
CFG scale
4
2
CFG scale
4
CFG scale
Figure 24 Effect of CFG on reward scores. Reward vs. CFG scale for four models (rows) and four reward metrics (columns). Black: base model (drawn on top of the DRL curves for visibility); colored lines: DRL at λ ∈ {1, 5, 10, 20} (darker = higher λ); purple dashed: DRL λ=1 with R1 =0 discriminator. DRL beats base at every CFG, and sometimes (e.g., SiT on ImageReward) even without CFG.
Base
DRL λ = 1
JiT
REPA
SiT
FD Inception ↓
DRL λ = 5
DRL λ = 10
FD DINOv2 ↓
90 60
DRL λ=1 (R1=0) 45 30
10
240 160
3
80
30
8 4
150 100
60 40
50
20
150 100
30 20
40
10
20
6 4
50 2
RAE
DRL λ = 20
FD DINOv3 ↓
40
2
30
6
20
4.5
2
CFG scale
4
15 30 20 10 30
7.5
3
2
4
CFG scale
FD SigLIP ↓
11 10 9 2
CFG scale
4
2
4
CFG scale
Figure 25 Effect of CFG on Fréchet Distance. FD vs. CFG scale for four models (rows) and four evaluation embedders (columns). All panels use a log y-axis. Each point shows the best result across the three CFG intervals defined in section I. Same legend as figure 24. Inception/SigLIP prefer smaller λ, DINOv2/v3 prefer larger; CFG also helps less as λ grows. Both are likely artifacts of fine-tuning only the conditional branch (required by adjoint matching).
71
K
Qualitative Samples
This section collects the qualitative sample grids referenced from the main text. All grids use matched noise seeds and class labels across conditions, so visual differences are attributable to model weights and guidance only. Base vs. Fine-tuned. Uncurated samples comparing the base model against the DRL fine-tuned model at
λ=10 for each architecture. Each figure shows two sections side by side: No CFG and Best CFG, where the latter selects the best CFG value according to DINOv3-L FD. We use DINOv3 as a held-out metric since the discriminator is trained on DINOv2 features. Effect of DRL Strength. Samples for each model under varying DRL strength (λ ∈ {1, 5, 10, 20, 40}) and three
guidance strategies. Each figure uses matched noise seeds and class labels across all conditions, so differences are purely due to model weights. Reading left-to-right shows the effect of increasing λ; the three sections compare no guidance (CFG = 1), standard CFG = 2, and the per-variant optimal guidance (DINOv3-L FD-optimal CFG or autoguidance, annotated above each column; see section J for context on autoguidance). As remarked in section J, autoguidance is often the best guidance setting at large λ, and noticeably improves the images compared to the CFG variant. RL Fine-tuning Sample Images. Generated samples for each model under varying RL strength (λPRL ∈ {1, 10, 40}) and guidance strategies (no CFG, CFG = 2, and autoguidance = 2; see section J for context on autoguidance). Each figure uses matched noise seeds across all conditions, so differences are purely due to model weights and guidance. The three column groups compare Base, DRL+PRL, and Base+PRL. As discussed in the main text, increasing λPRL makes the images progressively brighter and more distorted. This effect is particularly pronounced under autoguidance, which—as noted in sections J and J.4—yields the highest reward but the largest distortion.
72
K.1
Base vs. Fine-tuned JiT H-16 No CFG
Best Guidance base: CFG
recreational vehicle
Base
λ=10
balloon
Base
λ=10
zebra
Base
λ=10
leopard
Base
λ=10
coral reef
Base
λ=10
Figure 26 Base vs. DRL fine-tuned (λ=10) samples for JiT H-16.
73
λ=10: CFG
SiT-XL/2 No CFG
Best Guidance base: CFG
recreational vehicle
Base
λ=10
balloon
Base
λ=10
zebra
Base
λ=10
leopard
Base
λ=10
coral reef
Base
λ=10
Figure 27 Base vs. DRL fine-tuned (λ=10) samples for SiT-XL/2.
74
λ=10: CFG
REPA SiT-XL/2 No CFG
Best Guidance base: CFG
recreational vehicle
Base
λ=10
balloon
Base
λ=10
zebra
Base
λ=10
leopard
Base
λ=10
coral reef
Base
λ=10
Figure 28 Base vs. DRL fine-tuned (λ=10) samples for REPA SiT-XL/2.
75
λ=10: CFG
RAE DiTDH-XL No CFG
Best Guidance base: CFG
recreational vehicle
Base
λ=10
balloon
Base
λ=10
zebra
Base
λ=10
leopard
Base
λ=10
coral reef
Base
λ=10
Figure 29 Base vs. DRL fine-tuned (λ=10) samples for RAE DiTDH-XL.
76
λ=10: CFG
K.2
Effect of DRL Strength
JiT H-16 Base
λ=1
λ=5
λ=10
λ=20
λ=40
CFG
CFG
No CFG
CFG = 2
Best Guidance CFG
CFG
CFG
CFG
Figure 30 Effect of DRL strength (λ) on samples for JiT H-16.
77
SiT-XL/2 Base
λ=1
λ=5
λ=10
λ=20
λ=40
AG
AG
No CFG
CFG = 2
Best Guidance CFG
CFG
CFG
CFG
Figure 31 Effect of DRL strength (λ) on samples for SiT-XL/2.
78
REPA SiT-XL/2 Base
λ=1
λ=5
λ=10
λ=20
λ=40
AG
AG
No CFG
CFG = 2
Best Guidance CFG
CFG
AG
CFG
Figure 32 Effect of DRL strength (λ) on samples for REPA SiT-XL/2.
79
RAE DiTDH-XL Base
λ=1
λ=5
λ=10
λ=20
λ=40
AG
AG
No CFG
CFG = 2
Best Guidance CFG
CFG
CFG
CFG
Figure 33 Effect of DRL strength (λ) on samples for RAE DiTDH-XL.
80
K.3
RL Fine-tuning Sample Images
JiT H-16 Base ref
DRL + Pref. RL λ=1
λ=10
Base + Pref. RL λ=40
No CFG
CFG = 2
Autoguidance
Figure 34 RL fine-tuning samples for JiT H-16.
81
λ=1
λ=10
λ=40
SiT-XL/2 Base ref
DRL + Pref. RL λ=1
λ=10
Base + Pref. RL λ=40
No CFG
CFG = 2
Autoguidance
Figure 35 RL fine-tuning samples for SiT-XL/2.
82
λ=1
λ=10
λ=40
REPA SiT-XL/2 Base ref
DRL + Pref. RL λ=1
λ=10
Base + Pref. RL λ=40
No CFG
CFG = 2
Autoguidance
Figure 36 RL fine-tuning samples for REPA SiT-XL/2.
83
λ=1
λ=10
λ=40
RAE Base ref
DRL + Pref. RL λ=1
λ=10
Base + Pref. RL λ=40
No CFG
CFG = 2
Autoguidance
Figure 37 RL fine-tuning samples for RAE.
84
λ=1
λ=10
λ=40