Improving Flow Matching by Aligning Flow Divergence

Thank you for your thoughtful review and valuable feedback. We have revised the paper according to all reviewers’ feedback. In what follows, we provide point-by-point responses to your comments.

---

Q1. Lack of implementation details which may indicate unfair comparisons.

Response: In our submission, we have provided experimental and implementation details in Appendix B. To ensure fair comparisons with benchmarks, we use the same neural network architectures and model hyperparameters as the baseline methods. In particular, we use the same dataset and splitting, and employ the same search space for the batch size with the same training iterations.

---

Q2. The organization of the contribution part should be improved.

Response: Thank you for your suggestion. We have made the contribution section more concise and bulleted the key points in the revised paper.

---

Q3. Quantify the individual impact of the divergence loss term $L_{CDM}$ on performance.

Response: Our training loss is a weighted sum of $L_{CFM}$ and $L_{CDM}$ as shown in Equation (19) of our submission. Only using $L_{CFM}$ is the same as classic conditional flow matching, while, as we state in Lines 268-270 of our manuscript, directly minimizing $L_{CDM}$ cannot yield appealing results. Detailed explanations are given after those lines. In particular, we notice empirically that training the model by minimizing $L_{CDM}$ alone can be quite noisy.

Using a weighted sum of $L_{CFM}$ and $L_{CDM}$ can improve the performance of the classic conditional flow matching, solidifying the benefits of our proposed approach.

---

Q4. How does FDM’s training time and memory footprint scale with data dimensionality, especially for high-resolution videos?

Response: Compared to the standard conditional flow matching (CFM), our proposed loss function involves an additional divergence term and essentially has a similar scalability as CFM. In practice, instead of explicitly computing the Jacobian matrix, we follow the approach used in [1] and [2], which employs the Hutchinson estimator to approximate the divergence term (see Lines 307-309 in our submission). Then, we only need to compute the derivative of a scalar-valued output with respect to the input to construct the loss, which involves a single call to torch.autograd.grad() and adds just one additional backward pass. Since we're only doing one extra backward pass, the additional memory footprint and computational time of computing the divergence loss scales similarly to that of the original CFM loss with respect to the data dimensionality. We compared the training times and peak memory usage, and found that incorporating the divergence loss does not significantly hinder the model's usability, with computational time and memory footprint remaining mostly within 1.5 times that of the original CFM.

Regarding the video experiment mentioned by the reviewer, we first use a pretrained model to map the video data into a latent space. As a result, the computational cost of the flow matching model depends only on the dimensionality of the latent variables, not the raw video data.

[1] Lu, Cheng, et al. "Maximum likelihood training for score-based diffusion odes by high order denoising score matching." International Conference on Machine Learning. ICML, 2022. https://arxiv.org/pdf/2206.08265 [2] Lai, Chieh-Hsin, et al. "Fp-diffusion: Improving score-based diffusion models by enforcing the underlying score fokker-planck equation." International Conference on Machine Learning. ICML, 2023. https://arxiv.org/pdf/2210.04296

---

Q5. Can FDM be applied to non-Gaussian or non-OT probability paths?

Response: Yes, FDM can be applied to non-Gaussian paths. In our DNA sequence generation experiment, we use a Dirichlet probability path, which was designed in [3].

[3] Stark et al. Dirichlet Flow Matching with Applications to DNA Sequence Design, ICML, 2024.

---

Thank you for considering our rebuttal.

Thank you for your thoughtful review and valuable feedback. We have revised the paper according to all reviewers’ feedback. In what follows, we provide point-by-point responses to your comments.

---

Q1. It will be great if some knowledge can be shared regarding the relation between the choice of conditional vector field and the suggested choice of $\lambda$s.

Response: It is an interesting comment. Our design of the new loss function - Equation (19) of our submission - is based on the empirical observation and discussion in Lines 268-274 (left column) and 220-229 (right column) of our submission. During our submission, we did not consider the meaning of $\lambda$s in terms of data, but it is an interesting future direction to design a principle to choose $\lambda$s optimally. We have acknowledged this point in the revised manuscript.

---

Q2. While I find the localization of DM with CDM very inspiring, … It feels that, when the method is combined with the works like (Simulation-free Schrodinger bridges via score and flow matching), things would also align really well. Has this work been considered to be applied to Schrodinger Bridge?

Response: Thank you for the very insightful comments. We have not considered the Schrodinger bridge yet, while we believe it is indeed an interesting research problem. We have briefly discussed this and cited the related reference that the reviewer pointed out in the revised manuscript.

In the flow matching setting, the conditional probability path and the associated vector field are often designed on a per-sample basis. Therefore, our proposed CDM is tractable, similar to the baseline conditional flow matching.

---

Q3. How does the computational overhead scale by the addition of $L_{CDM}$? The gradient of $\nabla\cdot \hat v(x,\theta)$ would, in the process of backpropagation, require second derivatives, and that sounds a bit scary from the engineering perspectives when we think of scaling this up to the level of very large networks. Would you please comment on that?

Response: In practice, the divergence term $\nabla\cdot \hat v(x,\theta)$ is computed efficiently using the Hutchinson estimator (see Lines 307-309 in our submission). We follow the approach used in [1] and [2], which employs the Hutchinson estimator to approximate the divergence term. Then, we only need to compute the derivative of a scalar-valued output with respect to the input to construct the loss, which involves a single call to torch.autograd.grad() and adds just one additional backward pass. Since we're only doing one extra backward pass, the additional memory footprint and computational time of computing the divergence loss scales similarly to that of the original CFM loss with respect to the data dimensionality. We compared the training times and peak memory usage, and found that incorporating the divergence loss does not significantly hinder the model's usability, with computational time and memory footprint remaining mostly within 1.5 times that of the original CFM.

[1] Lu, Cheng, et al. "Maximum likelihood training for score-based diffusion odes by high order denoising score matching." International Conference on Machine Learning. ICML, 2022. https://arxiv.org/pdf/2206.08265 [2] Lai, Chieh-Hsin, et al. "Fp-diffusion: Improving score-based diffusion models by enforcing the underlying score fokker-planck equation." International Conference on Machine Learning. ICML, 2023. https://arxiv.org/pdf/2210.04296

---

Thank you for considering our rebuttal.

Thank you for your thoughtful review and valuable feedback. We have revised the paper according to all reviewers’ feedback. In what follows, we provide point-by-point responses to your comments.

---

Q1. At the end of page 6, … it is not clear that the equation at the bottom of page 6 is an exact relation, and would introduce extra errors.

Response: Thank you for your comment, which helped us identify areas where we can improve our presentation, and we have revised the paper accordingly.

To address your concern, we stress that our theoretical results do not depend on this relation for the learned vector field $\mathbf{v}_t(x,\theta)$. As the reviewer noted, this relation holds only for the ground-truth vector field $\mathbf{u}\_t$ and the score function of its corresponding probability flow $p_t$ (due to the uniqueness of solutions to the continuity equation). Our derivations rely solely on the continuity equation and initial conditions, not on this specific relationship, which aligns with the reviewer’s observation.

Our original intent in introducing this relation was to define a flow-matching model with the vector field $\mathbf{u}\_t$ corresponding to the score-based diffusion model with drift term $\bf f$ and noise coefficient $g$. Notice that approximating $\mathbf{u}\_t$ with a parameterized vector field $\mathbf{v}\_t(x,\theta)$ is enough to generate data effectively. However, for the numerical results (e.g., conditional probability computations in Table 2), we require an approximation of the score function $\nabla \log p_t$. Instead of computing $\nabla \log \hat p\_t$, we substitute the learned vector field $\mathbf{v}\_t(x,\theta)$ into the following relation to estimate the score: $\nabla \log p\_{1-t}(x) = 2\frac{\mathbf{u}\_t(x)+{\bf f}\_{1-t}(x) }{g^2_{1-t}}.$ While this introduces an approximation error, it is a practical trade-off for computational feasibility.

---

Q2. The bounds rely on an analytic form for the score $\nabla\log p_t(x)$ arising from solving the continuity equation up to time $t$. This quantity is not equivalent to the score arising from the model i.e. $\hat s_t\neq \nabla \log p_t$, nor is it something that can be written down by relating a learned velocity field to a score.

Response: Indeed, computing $L_{DM}$ requires $\nabla\log p_t(x)$, which is derived by solving the continuity equation. To address this challenge, we introduce its conditional variant $L_{CDM}$. For $L_{CDM}$, we only need access to $\nabla\log p_t(x|x_1)$. In practice, we design $p_t(x|x_1)$ such that this quantity is readily available by plugging in values, bypassing solving an ODE explicitly. This approach ensures computational feasibility while maintaining the integrity of our bounds; see Theorem 4.2.

---

Q3. The method seems to perform better, but only marginally. Given the fact that loss is more expensive …, how does one measure the advantage of the technique properly?

Response: The improvements for all tasks except trajectory sampling for dynamical systems are way over twice the standard deviation. For dynamical systems, we further perform a t-test with the null hypothesis that the mean negative log-likelihoods (NLLs) computed by two distinct models are equal. We use 32,000 test trajectories to compute each mean NLL. The t-test confirms that the improvement in the likelihood estimation of FDM over FM is significant.

Our approach introduces additional computational cost, but it gives better generation results and better likelihood estimation. For tasks requiring accurate likelihood estimation, our approach is especially more valuable compared to the baseline flow matching.

---

Q4. How does the method compare for forecasting videos to [1]?

Response: Thank you. We have cited this paper in the revision.

Indeed, the results in [1] are impressive. We only found the codes for Navier-Stokes and CIFAR at github.com/interpolants/forecasting, which is the reason we did not include this approach as a benchmark in our study.

[1] https://arxiv.org/pdf/2403.13724

---

Q5. The authors should probably also cite Liu et al, Flow Straight and Fast (2023) for its contributions to the flow matching literature as well.

Response: We have cited this paper in the revision.

---

Q6. Memory footprint in computing divergence. Does the method work when replaced with a Hutchinson estimator?

Response: In our implementation, we use the Hutchinson estimator to approximate the divergence term; see Lines 307-309 in our submission. It works well and does not significantly raise the memory footprint, which consistently stays within 1.5×.

For the 2D tasks, we test both computing the Hutchinson divergence estimator and the exact one; both result in nearly identical results. We use the Hutchinson estimator in all the other experiments due to memory constraints, following [2].

[2] https://arxiv.org/pdf/2206.08265

---

Thank you for considering our rebuttal.

Thank you for your thoughtful review and valuable feedback. We have revised the paper according to all reviewers’ feedback. In what follows, we provide point-by-point responses to your comments.

---

Q1. I would like to see the performance of their method on some more diverse video datasets or standard image benchmarks for generative modeling (CIFAR, Imagenet).

Response: During the rebuttal period, we conducted experiments to showcase the efficacy of our proposed FDM in improving flow matching (FM) for these reviewer-mentioned tasks. We first compare the performance of FDM against the baseline conditional FM (CFM) for generative modeling of CIFAR10 here, and we will show the advantages of FDM for other video generation tasks in Q2.

We follow the experimental settings in Appendix E of the flow matching baseline paper [1] to train the model for CIFAR10 generation. Due to the time constraint, we only compare our FDM against CFM for CIFAR10 generation using the better-performing optimal transport (OT) path in the baseline paper [1]. The following results confirm the advantages of FDM.

[1] https://arxiv.org/pdf/2210.02747

---

Q2. The KTH dataset isn't used usually for generative modeling benchmarks due to its limited diversity in terms of samples. It works as a simple example but some more diverse motion datasets such as BAIR (as done in Davtyan et. al (2023)) could be more interesting.

Response: We appreciate your feedback. For the BAIR dataset, we predict 15 future frames based on a single initial frame, with each frame having a resolution of $64 \times 64$ pixels. Due to the highly stochastic motion in the videos of the BAIR dataset, we evaluate the model as follows (following [2]): We randomly select 256 test videos, and generate 100 samples per test video where each model is conditioned on the same initial frame of the video. Finally, we compute the FVD for each model comparing the $256\times 100$ generated samples against the 256 test videos. Due to time constraints, we omit training the frame refinement network, which operates independently of the main model but could potentially enhance sample quality. The result is as follows:

As mentioned in [2], many models for the BAIR task are computationally expensive, whereas latent FM achieves a favorable trade-off between FVD and computational cost. Our approach further improves latent FM with acceptable additional computational overhead.

[2] https://arxiv.org/abs/2211.14575

---

Q3. Minor typos (does not effect valuation) L_{FDM} is used in the graph of Figure 1, but caption reads L_{DM}

Response: Thank you. We have fixed this in the revised manuscript.

---

Thank you for considering our rebuttal.