CAR-Flow: Condition-Aware Reparameterization Aligns Source and Target for Better Flow Matching

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.

• $**Reviewer comment (W1)**$: Performance on additional datasets.

Thank 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:

All CAR variants outperformed the baseline, demonstrating that the benefits of CAR generalize on different datasets.

• $**Reviewer comment (W2)**$: Convergence speed on high-dimensional data.

Thank 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:

We will add this table (and the corresponding learning curves) to the revision to document CAR’s convergence advantage on high-dimensional data.

• $**Reviewer comment (W3/Q3)**$: Exploring more complex transformations.

Thank 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.

• $**Reviewer comment (W3)$: Missing baseline samples in qualitative results.

We 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.

• \textbf{Reviewer comment (W4)**: Lenience to shift-prediction network learning rate.

Indeed, 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.

• $**Reviewer comment (Q1)**$: Completeness of failure modes in Claim 1.

Thank 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.

• $**Reviewer comment (Q2)**$: Applicability to continuous conditions.

We 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.

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.

• $**Reviewer comment (W1)**$: Clarify what is meant by "$\theta$-independent" degenerate solutions.

Thank 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.

• $**Reviewer comment (W2/W3)**$: Clarify the unbounded collapse cases and the proportional collapse case.

We appreciate the opportunity to clarify. Cases (iii)–(v) in Claim 1 describe purely theoretical collapse solutions derived from Eq. (26) in the Appendix:

1. 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.

2. 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).

As 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.

• $**Reviewer comment (Q1)**$: Can the proposed method benefit unconditional generation as well?

Yes--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.

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.

• $**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.

We 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}$.

Consider 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.). We will clarify this point and include the example above in the revision.

• $**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?

To 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.

To validate this, we ran a controlled experiment using the Sec. 4.1 1-D synthetic experiment setup that compares

(i) a learnable $**unconditional**$ source $p(x_0)$ with shift,

(ii) a $**condition-dependent**$ source $p(x_0\mid y)$ with shift dependent on y.

We 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.

• $**Reviewer comment (Q2)**$: Intuition for learning “transporting mass” and “encoding semantics”?

Thank 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:

1. 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.
2. 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.

We hope this clarifies our intuition.

• $**Reviewer comment (Q3)**$: Performance difference between target-only and source-only CAR?

Note 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:

1. On 1-D synthetic data (Sec. 4.1), source-only CAR yields a very slightly lower Wasserstein distance than target-only (Fig. 3(a)).
2. 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.

• $**Reviewer comment (Q4)**$: Meaning of cases (iii) and (iv)?

Cases (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.

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.

• $**Reviewer comment (W1)**$: Can CAR go beyond shift-only reparameterizations?

We 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.

• $**Reviewer comment (W2/Q4)**$: Request for ablation studies.

In the current manuscript, we have already conducted ablations to isolate CAR’s components:

1. CAR variants (Source-only, target-only, and joint CAR): Fig. 2 & Table 2 (synthetic data); Table 3 (ImageNet).

2. Learning-rate ablations: Appendix Fig. 5 (synthetic data) and Fig. 6 (ImageNet).

If 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.

• $**Reviewer comment (W3/Q5)**$: Limited dataset diversity—only ImageNet/FID.

Thank 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.

We 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:

All 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.

• $**Reviewer comment (W4/Q3)**$: Can CAR be applied to FFJORD or neural autoregressive flows?

We 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.

• $**Reviewer comment (Q1)**$: How does CAR perform under weak or noisy conditioning?

In 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.