CAR-Flow: Condition-Aware Reparameterization Aligns Source and Target for Better Flow Matching
Authors: Chen Chen
Pengsheng Guo
Liangchen Song
Jiasen Lu
Rui Qian
Tsu-Jui Fu
Xinze Wang
Wei Liu
Yinfei Yang
Alex Schwing
Venue: NeurIPS 2025 spotlight ()
Keywords: image generation
diffusion model
flow matching
Primary area: deep_learning

Abstract: Conditional generative modeling aims to learn a conditional data distribution from samples containing data-condition pairs. For this, diffusion and flow-based methods have attained compelling results. These methods use a learned (flow) model to transport an initial standard Gaussian noise that ignores the condition to the conditional data distribution. The model is hence required to learn both mass transport \emph{and} conditional injection. To ease the demand on the model, we propose \emph{Condition-Aware Reparameterization for Flow Matching} (CAR-Flow) -- a lightweight, learned \emph{shift} that conditions the source, the target, or both distributions. By relocating these distributions, CAR-Flow shortens the probability path the model must learn, leading to faster training in practice. On low-dimensional synthetic data, we visualize and quantify the effects of CAR-Flow. On higher-dimensional natural image data (ImageNet-256), equipping SiT-XL/2 with CAR-Flow reduces FID from 2.07 to 1.68, while introducing less than \(0.6\%\) additional parameters.

CAR-Flow: Condition-Aware Reparameterization Aligns Source and Target for Better Flow Matching
Chen Chen
Pengsheng Guo"
Liangchen Song
Jiasen Lu
Rui Qian
Xinze Wang
Tsu-Jui Fu"
Wei Liu"
Yinfei Yang
Alex Schwing
Apple Inc.
Abstract
Conditional generative modeling aims to learn a conditional data distribution from samples containing data-condition pairs. For this, diffusion and flow-based methods have attained compelling results. These methods use a learned (flow) model transport an initial standard Gaussian noise that ignores the condition to the conditional data distribution. The model is hence required to learn both mass transport and conditional injection. To ease the demand on the model, we propose Condition-Aware Reparameterization for Flow Matching (CAR-Flov) lightweight, learned shift that conditions the source, the target, or both distributions. By relocating these distributions, CAR-Flos --R-Flow shortens the probability path the model must learn, leading faster training in practice. On low-dimensional synthetic data, we visualize and quantify the effects of CAR-Flow. On higher-dimensional natural image data (ImageNet-256), equipping SiT-XL/2 with CAR-Flo -lo-wlow reduces FID from 2.07 to 1.68, while introducing less than 0.6% additional parameters.
Introduction
Conditional generative models enable to draw samples conditioned on an external variable- for example, class label, text caption, semantic mask. They are a key technology that has advanced
significantly in the last decade, from variational auto-encoders (VAEs) Kingma and Welling (014]
and generative adversarial nets [Goodfellow et al. 2014) to diffusion
Ho et al.) 2020, Song et al.
2021a b. Dhariwal and Nichol 2021 2021 Peebles and Xie 2023
Rombach et al.] 2022
Nichol and
Dhariwal 2021] and flow now-magching [Ma et al. 2024 Liu et al. 2023 Lipman et al.) 2023 Albergo
and Vanden- 2023 Albergo et al. 2023].
State-of-the-arr diffusion and flow-matching frameworks accomplish conditional generation by using trained deep net to trace probability path that progressively transforms samples from simple source distribution into the rich, condition-dependee distribution is single, condition-agnosse enters only through the network itself: in flow matching, for instance, the model predicts velocity field where the condition is commonly incorporated via embeddings or adaptive normalization layers. Although this strategy has enabled impressive results, this design forces the network to shoulder two tasks simultaneously-t transporting probability mass to the correct region of the data manifold, and (ii) encoding the semantic meaning of the condition. Because different conditions often occupy distant parts of that manifold, the dual burden stretches the learned trajectory, slows convergence, and can impair sample quality and diversity.
target distribution. A popular choice for the source standard Gaussian. Consequently. the conditioning signal
Work done while at Apple.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

volar. t, y)
kii Ps
P2
y)
(4.)
f(-9)
g( M)
(- l3)
(2t, L, y)
2 (1
Figure 1: Condition-Awars Reparameterization for Flow Matching (CAR-Flow). Illustration of the push-forward under standard conditional flow matching (direct mapping 25 - 2:1 versus Condition-Awaz Reparameterization (CAR-Flow) chain (xp + 2p + 21 + x1). In the standard setting, cundition-agnostio prior sample (red) is carried by the network's 's velocity field directly to each condition-dependent data manifold (blue), forcing it to juggle long-range transport and semantic injection at once. CAR-Flow, in contrast, employs lightweight source distribution map f(- y) and target distribution map g(-, y) to align the source and target distributions to relieve the network of unnecessary transport. During sampling, 21 is obtained via the (approximate) inverse map g (- y).
In this paper, we alleviate this burden via a condition-deependett mapping of the source and target distributions. First, we endow the source distribution with condition-awareness through lightweight map, i.c., the source distribution map. Hence, every condition has its own source distribution. The same idea can be mirrored the target distribution using another lightweight map, i.c., the target distribution map, that we require to be approximately invertible to keep sampling tractable. The resulting porw-forward chain is illustrated Figure [T] Intuitively. the target distribution map and its approximated inverse correspond to the encoder and decoder, commonly used in latent diffusion pipelines [Rombach et al. 2022]. Differently, here, the target distribution map explicitly conditional, aligning our formulation lation with recent latent-space conditioning approaches that embed semantic information directly into the latent representation [Wu et al., 2024
Leng et al, 2025 Yao
et al. 2025
Although fully general maps provide maximal flexibility, this freedom conceals a critical failure mode: the flow-matching objective admits zero-cost solutions. The predicted data distribution then collapses to single mode. A comparable drop in generation quality has been observed recent work when variational auto-encoder (VAE) is naively fine-tuned end-to-end with diffusion model [Leng al. 2025). We formalize these collapse routes and later verify them experimentally.
To preclude this trivial solution, we impose the simplest effective constraint--hiff--on conditioning. Concretely, reparameterize the maps to only translate. Relocating the start and/or end points while leaving scales untouched removes all trivial solutions, and still shortens the residual probability path. Note that volume preeeeeireervingg maps are also possible. We refer the resulting framework as condition-aware reparameterization for flow matching (CAR-Flow).
CAR-Flow can be applied to the source distribution, the target distribution, or both. Each variant shapes the learned path differently, and we find the joint version to perform best in practice. For both diw-dimensional synthetic benchmarks and high-dimensional natural image (1mageNet-256) data, CAR-FI Flow consistently improves performance: augmenting the strong SiT-XL/2 baseline [Ma] al. 2024) with CAR-Flow reduces FID from 2.07 1.68 while adding fewer than 0.6% additional parameters.
Our contributions are as follows:

---

1. We introduce CAR-Flow, simple yet powerful shift-only mapping that aligns the source. the target, or both distributions with the conditioning variable. CAR-FIe Flow relieves the velocity network of unnecessary transport while adding negligible computational overhead.
2. We provide theoretical analysis that uncovers zero-cost solutions present under unrestricted reparameterization and prove that they render the velocity collapse, thereby explaining the empirical failures of naive end-to-end VAB-diffusion training.
3. We validate CAR-Flow on low-dimensional synthetic data and the high-dimensional ImageNet-256 benchmark, consistently outperforming the standard rectified flow base- line.
2 Background and Preliminaries
We start with a brief review of the standard now-matching formulation.
2.1 Conditional Generation and Probability Paths
Conditional generation seeks to sample from a conditional data distribution 21 ~ (- y), where conditioning variable, e.g., class label text prompt. Diffusion and flow-based models tackle this problem by simulating a differential equation that traces a probability-densio path pt (21 zo. y). gradually transporting samples from simple source distribution 2n Pr into the target conditional distribution. A standard choice for the source distribution is the isotropic Gaussian pinit N(0, Ta).
The stochastic trajectory (Xi)osAg1 follows the SDE
dx, (X.. of Vlog (r v) dt + o.
(I)
where ue is the drift field, 0, is the diffusion coefficient, and W; is standard Wiener process. The term v log P. (X, y) denotes the score function, which can be written in terms of the drift field u, (Ma
al. 2024 In the deterministic limit o, O, this SDE reduces to the ODE dX, ue( y)
2.2 Gaussian Probability Paths
convenient instantiation is the Gaussian path. Let aa and 3e two continuously differentiable, monotonic noise scheduling functions satisfying the boundary conditions on -3 = 0 and a fo = 1. At time I, the conditional distribution
(- æ1, = N( IL v).
(2)
whose endpoints are po(: Ii, v) = (0, Ia) and pi (- [xi,)) = 6g1 Here denotes the Dirac
delta "distribution". Along this path, the state evolves by the interpolant x, = B,z0 + Ou with velocity field
4 BL00 d,æ1.
where overdots denote time derivatives.
2.3 Conditional Flow Matching and Sampling
Conditional flow matching trains neural velocity field vo f, v) to approximate the true velocity ur specified in Eq. B The commonly employed objective is
vo (Biz0 + (. (3:z0 + di.z) ||?
(()) E. yu-py, te-pr
(4)
sompl 21 21mep2
where py is the marginal distribution over conditions y, and pr is time density [0, 1].
After training, one draws 21 ~ pddaa by numerically integrating Eq. ① from 0 to with å vz f, y). A solver-agnostic procedure is outlined in Algorithm

---

Algorithm Sampling via conditional flow matching (standard Gaussian source)
Require: trained network vo; diffusion schedule {o=))ejo.11 number of steps N
Sample y ~ PY
2: Sample 250 ~ (0, Ia)
3: At t 1/N; æ - I0
4: for k to N do
5: f- At
6: u uo t,
1e drift */
7: Integrate SDE in Eq.
over [t, At] with (u, oz)
e.g., Muler-Maruyana //
8: end for
9: return Ø
1 sample at L = conditioned on y */
Condition-Aware Reparameterization
As detailed in Section 2 standard conditional flow matching initiates the probability path from a condition-agnostic prior, typically a standard Gaussian. Consequently, the velocity network 00(2,, t, must simultaneously learn two intertwined tasks- transporting mass and encoding se- mantics, imposing a dual burden. To alleviate this, we reparameterize the source and/or target distributions via explicit functions of the condition y. In this section, first reformulate conditional How-matching using general reparameterizations (Section [3.I). We then show that allowing arbi- trary reparameterizations leads to trivial zero-cost minima, collapsing distributions (Section 3.2). To both shorten the transport path and climinate these trivial solutions, we propose Condition-Awarr Reparameterization for Flow Matching (CAR-Flow) (Section 3.3
single mode
3.1 General Reparameterization
Formally, instead of drawing an initial value ty directly from the fixed source distribution pa iiitt we first apply
source distribution map f R mh (I, y) + such that
f v), To p
(5)
We hence obtain a sample from the modified source distribution pin via the push-forward 20 pinit( (- y) = #P
Similarly, define the target distribution map g= Th x + Rmmh (2, such that
21 g(01, v), pz (. 1)).
(6)
We characterize a sample from the target distribution in latent space pdata via the push push-forward pa (. (- y) = g(,, y) # (- y). Critically, to make sampling tractable, 9 must be approximately
invertible-i.e. we assume there exists g such that 2j R g- (21,
It is worth noting that our reparameterization framework subsumes both the classic VAE-based latent diffusion model [Rombach al., 2022] and more recent efforts to inject semantic awareness directly into the latent space, e.g., work by Leng et al. (2025). To recover the standard latent diffusion model, we leave the source untouched, i.c., f(zo, y) Zo. and set g( y) E(1),, g (2 3) 2 (z))) where E, D are the encoder and decoder of VAE trained via the ELBO (reconstruction loss plus KL. divergence) without explicit semantic supervision. In contrast, more recent works Yuu et al. 2025
Leng et al. 2025 Yao al. 2025) augment VAE training with a semantic alignment loss often using features from a pretrained DINO-V? V2 network- so that the encoder and decoder effectively depend on semantic embeddings of y. Both paradigms arise naturally as special cases of our general push-forward reparameterization framework. Moreover, our framework readily accommodates further reparameterization of any pretrained VAE. Concretely, one can introduce an invertible map g': R x x + mm and set
g( 2 a ()((x))),)), (=1()) D((
thus embedding semantic conditioning directly into both the encoder and decoder.
Flow-matching loss and sampling under reparameterization. When plgit( 1) Gaussian, the probability path 2g + remains Gaussian under the interpolation 2t = 3,20 01 44

---

Algorithm 2 Sampling via conditional flow matching (reparameterized)
Require: trained to; diffusion schedule {ar}: steps N
Sample y ~ py
2: Sample 2 ~ (0, Ia) and set = <- f( v)
3: At - 1/N
4: for k 0 to N do
5: +
6: u u uo(2, t, y)
7: Integrate SDE Eq. over [t, + Ad] with 0.)
8: end for
9: I 1- (z,
1 invertible needed here */ 1 sample at conditioned y %/
return
8:20 2, The evolution of Z, then follows the SDE
of Vlog pelZe w) di + o, a.
dz.
U1 (Z.
(8)
where pt is the Gaussian path in z-space. Note that the conditional score functions generally differ: V log pe (Ze y) * V log Px y), unless f(zo, y) = zo (see Appendix A for details). The resulting now-matching loss becomes
(()) = E y--py: txpr sarups
10 (3:20 + 42214,y) (B:20 + 0121)
Once trained, sampling proceeds by first drawing an initial starting point p and computing the condition-awart latent 20 f(ro. y). We then integrate the reparameterized SDE Eq. (8) from t = 0 to to obtain and finally map back to data space via 2 (21, v). Algorithm provides a pseudo-code that summarizes the process.
3.2 Mode Collapse under Unrestricted Reparameterizzation
Under the general reparameterization framework of Eq. (9). the maps f and enjoy full flexibility. Unfortunately, this expressivity also admits several trivial zero-cost minima: analytically, the loss can be driven to zero by degenerate shift solutions, resulting in a collapse of the generative distribution to single/improper mode.
Claim 1. Let f. g: R" Y + Rm be arbitrary maps. If any of the following holds (for some functions cly) Th scalars k(y) R)::
(i) Constant source: [00, v) = cly) for all do-
(i) Constant target: g( v) cly) for all TT.
(ii)) Unbounded source scale: [ 8)ll + 20,
(iv) Unbounded target scale: (1g(21,9) y)1l + oc,
(v) Proportional collapse: f v) = k(y) g(m,,)),
then the dow-matching loss in Eq. admits trivial minimum in which the optimal velocity field takes the form
Uo Ze, t, 7t, v) nnt, v).
causing all probability mass to collapse to single/improper mode.
The proof of Claim
together with closed-form expressions for ((t, y) and 17(t, y). is deferred the
Appendix B
To illustrate the collapse concretely, we consider the linear schedule = t, a = and an affine reparameterization
[(to y) = ao (v) so + po(v), g(z1,y) a (v) क + pn (y),
(10)
analogous to the standard Gaussian trick. Table [ summarizes each collapse mode, listing the maps, the closed-fome form collapsed velocity (x (t, v) + n(t, y), and the resulting forhaforward distributions piitt and pataa
5

---

Table 1: Zero-cost collapse modes under the linear schedule a₁ =t, t, = using affine maps
f(0) = Jozo + No and g = 01 + f1. For brevity, we omit explicit y y-dependence.
Case
f
(1)
s((1)
Vo (21,1)
init Ps
Ps
μ4 (ii) arbitrary
arbitrary Mi arbitrary
6Rm
m
6pz
T
(ii) 80
Uniform((R
T-t
(iv)
arbitrary
20
0
Uniform (RR)
VV) M
kpo
(A: - 1)40
(k - 1)po
pz
We note that in the constant-map scenarios (cases and (ii)), the affine transforms or effectively collapse the variance of pinit or pdata to zero. In particular, case (ii) mirrors the finding in REPA-E (Table 1) [Leng et al. 2025], where end-to-end VAE-diffusion tuning favored simpler latent space with reduced variance. In contrast, the unbounded-scule modes (iii) and (iv) push the distributions toward an improper uniform distribution, eliminating any meaningful localization. Proportional collapse (v) yields degenerate flow with p puaat and the velocity field being constant.
Empirically, we do not observe cases (iii)-vv). In particular, cases (iii) and (iv) correspond to unbounded collapse solutions: when the scale of either the source or target distribution tends to infinity, the counterpart distribution collapses relative it, yielding zero cost. These solutions require unbounded weights and are therefore unstable any perturbation pushes the optimization back toward cases (i) or (ii). Case (v), in contrast, assumes exact proportionality between maps f and g. which cannot occur under independently sampled inputs except in trivial constant-map settings that reduce to cases (i) or (ii). In practice, the optimizer thus follows the easiest "shortcut" the constant-mup collapse-sing setting ʃ or to fixed constant immediately zeroes the loss. In Section 4 we empirically verify not only that such constant-map collapses occur in real models, but also how these behaviors manifest in practice.
3.3 Shift-Only Condition-Aware Reparamcterization
To eliminate the trivial zero-cost solutions while retaining the benefits of condition aware reparam- eterization, we restrict both reparameterizations to additive shifts only, while noting that volume- preserving maps are possible. For simplicity we use
( y) = zo + po(y); g( = + (1
(11)
Here, Ho: N1: 3 Th are lightweight, learnable, condition detttioo--deedeent shifts. By preserving scale, we block every collapse mode given in Claim yet we still move zg and 2ι closer in latent space. In particular, when used with a pretrained VAE, Eq. (7) becomes
g( v) & E(x1) + po(y),
g g a D(21 -n)))).
(12)
CAR-Flow admits three natural variants:
Source-only: μn = 0, so only the source is shifted.
Target-only: #o = 0, so only the target is shifted.
Joint: both μ4n and (41 active.
Each variant alters the probability path र = 3r (xo + (o(x)) + az (æ1 + (1())) in a distinct way. Note that shifting the source cannot be replicated by opposite shift on the target. In fact, no nonzero constant shift on the source can ever be matched by constant shift on the target-except in the trivial case:
Claim 2. Shifting the source by pn is equivalent to shifting the target by μsd and only if = p1 0.
Proof. With a source-only shift (po, p1) (no, 0). the interpolant is =[(0) = 3r (3o + po) ar 2i. With a target shift (Po, /1) = (z) it is 2,0) 3e xo + 0t (#) + 11). Equating these gives for Vt, Be μn as P1. Since μg and μ1 are t-independent functions of y. it forces p4 = #n 0.
For Source-only, the interpolant simplifies to 21 Ia + (yyooy). When poly) R po(3') for two conditions y, v/. their paths share common starting region, simplifying the early-time flow. For

---

X1
Pred A
Pred B
How Stage
Shift
traj.
X0
-2
-2
2
-2
2
-2
2
(a) Baseline
(b) Source-only
(c) Target-only
(d) Joint
Figure 2: Learned flow trajectories on 1D synthetic data. Each panel shows trajectories from source 2e (bottom) target 21 (top) for (a) baseline and CAR-Flow variants-(0 (b) source-only, (c) target-only, and (d) joint. Intermediate stages 2g and 21 reflect reparameterized coordinates. Colored densities represent predicted and ground-trute truut class distributions (red/blue: prediction; magenta/cyan: ground truth). Thin lines illustrate individual sample trajectories between in and 21. Dashed vertical lines mark 30 for each shift. The source-only CAR-Flow relocates the source distribution per class, while the target-only variant unifies the trajectory endpoints. The joint variant combines both and achieves the best alignment and flow quality.
Target-only, the trajectory becomes 2, = 24 04 #1 (y). When μu (v) R3 fı (v'), ). the network only needs to learn shared "landing zone", easing the late-time flow. For Joint, have 2 2t 14 (v) where (y) 3e no + 0z (y), so the umee-aryynng shift aligns both endpoints, minimizing the overall transport distance and reducing the burden Ey throughout the entire trajectory.
Empirically, we find the joint CAR-Flow variant-allowing both μo and μn to adapt-yields the largest improvements in convergence speed and sample fidelity (see Section
4 Experiments
In this section, we evaluate the efficacy of the Condition- Aware Reparameterization for Flow Matching (CAR-Flow) under linear noise schedule 1 Ou and compare classic rectified flow. All experiments are done using the axlearn framework Detailed implementation settings can be found in Appendix
4.1 Synthetic Data
For our synthetic-data experiments, consider one-dimensional task where the source distribution (O, 1) and the target distribution a two-class Gaussian mixture, with class data distribution \/-1.5, 0.2?) and class B data distribution ((+1.5, 0.2?).
We encode 2t, y, and using sinusoidal embeddings before feeding them the network. In the baseline rectified-flow flow model, these embeddings are concatenated and passed through a three-layer MLPP(1,995 (1,993 parameters total) predict the velocity field uo( y, t). Training uses the loss in Eq. (4). and sampling follows Algorithm[1 [ For CAR-Flow we augment this backbone with two lightweight linear layers that map class embedding to the shifts po(y) and/or f1 (y), each adding only parameters. In the source-only variant predict μr (with μo = 0); and the joint variant we predict both. These shifts are applied via Eq, (II). training proceeds with the loss in Eq. (9), and sampling uses Algorithm 2 All models- and CAR variants set o. = (reducing the SDE to an ODE) and employ a 50-step Euler integrator for sampling. To ensure robustness, each configuration run three times, and report the average performance, noting that variance across runs is negligible.
predict 14 (with μ4 0); in the target-only variant
Figure 2 shows how CAR-Flow alters the learned flow. In Figure 2a the baseline must both transport mass and encode class information, yielding the longest paths. The source-shift variant in Figure 2b

---

Table 2: Average trajectory length |20 81
with 2o enror bounds.
Source-only CAR-Flow
Target-only CAR-Flow
Joint
Baseline
CAR-Flow
Length
1.5355 = 0.0024
0.7432 + 0.0019
0.7129 0.0010
0.7121 = 0.0011
Source-only
Wasserstein Distance
10a
Baseline
Shift
Source-only
Source-only
Joint (B)
10-1.
Target-only Joint
Joint (A)
Target-only (A)
Shift A
Target-only (8)
0
Joint (B)
Joint (A)
10
20
30
40
50
20 30
40
Training Steps (K)
Training Steps (K)
(a) Wasserstein distance
(b) Learned po and μr shifts
Figure 3: Comparison of convergence and learned shifts. (a) shows the Wasserstein distance between predicted and ground-truth distributions in symlog-scale. Joint CAR-Flow achieves both the fastest convergence (b) plots the evolution of the learned shifts di (top) and Yi (bottom) for two classes.
cleanly relocates each class's s start, the target-shift variant in Figure 2c leaves the source untouched and merges endpoints, and the joint variant in 2d aligns both start and end with minimal source shift- producing the shortest trajectories (Table
Figure 3a summarizes convergence measured by Wasserstein distance over training: joint CAR-Flow converges fastest and reaches the lowest error, followed by source and target-only, all outperforming the baseline. Figure 3bb traces the learned μo and μ4 for each class and variant: joint CAR-Flow yields the most moderate, balanced shifts, explaining its superior convergence and flow quality.
Mode Collapse. To empirically validate the mode-collapse analysis described in Section 3.2 we
reuse the setup but now allow both shift and scale parameters be learned simultaneously (Eq We train two separate models: one reparameterizing the source distribution with learned parameters (No: do), and another morphing the target distribution with (1), Results are presented Figure Specifically, Figure 4a| shows the rapid evolution of the learned standard deviations a, clearly indicating the network quickly discovers a "shortcut" solution by shrinking a zero. Figure 4b plots the expected norm gap El[vg- v. 2| demonstrating convergence to zero, which indicates the network's velocity prediction aligns closely with the analytic zero-cost solutions derived Table Moreover, unrestricted parameterization of the source distribution triggers mode-collapse case (i) as described in Claim I. where the flow degenerates to predicting the class-wise mean the target (see Figure 4c). Conversely, unrestricted parameterization of the target collapses the distribution nearly to a constant, and consequently, the predicted # degenerates into an improper uniform distribution due 01 0, illustrated in Figure 4d
4.2 ImageNet
To benchmark on a high-dimensional, large-scale dataset, we conduct experiments on ImageNet 256 x 256 data using v6e-256 TPUs. Our baseline is SIT-X1/2 [M et al.) 2024]. re-implemented in JAX [Bradbury et al. 2018]: we strictly follow the original training recipe from the open open-source SiT repository to replicate the results reported in the paper. For CAR variants, we apply the shift-only reparameterization to the sd-vae-ft-ema VAE backbone used by SiT (see Eq. (TT- (22). We introduce two lightweight convolutional networks (%= 2.3M parameters each) to predict μi and H4 from the class embeddings, projecting them into the latent space. All models are sampled using the Heun SDE solver with 250 NFEs.
Table presents the quantitative results. Augmenting SiT-XL/2 with CAR-Flow consistently outper- forms the baseline across all variants. In particular, the joint-shift variant achieves the best result, reducing FID from 2.07 to 1.68 while adding fewer than 0.6% parameters. These results underscore the importance of explicitly conditioning the source and target distributions: simple shift reparam-

---

1
Scale a
Source-only 0glA) Target-only (1(A)
Seurce eg(B)
Target-only oy(B)
X1
Pred
Pred
Flow Stage
GT A
0
20
40
60
80
100
GT B
(a) .earned scale
Shift (3d)
Flow traj.
Source-only Target-only
Z0
2
Xo
0
40 60
100
Training Steps (K)
(b) Validation emror
(c) Source-only
(d) Target-only
Figure 4: Mode collapse diagnostics with scale reparameterizatiion (a)(b) Evolution of learned and validation error. (c)-(d) Learned flows when allowing shift+scale on source vs. target.
Table 3: Benchmarking class-conditional image generation on ImageNet 256 x 256. All CAR-Flow variants surpass the SiT-XL/2 baseline.
Method
Params (M) Training Steps FID 4 IS f sFID I Precision
Recall f
SYT-XL/2
675
400K
17.28
78.88
5.71
0.67
0.61
SIT-X1/2 Flow Source only
677
400K
15.15 83.42
5.56
0.68
0.61
SiT-XL/2CAR CAR-Hoo Ao-Hoow Targetoonly
677
400K
14.44
85.66
5.59
0.68
0.61
SiT-XL/2CAR Flow Flow Joint
679
400K
13.91
87.96
5.38
0.68
0.62
SiT-XL/2eff=1.55 JiT-XIL/esR-Foooobbbnneeeell Jointsefg=1.5
675
7M
2.07
280.2
4.46
0.81
0.58
679
7M
1.68
304.0 4.34
0.82
0.62
eterization not only improves sample fidelity but does so with minimal computational overhead, facilitating easy integration into existing large-scale generative frameworks.
Table presents the quantitative results. Augmenting SiT-XL/2 with CAR-Flow consistently outper- forms the baseline across all variants. In particular, the joint-shift variant achieves the best result, reducing FID from 2.07 to 1.68 while adding fewer than 0.6% parameters. Beyond final performance, CAR-Flow also accelerates optimization: our convergence analysis on ImageNet-256 shows that all CAR-Flow variants consistently reduce FID faster than the baseline across training steps (see Fig. 5). These results underscore the importance of explicitly conditioning the source and target distributions: simple shift reparameterization not only improves sample fidelity but does so with minimal computational overhead, facilitating easy integration into existing large-scale generative frameworks.
Related Work
Generative modeling has vdncee significantly the last decade from variational auto-encoders
(VAEs) [Kingma and Welling 2014), Generative Adversarial Nets [Goodfellow et al. [014]. and normalizing flows [Rezende and Mohamed 2015). More recently, score matching [Song and Ermon
2019 Song al. 2020]. diffusion models Ho al. [020], and flow matching [Liu et al. 2023
Lipman et al.) 2023 Allergo and Vanden Vanden-Eijnden 20233 Alberro et al., [2033 was introduced. The be viewed numerically solving
latter three frameworks are related: sampling at test-time
transport (ordinary) differential equation by integrating along learned velocity field from the source
distribution at time
the target distribution at time one.
For learning the velocity field, various approaches to interpolate between samples from the source distribution and the target distribution have been discussed [ipmman et al. 2023 Liu et al. 2023

---

40
Baseline
35
Source-only
Target-only
30
Joint
25
20
15
100000
150000
200000
250000
300000
350000
400000
Training steps
Convergence on ImageNet 256 256: FID vs. training steps. CAR-Flow variants consis- Figure 5:
tently converge faster than the baseline.
Tong et al. 2024). Among those, rectified flow matching was shown lead to compelling results on large-scale data Ma et al. 2024 Esser et al. 2024].
For use on large-scale data, flow matching is typically formulated in latent space by compressing data via the encoder of a pre-trained and frozen VAE [Rombach et al.) 20022.. These mappings differ
from the discussed CAR-Flow, as they are typically independent of the conditioning variable. As mentioned before, CAR-Flow can be applied on top of pre-trained and frozen projections on the latent space.
More recently. Yu et al. 2025
Yao et al. 2025 proposed to align representations within deep nets
that model the velocity visual representations from vision foundation models. Specifically, Yu et al.
[2025] align early layer features of DiT [Pebles and Xie 2023) and SiT (Ma et al. 2024) models with representations extracted from DINOv2 [Oquab al. [204]] and CLIP [Radford al. 2021]
In contrast, Yao et al. (2025) aligns the latent space of VAE with representations from pre-trained
vision foundation models, which are then frozen for diffusion model training. These approaches differ from our approach, which learns to transform the source and target distributions rather than encouraging feature alignment.
Most related to our work is the recently introduced REPA-E [Lng et al. 2025]. In REPA-E, Leng al. (0255) study end-to-end training of diffusion models and VAE encoders/decoders. which map data to/from a latent space. Differently, in this paper, we formalize failure modes reported by) Leng al. 2025 and identify them trivial solutions that arise when jointly training flow matching model and a target distribution mapping. We further introduce source distribution mapping. Finally, we impose simple restrictions that preclude those trivial solutions.
6 Conclusion
We propose and study condition-awer reparameterization for flow matching (CAR-Flow), which aligns the source and target distributions in flow-matching models. We find CAR-Flow to alleviate the burden of classic flow-matching, where model simultaneously transports probability correct region the data manifold, while also encoding the semantic meaning the condition.
the
Limitations and broader impact. This work characterizes the failure modes when jointly training flow matching model, source distribution mapping, and a target distribution mapping, and identifies them as trivial solutions of the objective. We further study the simplest effective approach to avoid these trivial solutions. While we find this simple approach to lead to compelling results, we think more general mappings will likely improve results even further. We leave the identification of more general suitable mappings to future work.
Improving the expressivity of generative models has significant broader impact. On the positive side, modeling complex distributions more easily saves resources and enables novel applications. the negative side, efficient generative modeling can be abused to spread misinformation more easily.
10

---

Acknowledgements
We begin by thanking Bycongjoo Ahn, Wenze Hu, Zhe Gan, and Jason Ren for their thoughtful discussions and rapid feedback, which significantly shaped this study. We also acknowledge the Apple Foundation Model team for providing the critical infrastructure that powered our experiments Finally, we are profoundly appreciative of Ruoming Pang and Yang Zhao for their strategic vision, steadfast guidance, and unwavering encouragement throughout this project.

---

References
M. Albergo and Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In Proc. ICLR, 2023.
M. Albergo, N. Boffi, and E. Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.00 08797, 2023.
Bradbury, R. Frostig. P. Hawkins, M. Johnson, Leary, Maclaurin, Necula, Paszke,
J. VanderPlas, S. Wanderman-Milgg and Q. Zhang. JAX: composable transformations of Python+Nummy programs, 2018. URL htttthttt//gieee
com/ mam/jax-m/// ml/jax
Dhariwal and A. Nichol. Diffusion models beat gans on image synthesis. Advances in neural 8794, 2021.
information processing systems,
P. Esser, S. Kulal, A. Blattmann, R. Entezari, J. Müller, H. Saini, Y. Levi, D. Lorenz, A. Sauer, F. Boesel, al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first international conference on machine learning, 2024.
I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio. Generative adversarial nets. In Proc. NeurIPS, 2014.
J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-68511 2020.
D. Kingma and M. Welling. Auto-Encce Encoding Variational Bayes. In Proc. ICLR, 2014.
X. Leng. J. Singh, Y. Hou, Z. Xing. S. Xie, and L. Zheng. Repa-e: Unlocking vae for end-to-end tuning with latent diffusion transformers, 2025. URL httts:://aaiiiies brggabs//50tl 10483
Y. Lipman, R. Chen, H. Ben-Hamu, M. Nickel, and M. Le.. Flow matching for generative modeling. In Proc. ICLR, 2023.
X. Liu, C. Gong, and Q. Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In Proc. ICLR, 2023.
N. Ma, M. Goldstein, M. S. Albergo, N. M. Boffi, E. Vanden- Eijnden, and S. Xie. Sit: Exploring flow and diffusion-baan based generative models with scalable interpolant transformers. European Conference on Computer Vision, pages 23 40. Springer, 2024.
A. Q. Nichol and P. Dhariwal. Improved denoising diffusion probabilistic models. In International conference on machine learning, pages 8162-8171. PMLR, 2021.
M. Oquab, T. Darcet, T. Moutakanni, H. Vo, M. Szafraniec, V. Khalidov, Fernandez, D. Haziza, F. Massa, A. El-Nouby, M. Assran, N. Ballas, W. Galuba, R. Howes, P.-Y. Huang, S.-W. Li, Misra, M. Rabbat, V. Sharma, G. Synnaeve, H. Xu, H. Jegou, J. Mairal, P. Labatut, A. Joulin, and P. Bojanowski. Dinov2: Learning robust visual features without supervision, 2024. URL http://aaaai...rr////////0//// htw.g//bb/330..000 07193
W. Peebles and S. Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4195-4205, 2023.
A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, Clark, G. Krueger, and Sutskever. Learning transferable visual models from natural language supervision, 2021. htttttt/s//xxiiit http://axxivo.oor/a/a//////o/ara
D. Rezende and S. Mohamed. Variational inference with normalizing flows. In Proc. ICML, 2015.
R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 1669-11695, 2022.
J. Song, C. Meng, and S. Ermon. Denoising diffusion implicit models. In Proc. ICLR, 2021a.
Y. Song and S. Ermon. Generative modeling by estimating gradients the data distribution. In Proc. NeurlPS, 2019.
12

---

Y. Song, J. Sohl-Dicksteing D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score- generative modeling through stochastic differential equations. arXiv preprint arX1v:2011.13456, 2020.
Y. Song, J. Sohl-Dickstein, D. Kingma, Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations. In Proc. ICLR, 2021b.
A. Tong, N. Malkin, G. Huguet, Y. Zhang, J. Rector- Brooks, K. Fatras, G. Wolf, and Y. Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. TMLR, 2024.
Y. Wu, Z. Zhang, J. Chen, H. Tang, D. Li, Y. Fang, L. Zhu, E. Xie, H. Yin, L. et al. Vila- u: a unified foundation model integrating visual understanding and generation. arXiv preprint arXiv:2409.00 244,0499 2024.
Yao, B. Yang, and X. Wang. Reconstruction vs. generation: Taming optimization dilemma in latent diffusion models, 2025. Btt.ppttss://axxii http://axxvvvoooaaaaaa/ abs/abasa111 01423
S. Yu, S. Kwak, H. Jang, J. Jeong, J. Huang, J. Shin, and S. Xie. Representation alignment for generation: Training diffusion transformers is easier than you think. Proc. ICLR, 2025.
13

---

Appendix - CAR-Flow: Condition-Aware Reparameterization Aligns Source and Target for Better Flow Matching
The appendix is organized as follows:
In Section we derive the effects of our reparameterization on the score function. This is important for correct sampling, particularly when using an SDE solver.
In Section B
we prove Claim] [
In Section we provide additional implementation details for both synthetic and ImageNet
experiments.
In Section D we discuss some findings.
In Section we illustrate additional qualitative results.
A Score Function Under Reparameterization
In this section show that
The conditional score 8, Vlogp(z; y) generally changes, once the source-space map is applied, unless [x, v) 2;;
2. For the practically important shift-only map f(z, y) = x + polv) under a Gaussian path, both the density and the score admit closed forms; and
3. The score couples to the drift (velocity field) that appears in the SDE used for sampling.
A.1 Change of the Conditional Score
Recall that the source transform is f: R" x RT (I, + Let (2 R denote its Jacobian. By the thange-of-variable formula,
px(# y) = p2(((r v) |det Jr(
(13)
Taking Ve log of Eq.
13 gives
V
lo y) J5(x,y) y)
V log v) +V log det ( v)
(14)
score 2-space
The first term transports the z-space score back to the x-space tangent via the Jacobian transpose; the second term corrects for the local volume change introduced by f.
Volume-preserving maps. If f is volume preserving, det Jy = 11 and hence V log det Jyl Eq. (14) then reduces to V log px V log Pz. In the special case f(x, y) = x, one has JJ I and therefore v log px Va log pz, recovering the classical setting without reparameterization.
A.2 Shift-only Transform
Assume initial distribution pi = (O, f(x y) = + Mo (v). Along a Gaussian path parameterized by ou, Br (with boundary conditions ag = 0, Bu = and a1 1, B1 = 0), the conditional density at time is
(0, Ia) is standard Gaussian and is a shift-only map, i.e.,
p-(- 2 = N(422) 21 + Sepps 32 La y).
(15)
with both endpoints being
p- | 21, v) = N NNo: Ia v),
pi(- Æ1, 3) = 62110-
(16)
Consequently, by using the form of the Gaussian probability density, the (conditional) score reads
(x | 21,9) = Vz. lo 21.1) = 04 + (.poo 2t
(17)
8,2
14

---

A.3 Link Between Score and Velocity Field
Fix conditioning label y and a pair (20, 21) drawn from the Gaussian endpoint distributions introduced in Section. B.1 The flow is defined as
(/ 2),
Br 20 +
0 0515
(18)
so that Vo = 20 and 1 = 21. Because 1h satisfies the ODE d = ue (1/2 (11 v), differentiating
Eq. [88 in gives the conditional velocity field
u 3.20 2)
3. 20 + de 21.
(19)
Since 2; (z0 21 21.2), 20 (zt 0/211/331 substituting these into Eq. 19 yields
ue (z 21,3) = -Be 21
0x
+ di 21
(20)
31
Eq. 17 links 21 and the conditional score Sc(ze Zt z1, y). Solving Eq. result into Eq. 20 yields
17 for 21 and inserting the
át
88 0t Belr) at
04 + BeHo(())
(o(())
u. (2 21 v)
24 +
(21)
B7
a₁
Se(2e(21.))
bias correction
The first term under the brace is precisely the conditional score from Eq. 17 while the second term compensates for the mean shift po(y) introduced by the shift-only transformation.
Eq. (21 can be rearranged to express the score through the velocity:
a4 u2 21 y)
âi 2t
+ (o(v)
St (z 21,8) v)
(22)
8. a,3,3
B
To translate the conditional identity given in Eq. (22) to the marginal setting used inference time, we integrate over the target endpoint 21. For this we introduce the posterior
p 21 v) p 1z) Pt (= y)
f (r(21 I ze. y) dæl =
(= Re, 2),y)
with
The marginal velocity and score are simply expectations under this density, i.c. ue(2e y) E.g. [u<(z v1. se(2; v) [e 21, v1.
Applying Eq.
22 inside the expectation and using linearity yields
O4 (ze v) 2i 2t
(o(y)
(< (v)
(23)
3,d, ar
Be
the exact coupling used in the latent-space SDE (Eq.8) to express the score with the drift during sampling.
B Proof of Claim
In this section we prove Claim in two steps: (i) We show that any parameter choice driving the objective Eq. (9) to its global minimum forces the velocity field affine the interpolant, 29(2,, t, y) (t, y) 2 + nt, y). (ii) For each collapse pattern (i)-(v) in the claim, we verify that the loss indeed attains this trivial minimum and we state the resulting (t,, y) and 1/(1, y) in closed form.
Step Affine Form at Zero Loss
Let B:00 + 0x If the loss in Eq. 9 vanishes almost surely, the integrand must be identically
zero:
to t,y) B, z0 â ((20.21,t,1)
(24)
15

---

Fix (t. y) and apply the chain rule w.r.t. and 21:
Ovg 2z.
doo
B2 = Br,
Cre-
Ōz.
The right-hand sides are constant in (zo, =1), hence Ova/Oza = (t, does not depend on 2T. Integrating once 2 gives the affine form
(0(#, t, v) = (t, y) 2 +t,
(25)
with ntt,, y) an integration constant. Thus any zero-loss solution necessarily has the affine form given Eq. (25 and already independent of f. This completes Step
Step 2. Each Collapse Pattern Attains the loro-loss Affine Field
Write == [z0, v) and z1 g(v1,)) v). Insert the ansatz given in Eq. 25 into the pointwise identity given in Eq. (24) and substitute 2, = 3.20 + a421:
[t + a (tt v) 3r + di 21
(26)
Eq. (26) is linear in (20. 21); each collapse scenario reduces it to one or two scalar conditions, from which y and 17 obtained explicitly.
(i) Constant source. 2g cly) is fixed while varies freely. Matching the 21 coeeffffiee in Eq. 26 forces dx/Oa, The remaining scalar equation fixes ') = c(y) (3t 7Be).
(ii) Constant target. Symmetric case (i): y / Br and 1) c(y) (d: 7as).
(iii) Unbounded souree scale. As || + do with 21 bounded, the zg-terms in Eq. (26 dominate; finiteness of the left-hand side requires 7/t = 3, = Br/ With this choice the entire identity holds for all (zo, 21) when set 17 O.
(iv) Unbounded target scale. Analogous to (iii) with roles exchanged: 7 = á and 1)
(v) Proportional collapse. Suppose k(y) Substituting into Eq. (26 yields a single free variable 21: [B:k(y) - ag) œ1 [3,k(z) + áæ} 21. Hence y(t, = (B, (v) + (.)/(B,k(y) + a.) and rt, y) 0.
In all five situations 7 and 1y depend only on (1 and the collapse maps (f.g). Consequently the optimizer can reach trivial minimum in which no longer guides a meaningful flow and the generated distribution collapses to a single/improper mode.
C Implementation Details
C.1 Synthetic Data
The velocity network consumes three sinusoidal position embeddings that encode the latent state dy- the class label and time f. Each embedding has dimensionality 8, and concatenating them yields 24-dimensional feature vector. This vector processed by a three-layer MLP whose hidden layers are all 24 - 24 linear projections followed by GELU activations. A final linear layer maps 24 + 1, producing the scalar velocity. The entire podeel-nncludng all embedding 1 993 trainable parameters.
To implement the additive shifts po(y) and (y) we introduce two lightweight condition networks. Each consists of a single linear layer that maps the 8-dimensional class embedding to scalar shift, plus bias term, for parameters per network. Both linear layers are initialized with all weights and biases set to zero, ensuring the additive shifts are identically zero at the start of training.
We train with batch size of 1 022 using AdamW with 31 0.9 and B3 0.95. Learning rates fine-tuned per parameter group: 1 x 10 for the source shift network, 1 x 10-4 for the target shift network, and 10-3 for all remaining parameters. Unless noted otherwise, models are trained for 50k steps; mode-collapse experiments are extended to 100k steps to ensure convergence.
All the synthetic data experiments were executed on the CPU cores of an Apple MI Pro laptop.
16

---

C.2 ImageNet
We re-implemented the open-source SiT code-base in JAX and reproduced the SiT/XL-2 config- uration on ImageNet at 256 x 256 resolution as our baseline. The exact architecture, datapipeline, optimizer (Adam () and learning-rate schedule are identical to the original code. Training per- formed on single v6e-256 TPU slice.
For both source and target CAR-Flows, append lightweight label-conditioning network that maps the 1152-dimensional class embedding to a latent tensor of shape 32 x 32 x 4:
Dense: 1152 + 128 x x 4.
Upsampling: three repeats of Conv'T (kerre 2
+ GroupNorm +
ReLU; shapes 4 x
128 + 8 x 64 - 16 16 x 32 32 x 32 x 16.
Head: x 3 convolution (padding 1) to 32 x 32 x 4, initialized with all weights and bias set to zero ensure the no shifts at the start of training.
Each network contains 2.4M parameters (~~ 0.3% of the SiT/XI.-2 backbone) incurring negligible overhead.
We inherit the original AdamW optimizer for the SiT backbone with learning rate of x 10-4 4 and global batch size of 256. Both label-conditioning networks are trained with a higher learning rate I x 10-1: all other hyper-parameters are unchanged.
D Discussion of Findings
D.1 Relative Learning Rates
During experiments, we discovered that the relative learning rate between the lightweight condition networks (which learn the additive shifts po(y) and (1)) and the velocity- backbone has a first-order effect both the magnitude of the learned shifts and the convergence speed-a "race condition" between them.
To study this, we reuse the synthetic-data example: we fix the backbone's 's learning rate at 1 x 10-5 and train source-only and target-only CAR-Flow models while sweeping the shift-network learning rate across three orders of magnitude (1 x 10-5, 1 x 10-1, and 1 x 10-5). Figures 6a and 6c plot the Hi and μl trajectories for the source-only and target-only variants, respectively. At the smallest rate, the shifts remain near zero and the backbone carries the conditioning, yielding slower alignment: the intermediate rate, the shifts grow steadily-though they can still pulled by the backbone (c.g.. Moly B) starts positive but gradually becomes negative, reducing [oO A) po B)); and at the largest rate, the shifts rapidly attain larger magnitudes and drive the quickest convergence. Figures (6b and [dd report the Wasserstein distances, confirming that the highest shift- network rate achieves the fastest distribution alignment.
These observations suggest that empowering the shift networks with higher learning rate can substantially accelerate alignment. practice, one should choose a shift-network rate that sufficiently high to speed convergence while preserving robust inter-class separation.
To verify that these trends extend to large-scale data, we performed the same learning-rate sweep on ImageNet using the SiT-XL//2 backbone. We fixed the backbone's 's learning rate at its default x 10 and trained the CAR-Flow condition networks under three configurations -ource-only. target-only, and joint- wwile varying their learning rate from 1 x 10-3 to x 10 Figure 7 shows the FID at 400K steps for each variant, with SiT-XL/ baseline of 17.2 indicated by the dashed line.
At the lower condition-network rate (~ 10-3), source-only and target-only variants remain on-par worse than the baseline, whereas the joint variant already outperforms it. As the rate increases, all three variants deliver substantial improvements-taan teameentt-targe-onnll achieves the best FID of 12.77 at 2 x 10-2, and the joint variant exhibits the most stable performance across the sweep. Even the highest rate (10-1), all configurations remain well below the baseline.
These large-scale results mirror our synthetic-data findings: increasing the accelerates alignment and improves sample quality, although the exact optimum varies by variant and requires some tuning. For simplicity and robust performance, we therefore adopt rate of 1 x 10-
learning rate
17

---

I x 1
1 10
Wsme== x 10
Class
Class B
Wasserstein Distance
10a
10
Hl
0
10
20
30
40
50
40
50
Training step (K)
Training Steps (K)
(a) Learned μi shift for source-only
(b) Wasserstein distance for source-only
Wasserstein Distance
10
M1 shift
10-1
0
10
20
30
40
50
10
20
30
40
50
Training step (K)
Training Steps (K)
(c) Learned μn shift for target only
(d) Wasserstein distance for target-only
Figure 6: Effect of varying the shift-network learning rate relative fixed backbone of x 10-5 Panels (a) and (c) show the trajectories of and μn for three learning rates; panels (b) and (d) plot the corresponding I-D Wasserstein distances. At the lowest shift-network rate, the learned shifts remain negligible (slow alignment); at the intermediate rate, they grow steadily without instability; and the highest rate, they overshoot then damp, yielding the fastest overall convergence despite early oscillations.
18
FID @400k Steps
Source only
16
Target-only
Joint SiT/XL-2 Baseline
14
10-3
10-2
10
Condition Network Learning Rate
Figure 7: ImageNet FID at 400K steps vs. condition-network learning rate for SiT-XL/2 CAR-Flow variants: source-only (blue), target-only (orange), and joint (green). The dashed horizontal line marks the bascline FID of 17.2. All variants improve over the baseline as the shift-network rate increases.
in our main experiments. Figure presents example outputs from each CAR-Flow variant at this rate, demonstrating that the joint variant attains superior visual fidelity compared to the source-only and larget-only models.
D.2 Condition-aware source versus unconditional shift
Complementary to the observations above, we examine whether making the condition-awant is beneficial compared to using single unconditional shift. For intuition, consider the source-only CAR-Flow variant: it reparameterizes the source with shift that depends on y. When the target distribution varies with the conditioning variable, aligning the source per condition should reduce transport effort relative to an unconditional source [Albergo et al. 2023
We test this in the 1-D synthetic setup of Sec. 4.1 by comparing (i) learnable unconditional source with a global shift and (ii) condition-aware source with a y-dependent shift. The Wasserstein distances are 0.058 for (i) versus 0.041 for (ii), indicating that per-condition alignment reduces transport. This supports CAR-Flow's design choice and clarifies its distinction from an unconditional "learnall source".
18

---

Source-only
Joint Target-only
Figure 8: Qualitative ablation of CAR-Flor variants on SiT-XL/2 at 400K steps. Each sponds to one variant variant-(top) source-only, (middle) target-only, and (bottom) joint CAR-Flow and each column shows generated sample for different class from the same noise using cfg 1. The joint model produces the most realistic and semantically accurate images across all scenarios.
Table 4: Class-conditional image generation on CIFAR-10 10 (FID ().
Source-only
Target-only CAR-Flow
Joint CAR-Flow
Baseline
CAR-Flow
FID +
13.8
7.5
11.1
10.6
E
Additional Results on CIFAR-10
To assess generalization beyond ImageNet, trained SiT-XL/2 baseline and CAR-Flow variants for 400k steps on CIFAR-10 using pixel-space diffusion (VAE omitted due to CIFAR-10's low resolution 32 x 32). All CAR-Flow variants outperformed the baseline, demonstrating that the benefits of CAR-FI Flow generalize across datasets.
19

---

Qualitative Results
Qualitative results are provided in Figure g
Figure 9; Randomly selected samples generated from our SiT-XL/2CAR-FIN Flow Joint model trained on ImageNet 256 x 256 data using cfg 4.
20

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our three contributions are summarized at the end of Section I Points 1 and are detailed in Section while point 3 detailed in Section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The limitations are discussed in Section 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
21

---

Justification: We have two claims (Claim I and Claim 2). The proof of Claim is provided in the Appendix B
The proof of Claim provided in Section 3.3
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
Implementation details are provided in Section 4 and Appendix C
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
22

---

Answer: [No]
Justification: Implementing the proposed approach is straightforward. We follow prior work w.r.t. papae--ar settings.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm ec/
public/guides https//oooeouuuossoooooooooooooo
for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
//nips.cc/pp cc/public/gut
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Details provided in Section 4 and Appendix C
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Details are provided Section
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
23

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Information is provided in Section 4 and Appendix C
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics http:p///..eerr....cc//////////a
Answer: [Yes]
Justification: We verified that the research conducted in the paper conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Broader impacts are discussed Section 6 Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Our contribution is mostly theoretical. Direct risk of misuse is minimal.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Baselines and datasets are properly cited and licenses properly respected.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
25

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
No new assets will be released.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details of the dataset/cooe/mmodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
The reported research does not involve crowdsourcing nor research with human subjects.
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The reported research does not involve crowdsourcing nor research with human
subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
26

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: Method development in this research does not involve LLMs as any important, original, or non-standard component.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components. ///Coonferences/22222
Please refer to our LLM policy for what should or should be described.
27


Reviews:
[{"review_id": "D1iKooEvvE", "reviewer": "Reviewer_naoP", "summary": "* The paper introduces Condition-Aware Reparameterization (CAR) which is a new way of doing conditional generative modeling. CAR explicitly incorporates conditioning information into generative models via learned, shift-only reparameterizations of either the source or target distributions. This simplifies the training of flow-based generative models and enhances their sampling efficiency.\n\n* The authors provide theoretical insights, demonstrating that allowing certain families of reparameterizations can lead to degenerate, trivial solutions. By restricting reparameterizations to shift-only transformations, CAR avoids these degenerate cases.\n\n* In the synthetic dataset example, trajectory lengths being shorter does not necessarily imply that generation quality improves. Indeed, the CAR method can be seen as decomposing the trajectory into different stages x_0 -> z_0 -> z_1 -> x_1, so naturally the length of the z_0 -> z_1 path would be shorter. However, by measuring the Wasserstein divergence of the generated distribution and the ground truth, the authors verify that the quality of the generated samples is better when using CAR.\n\n* Empirical validation shows CAR reduces trajectory lengths and improves sampling quality. Specifically, it achieves gains on ImageNet generation benchmarks (e.g., improving FID from 2.07 to 1.68 on ImageNet-256) with minimal additional parameters.", "questions": "* Do you have any examples (theoretical or empirical) where learning the distribution for y performs better than just learning it for x_0? (also see first weakness above)\n* You mention that your method decouples the “transporting mass” and “encoding semantics.” However, I don’t see any reason for it being problematic to learn these two behaviors at the same time. Is there any intuition you could provide for what you have in mind here?\n* Do you have any intuition on why the target-only CAR leads to paths with smaller lengths and lower FIDs than the source-only CAR?\n* The meaning of cases (iii) and (iv) is not very clear from what is written in Claim 1.", "limitations": "Yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths\n* The authors do a good gob laying out the foundations and formalizing their framework. The presentation overall is mostly clear.\n* The paper provides both a simple demonstration of their proposed algorithm (CAR) on synthetic dataset. This is useful to understand that CAR reduces the trajectory length in latent space, potentially leading to less error due to discretization. On the other hand, the authors also show that CAR improves FID performance on ImageNet, which is evidence of the practical use of this method.\n\nWeaknesses\n* Although the theoretical claims (Claim 1 and Claim 2) are interesting, they are mostly characterizing degenerate behavior. This leaves the reader without a sense of what are optimal \\mu_0 and \\mu_1 in non-degenerate cases. Understanding theoretically the case where the target is a two-class Gaussian mixture would be ineresting. If that is too hard, a simplification would be to understand the optimal \\mu_0 for the case where the target is a symmetric sum of diracs (which I believe doesn’t match the triviality conditions from Claim 1)\n* It is known in the literature that changing the source distribution can lead to lower transport costs (e.g. Albergo and Vanden-Eijnden 2023). Instead of learning a distribution for the y, we could just learn a distribution for the x_0 instead of  using the N(0, Id). Note that there are ways of doing this using only a few parameters like still parameterizing x_0 by a Gaussian, but learning the mean and covariance (where we could assume the covariance is diagonal if we want to save parameters.) Although this wouldn’t decouple “transporting mass” and “encoding semantics”, it is not clear that this is needed, and having a learnable distribution for x_0 may already explain the performance increase.", "was_revised": true, "final_justification": "The authors did a good job clearing up a few questions and confusions I had about the paper. Thanks to the authors' answer I have updated the significance score and the overall score.", "rebuttal": "Thanks a lot for your time and valuable feedback. We appreciate the reviewer highlighting that the paper does “a good job laying out the foundations and formalizing the framework”. We are glad to see that the reviewer appreciates that we show a combination of both simple demonstrations and practical impact. Below we address each of the reviewer’s comments and questions in turn.\n\n----\n\n- $\\textbf{Reviewer comment (W1)}$: Claims 1 & 2 mostly rule out degenerate cases; they don’t tell us what the optimal $\\mu_0$, $\\mu_1$ are in non-degenerate cases.\n\nWe appreciate this insightful observation. We would like to point out that there is generally no unique optimal $\\mu_0$, $\\mu_1$. Instead they live in a family of equally valid solutions because they are jointly learned with the velocity field $v_{\\theta}$. \n\nConsider the simplest 1-D example, where the source and target are Dirac masses at $x_0$ and $x_1$. In the case of CAR (with source shift learned), the optimal velocity is $v^*=x1-(x0+\\mu_0)$. Hence, for any choice of $\\mu_0$, there exists a corresponding zero-cost velocity field. Once training finds any member of this family, the gradient on $\\mu_0$ vanishes. In practice, its final value is governed by many factors (learning-rate, model architecture, flow complexity, etc.). \nWe will clarify this point and include the example above in the revision.\n\n----\n\n- $\\textbf{Reviewer comment (W2/Q1)}$: Why not simply learn an unconditional source distribution $p(x_0)$ (Albergo & Vanden-Eijnden, 2023)? Do you have any evidence that learning the distribution for $y$ is actually better?\n\nTo make sure, note that we don’t learn a distribution for the condition $y$; rather, we learn a re-parameterized source that depends on the condition $y$. Intuitively, when the target distribution varies with the conditioning variable, aligning the source per condition should reduce transport effort beyond what a single unconditional source can offer. \n\nTo validate this, we ran a controlled experiment using the Sec. 4.1 1-D synthetic experiment setup that compares\n\n(i) a learnable $\\textbf{unconditional}$ source $p(x_0)$ with shift, \n\n(ii) a $\\textbf{condition-dependent}$ source $p(x_0\\mid y)$ with shift dependent on y. \n\nWe got Wasserstein distance of (i) 0.058 vs. (ii) 0.041. This supports our claim: aligning the source per condition leads to reduced transport effort. We will add this result to the final version and clarify the distinction between an unconditional “learnable source” and CAR’s condition-aware variant.\n\n----\n\n- $\\textbf{Reviewer comment (Q2)}$: Intuition for learning “transporting mass” and “encoding semantics”?\n\nThank you for this question. We do not mean to imply that a single velocity field $v_{\\theta}(x_t, t, y)$ cannot learn both mass transport and semantic injection simultaneously. Rather, our intuition and goal is to simplify what the flow is asked to model:\n1) In standard flow matching, a single, condition-agnostic source must be carried to potentially very different target manifolds for each $y$. As a result, the network must learn both long-range transport and the precise semantic shifts encoded by $y$, which can lead to intricate trajectories and slower convergence.\n2) In CAR, instead of solely relying on $y$ to inject the semantics, we explicitly reparameterize the source (or target) distribution using condition $y$. This can make the resulting flow trajectory simpler because the source distribution depends on the condition $y$—transporting mass becomes easier to learn.\n\nWe hope this clarifies our intuition.\n\n----\n\n- $\\textbf{Reviewer comment (Q3)}$:  Performance difference between target-only and source-only CAR?\n\nNote that neither source-only nor target-only CAR consistently outperforms the other in our experiments. While both variants reliably improve upon the unconditional baseline, we find their relative advantage is small and depends on the dataset and hyper-parameters:\n1) On 1-D synthetic data (Sec. 4.1), source-only CAR yields a very slightly lower Wasserstein distance than target-only (Fig. 3(a)).\n2) On ImageNet, Table 3 shows target-only CAR delivers a small FID improvement over source-only. However, in our learning-rate ablation (Fig. 6, Appendix), neither variant shows a statistically significant edge.\n\n\n----\n\n- $\\textbf{Reviewer comment (Q4)}$: Meaning of cases (iii) and (iv)?\n\nCases (iii) and (iv) in Claim 1 are theoretical degenerate \"collapse\" solutions that can be derived from Eq. (26) in the Appendix and require unbounded weights. When the scale of either the source or target distribution tends to infinity, it causes the counterpart distribution to appear infinitesimally narrow in comparison—effectively collapsing it to a point, leading to zero-cost loss. As noted in L164–165, such unbounded weights are not realistic in practice, so these solutions are unstable—any small perturbation causes the optimization to revert to cases (i) or (ii). We included them in Claim 1 for a theoretically comprehensive discussion and will add a remark in the revised manuscript to emphasize this."}, {"review_id": "V6LHPJt07O", "reviewer": "Reviewer_rueW", "summary": "The authors in this paper propose Condition-Aware Reparameterization (CAR) as a method for imprving conditional generative models especially in diffusion and flow-matching frameworks. Standard conditional generation methods rely on a standard Gaussian prior, forcing the model to handle both semantic conditioning and distribution matching, leading to inefficient, long probability paths during training. Proposed method mitigates this by shifting the source and/or target distributions in a condition-aware manner aligning distributions to reduce the burden on the main model. The authors provide theoretical analysis demonstrating that unrestricted reparameterizations can cause mode collapse. By restricting to additive shifts only, CAR avoids these trivial minima while maintaining benefits in path shortening. Experiments on synthetic data and ImageNet-256 show that CAR consistently improves convergence speed and sample fidelity.", "questions": "1. How does CAR perform when conditioning signals are noisy or weak?\n2. Have you explored combining CAR with various transformations under constraints?\n3.Can you discuss how this will extend to models like FFJORD or neural autoregressive flows?\n4.Some ablation study results will be appreciated.\n5. Can CAR be extended to multi-modal or multi-task settings?", "limitations": "The authors can discuss the limitations in some more detail, covering multi modal modeling, more flow models especially like FFJORD and neural autoregressive flows, cases when conditioning signal is weak, transformation under constraints, etc.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The paper proposes an interesting solution for reducing the complexity of conditional flow matching by aligning distributions via learnable shifts.\n2. The theoretical analysis of the paper is well complimented by empirical analysis.\n3. CAR is easy to integrate into existing frameworks, adding minimal overhead, and provides consistent improvements in convergence and sample quality.\n\nWeaknesses:\n1. The method is currently limited to shift-only reparameterizations.\n2. Detailed ablation study should be performed to quantify contribution of various components.\n3. Empirical validation focuses on FID and ImageNet; more diverse datasets and modalities (e.g., text, audio) are not explored.\n4. How CAR works for models such as FFJORD or neural autoregressive flows is yet to be covered. (Grathwohl, Will, et al. \"Ffjord: Free-form continuous dynamics for scalable reversible generative models.\" arXiv preprint arXiv:1810.01367 (2018)., Huang, Chin-Wei, et al. \"Neural autoregressive flows.\" International conference on machine learning. PMLR, 2018.)", "was_revised": true, "final_justification": "Based on the authors rebuttal I am upgrading my score.", "rebuttal": "Thanks a lot for your time and valuable feedback. We are glad the reviewer finds our solution to be “interesting”, the “theoretical analysis of the paper” to be “well complemented by empirical analysis”, and our method to be “easy to integrate into existing frameworks” while “adding minimal overhead” and while leading to “consistent improvements in convergence and sample quality”.  Below we address each of the reviewer's comments and questions in turn.\n\n\n----\n\n- $\\textbf{Reviewer comment (W1)}$: Can CAR go beyond shift-only reparameterizations?\n\nWe appreciate the suggestion. CAR focuses on shift-only reparameterizations because they are the simplest volume-preserving transforms. Allowing more expressive transformations—especially ones that can change volume—introduces the risk of degenerate solutions (as characterized in Claim 1). By constraining ourselves to shifts, we sidestep such instability while still observing strong empirical gains with negligible overhead. Exploring richer yet stable, volume-preserving reparameterizations is an exciting direction for future work.\n\n----\n- $\\textbf{Reviewer comment (W2/Q4)}$: Request for ablation studies.\n\nIn the current manuscript, we have already conducted ablations to isolate CAR’s components:\n\n1) CAR variants (Source-only, target-only, and joint CAR): Fig. 2 & Table 2 (synthetic data); Table 3 (ImageNet).\n\n2) Learning-rate ablations: Appendix Fig. 5 (synthetic data) and Fig. 6 (ImageNet).\n\nIf there are additional components the reviewer would like us to examine, please let us know. We are happy to incorporate any further ablations in the revision.\n\n\n\n----\n- $\\textbf{Reviewer comment (W3/Q5)}$: Limited dataset diversity—only ImageNet/FID.\n\nThank you for this suggestion. While ImageNet remains a widely used benchmark for high-dimensional generation, we additionally ran a new ablation on CIFAR-10 to assess generality.\n\nWe trained a SiT/XL-2 baseline and CAR variants for 400k steps, using pixel-space diffusion (VAE omitted due to CIFAR-10’s resolution). Results:\n\n\n| Model            | FID  |\n|-------------------|------|\n| Baseline          | 13.8 |\n| Source-only CAR   | 7.5  |\n| Target-only CAR   | 11.1 |\n| Joint CAR         | 10.6 |\n\n\n\nAll CAR variants outperform the baseline, indicating generalization beyond ImageNet. We believe CAR can also extend to other modalities (e.g., text, audio), but compute constraints prevent us from including such experiments in this submission.\n\n----\n- $\\textbf{Reviewer comment (W4/Q3)}$: Can CAR be applied to FFJORD or neural autoregressive flows?\n\nWe see no fundamental barrier to applying CAR in frameworks like FFJORD or neural autoregressive flows. However, each new setting may introduce its own degenerate collapse modes (analogous to the ones characterized in Claim 1), and a detailed theoretical and empirical study is required to identify and address those cases. \n\n----\n- $\\textbf{Reviewer comment (Q1)}$: How does CAR perform under weak or noisy conditioning?\n\nIn conditional flow-matching models, when the conditioning signal $y$ is weak or noisy, both the baseline and CAR-augmented models are affected—there is simply less reliable information to guide the learned velocity field $v_{\\theta(x_t, t, y)}$. This limitation is not unique to our approach but is inherent to all conditional generative methods. Improving robustness to weak or noisy conditioning is an interesting but orthogonal direction to our contribution."}, {"review_id": "JNGlFFHpCZ", "reviewer": "Reviewer_n8K3", "summary": "The paper proposes CAR, condition-aware reparametrization, a technique to train source and target transformations simultaniously with the velocity field in the flow matching frameworks. This approach allocates some capacity for processing the conditioning separately from the main flow matching objective, which fosters the alignment between the source and the target distributions, and thereby eases the optimization of the velocity field. The authors provide theoretical analysis of the degenerate solutions in the general unresticted settings and derive the shift-only CAR to avoid those. Experiments on synthetic and real-world data support the claims made in the paper and demonstrate superior performance compared to the prior work.", "questions": "1. While the method is motivated for the conditional case, I think the unconditional generation can also benefit from the proposed source/target alignment. Could the authors provide any insights on why the proposed approach should/sould not work for the unconditional setting?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strengths\n1. The paper is well-written and the motivation is clear. The method is simple and theoretical foundations are provided. The results on both synthetic and real data support the claims made in the paper.\n\n### Weaknesses\n\n1. The main weakness of the paper is the lack of clarity in delivering the theoretical results. It is not clear why the degenerate 0-cost solutions are presented as \"$\\theta$-independent\". These are analytical solutions that exist in the general setting, but if the $\\theta$-depended model converged to these during the optimization process they still depend on parameters. \n2. The proportional collapse case is unclear. In the loss (equation 9), $x_0$ and $x_1$ are sampled independently from each other. However, case *(v)* somehow suggests that $z_0$ and $z_1$ are coupled. Could the authors clarify this?\n3. The unbounded case also requires further clarifications. In general, the 0-cost solution of the optimization problem in equation 9 is feasible if the are no intersections between the trajectories that interpolate $z_0$ and $z_1$. While in cases *(i)* and *(ii)* this is clear, in the case where both the target and the source distributions have non-trivial supports, there will be intersections, unless the source and the target are non-trivially coupled. I would like to see more discussion of this issue in the cases *(iii)-(v)*.", "was_revised": true, "final_justification": "The paper addresses an interesting problem of jointly training source and target transformations along with the velocity field in flow matching models. The proposed approach and results appear sound. The rebuttal addressed most of my concerns. However, I believe the clarity of the theoretical section could still be improved (in line with the discussion with the authors), hence my rating.", "rebuttal": "Thanks a lot for your time and valuable feedback. We appreciate that the reviewer finds the paper to be \"well-written and the motivation is clear\". We are glad to see that the reviewer appreciates the provided theoretical foundations and acknowledges that the results on both synthetic and real data support our claims. Below we address each of the reviewer's comments and questions in turn.\n\n----\n\n- $\\textbf{Reviewer comment (W1)}$: Clarify what is meant by \"$\\theta$-independent\" degenerate solutions.\n\nThank you for highlighting this confusion. We agree that our phrasing was unclear. Our intention was to indicate that the degenerate solutions $v^*$ yield zero cost analytically—by construction—when paired with corresponding source/target shifts. However, the learned velocity field $v_{\\theta}$ must still approximate this ill-conditioned optimum via parameter $\\theta$. In that sense, the solution is not truly $\\theta$-independent. We will revise the wording accordingly in the final version.\n\n----\n\n- $\\textbf{Reviewer comment (W2/W3)}$: Clarify the unbounded collapse cases and the proportional collapse case.\n\nWe appreciate the opportunity to clarify. Cases (iii)–(v) in Claim 1 describe purely theoretical collapse solutions derived from Eq. (26) in the Appendix:\n\n1) Cases (iii) and (iv) (unbounded collapse): When the scale of either the source or target distribution tends to infinity, the counterpart distribution collapses relative to it. This leads to minimal trajectory overlap between $z_0=f(x_0)$ and $z_1=g(x_1)$, resulting in near-zero cost.\n\n2) Case (v) (proportional collapse): This case assumes exact proportionality between $z_0$ and $z_1$. As the reviewer correctly notes, this cannot occur when $x_0$ and $x_1$ are sampled independently—except in trivial cases where both $z_0$ and $z_1$ are constants, which reduces to cases (i) or (ii).\n\nAs discussed in L164–165, these collapse modes do not arise in practice. Cases (iii) and (iv) require unbounded weights, and are therefore unstable—any  perturbation pushes optimization causes the optimization to revert to cases (i) or (ii). Case (v) requires implausible, perfectly aligned reparameterizations under independently sampled inputs. We included these cases for theoretical completeness and will revise the manuscript to clarify their pathological nature.\n\n----\n\n- $\\textbf{Reviewer comment (Q1)}$: Can the proposed method benefit unconditional generation as well?\n\nYes--thank you for raising this point. We believe the same source/target reparameterization can benefit unconditional settings. Even without explicit conditioning, allowing learned shifts in the source and/or target distributions gives the model more flexibility to reduce transport cost, thereby simplifying the velocity field it must learn. A related idea—adjusting the prior to ease conditional transport—was briefly mentioned by Albergo & Vanden-Eijnden (2023), but not studied in detail."}, {"review_id": "GNFwfbNBbq", "reviewer": "Reviewer_zCmc", "summary": "This work introduces CAR, a technique to improve conditional flow matching models by learning a transformation of the latent and/or data distribution that depends on the condition. The flow matching model is then applied to map between these two transformed distributions. More specifically, the authors first identify five classes of degenerate zero-cost solutions of flow matching that can be achieved when using arbitrary transformations. They then notice that it is impossible to create these solutions with only a shift of both distributions, and thus propose to include a learned shift before and after the mapping induced by flow matching. They then verify in a toy example that the proposed shift transformation learns a non-degenerate solution, whereas including scaling also leads to one of the zero-cost solutions in practice. Finally, they demonstrate that even this simple addition to the FM framework yields significant improvements on conditional ImageNet 256x256.", "questions": "- Do you have a notion of how complete the list of failure modes in Claim 1 is? Are there other problems beyond mode collapse that could occur?\n- Have you tried to apply CAR for continuous conditions? Can you observe similar benefits?\n- Have you tried more general transformations (e.g. volume-preserving transformations would seem to avoid all the trivial solutions as well)?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**  \n- The presented method is simple and seems applicable to most conditional flow matching scenarios\n- The results achieved on ImageNet 256x256 are very convincing\n- The identified trivial minima are not only relevant for this work but might help to illuminate problems that occur in adjacent work which performs joint training of flow matching and pre-/post-hoc transformations.\n- The method is very lightweight and seems to have almost no implementational and computational overhead  \n- The writing is very clear, proofs are formulated understandably and the toy example is quite instructive.\n\n**Weaknesses**\n- The method's quantitative and qualitative performance was only compared against baseline flow matching on a single dataset. While the results are convincing there is still room to explore how it performs for different dataset and condition types\n- The improvements in convergence speed were only shown for the 1-d toy example. It would be interesting to see if improved convergence speed is also observed on high-dimensional data.\n- The authors only consider the simplest transformation that avoids the identified trivial solutions. More complex transformations might result in much better performance \n\nMinor things:\n- The qualitative results in the appendix lack a baseline comparison\n- The shift prediction network learning rate needs to be fine-tuned to see the beneficial effects of CAR (but according to Fig. 6 it seems very lenient)\n\nIn conclusion, while the extent of performance gain that can be expected from CAR (and possibly more advanced versions) might not be quantified entirely from the limited experiments in this paper, due to the mentioned strengths it can be expected to overall contribute positively to CFM training.", "was_revised": true, "final_justification": "This paper presents a simple, universally applicable method to improve the conditional flow matching framework. Being one of the most relevant frameworks in conditional generative modelling, this method can be expected to have a significant impact on the field of generative modelling as a whole, which is supported by solid (but not exhaustive) theoretical and experimental evaluation. I do not rate this work higher (i.e. groundbreaking) as it only considers a first analysis of the simplest settings (one type of data/condition in the experiments, only the simplest transformation that avoids the identified trivial solutions).", "rebuttal": "Thank you for your time and thoughtful feedback. We’re excited to see that the reviewer found our method “applicable to most conditional flow matching scenarios,” our ImageNet results “very convincing,” our identified trivial solutions broadly relevant, and our approach both lightweight and clearly presented. We address the reviewer’s questions and suggestions below.\n\n----\n\n- $\\textbf{Reviewer comment (W1)}$: Performance on additional datasets.\n\nThank you for the suggestion. To assess generalization beyond ImageNet, we conducted an additional experiment on CIFAR-10. We trained a SiT/XL-2 baseline and CAR variants for 400k steps, using pixel-space diffusion (omitting the VAE due to CIFAR-10’s resolution). Results are shown below:\n\n| Model | FID | \n|-------------------|------| \n| Baseline | 13.8 | \n| Source-only CAR | 7.5 |\n| Target-only CAR | 11.1 | \n| Joint CAR | 10.6 |\n\n\nAll CAR variants outperformed the baseline, demonstrating that the benefits of CAR generalize on different datasets.\n\n----\n\n- $\\textbf{Reviewer comment (W2)}$: Convergence speed on high-dimensional data.\n\nThank you for pointing this out. We conducted a new convergence analysis on ImageNet-256 using the SiT-XL/2 backbone. The table below reports FID over training steps, showing that CAR variants consistently converge faster than the baseline:\n\n\n| **Steps** | **Baseline** | **Source-only** | **Target-only** | **Joint** | \n|-----------|--------------|-----------------|-----------------|-----------| \n| 50k | 75.76 | 74.44 | 74.11 | 71.18 | \n| 100k | 40.60 | 39.96 | 40.59 | 36.60 |\n| 150k | 30.15 | 27.57 | 28.52 | 26.33 | \n| 200k | 24.66 | 22.59 | 22.69 | 21.38 |\n| 250k | 21.30 | 19.72 | 19.73 | 18.16 | \n| 300k | 19.11 | 17.72 | 17.31 | 16.19 | \n| 350k | 17.46 | 16.24 | 15.76 | 15.08 | \n| 400k | 16.28 | 15.15 | 14.44 | 13.98 |\n\n\nWe will add this table (and the corresponding learning curves) to the revision to document CAR’s convergence advantage on high-dimensional data.\n\n\n----\n\n- $\\textbf{Reviewer comment (W3/Q3)}$: Exploring more complex transformations.\n\nThank you for raising this excellent point. We focused on shift-only reparameterizations as they are the simplest volume-preserving transforms that already yield measurable benefits with minimal overhead. We agree that richer—but still volume-preserving—mappings (e.g., learned orthonormal transforms) could potentially deliver further improvements while avoiding the trivial minima identified in Claim 1. Systematically exploring these more expressive options is a promising next step.\n\n----\n\n- $\\textbf{Reviewer comment (W3)$: Missing baseline samples in qualitative results.\n\nWe will add side-by-side qualitative comparisons with the baseline in the appendix of the revised manuscript. Although differences are subtle at later training stages (e.g., 7M steps with classifier-free guidance), they are more noticeable earlier in training. For instance, CAR-generated Chihuahua samples at 400k steps (without classifier-free guidance) exhibit stronger structural coherence than those from the baseline. Due to the text-only rebuttal format, we cannot include images here, but we will make them available in the final version.\n\n----\n\n- $\\textbf{Reviewer comment (W4)}$: Lenience to shift-prediction network learning rate.\n\nIndeed, in practice we found CAR settings to be quite lenient: its gains hold across a wide span of shift-network learning rates, so only coarse tuning is needed.\n\n----\n\n- $\\textbf{Reviewer comment (Q1)}$: Completeness of failure modes in Claim 1.\n\nThank you for the question. Cases (iii)–(v) in Claim 1 are theoretical degenerate “collapse” solutions that can be derived from Eq. (26) in the Appendix. As noted in L164–165, they require extreme conditions that do not arise in practice. While additional failure collapse modes may exist in theory, we believe the practically relevant ones are already captured by cases (i) and (ii), which we observe empirically (see Fig. 4). In practice, beyond these mode-collapse scenarios, we have not encountered other failure modes in our experiments.\n\n----\n\n- $\\textbf{Reviewer comment (Q2)}$: Applicability to continuous conditions.\n\nWe have not yet run experiments with continuous‐valued conditions. Conceptually, CAR should extend straightforwardly—for instance, by modeling $P(\\mu | y)$ where  $y$ is continuous. We expect similar benefits in such settings, though verifying this would require additional experimentation beyond the scope of this submission."}]

Author rebuttal:

Meta-review: (a) This paper introduces Condition-Aware Reparameterization (CAR), a method to improve conditional generative modeling in flow-matching and diffusion frameworks. CAR introduces condition-dependent shifts of the source and/or target distributions, thus reducing trajectory lengths. The authors provide a theoretical analysis that identifies trivial zero-cost solutions under general reparameterizations and show that restricting to shift-only avoids these degeneracies. Experiments on synthetic data and ImageNet-256 demonstrate that CAR improves sample quality (e.g., they reduce FID from 2.07 to 1.68) with minimal computational overhead, while also shortening generative paths.

(b) The proposed method is clear and simple and the paper is well-written. CAR is easy to implement, adds minimal overhead, and integrates naturally into existing conditional flow-matching setups. The authors provide theoretical insights, demonstrating that allowing certain families of reparameterizations can lead to degenerate, trivial solutions. They show that, by restricting reparameterizations to shift-only transformations, CAR avoids these degenerate cases. This work has practical impact. The authors demonstrate experimentally that the method achieves strong improvements on a competitive ImageNet benchmark. 

(c) The reviewers raised several concerns. One limitation is in the scope of the experiments. Experiments are mainly on ImageNet-256, and performance on other datasets and modalities (eg, text and audio) remains unexplored. The presentation of the theoretical analysis needs more work. Some parts of the analysis (eg, degenerate solutions, proportional collapse, independence assumptions) were found confusing or underspecified. The reviewers pointed out missing ablations and comparisons. 

(d) The method is simple, broadly applicable, and well-motivated both theoretically and empirically. Despite the issues on incomplete experiments and the unclear theory presentation, the proposed method is found valuable. The performance improvement on ImageNet-256 is significant. Reviewers consistently valued the clarity, theoretical framing, and practical relevance, with two recommending full acceptance and the others borderline accept. The main limitations leave room for future work but do not undermine the contribution. The paper is a candidate for spotlight, as it presents an approach that would be of interest to a broad audience.

(e) The main points raised concerned: 1) the scope of experiments (only ImageNet and toy data), 2) the clarity of theoretical results on degenerate solutions, 3) whether unconditional generation could benefit, and 4) whether more general transformations could be explored. The authors provided clarifications during rebuttal, which reviewers felt addressed most of the concerns. Several reviewers raised their scores (e.g., naoP updated significance and score upward, rueW upgraded to accept, n8K3 confirmed rebuttal addressed concerns). The consensus converged to acceptance, with reviewers acknowledging CAR as a simple and effective addition to conditional flow matching despite limited experimental breadth.
Decision: Accept (spotlight)