Local Flow Matching Generative Models

1. (Weakness) Some experimental details in the paper may need further clarification. The data in Table 2 is used to demonstrate that LFM requires less computational cost, but the model sizes of both LFM and InterFlow are not provided. Additionally, it is unclear whether the training data shown in the table refers to the data used by each sub-model individually or by all sub-models collectively.

Please refer to the second question in the common response regarding the same number of model parameters being used throughout experiments. Meanwhile, the batch size in Table 2 refers to the number of data used by each sub-model individually, each of which is $K$ times smaller than the global flow, assuming $K$ local flows are used in LFM.

2. (Question) Could you provide more detailed data regarding the issues mentioned in the weaknesses, such as the specific number of parameters for each sub-model in LFM, the number of sub-models, and the parameter count of the InterFlow model (or explain why the total number of parameters across all sub-models is equal to that of InterFlow). Could you provide the amount of data required to train each sub-model?

Information regarding “the specific number of parameters for each sub-model in LFM, the number of sub-models, and the parameter count of the InterFlow model” is provided in Appendix Tables A.1 - A.3. The total number of parameters across all sub-models is equal to that of InterFlow because we scale the width of hidden layers of LFM, so that each block is proportionally smaller.

The amount of data required to train each sub-model is the same as that used to train the global flow. More specifically, training data of the $n$-block is the empirical push-forward distribution by the $(n-1)$-th block, which maps its training data through its ODE solution map to obtain the push-forward distribution.

Please see our first question in the common response regarding the comparison of LFM with FM.

1. (Weakness) FM simplifies the diffusion process into a single step, transforming a trajectory from a curve into a straight line, thus reducing training costs and improving sampling efficiency. However, this paper reverses that advantage by breaking this single step into multiple segments. The division of the trajectory into multiple steps could introduce added computational complexity, potentially negating the efficiency gains that FM originally aimed to provide.

We appreciate the comment and address it below:

• "FM...transforming a trajectory from a curve into a straight line...thus reducing training costs and improving sampling efficiency": Theoretically, as argued in (Lipman et al., 2023), the linear interpolation by FM is the optimal transport map, leading to "straight line trajectories with constant speed". However, we empirically found that the trained velocity field yields curvy trajectories when the data distribution $P$ is close to the target Gaussian distribution. We tested this by setting $P=\mathcal{N}(\mu, I_2)$ for $\mu=[-0.05, -0.05]$ and training an FM model $v_{\theta}$. Visualizing the ODE trajectory (see here) given 20 random starting points using $v_{\theta}$ shows that the trajectory can be highly curvy. Thus, it remains unclear under what conditions the empirically trained FM model yields straight trajectories and consequently "reduces training costs and improves sampling efficiency" as suggested.
• "The division of the trajectory into multiple steps could introduce added computational complexity, potentially negating the efficiency gains that FM originally aimed to provide": Having addressed why FM does not necessarily yield straight trajectories and thus establish efficiency gains, we argue that dividing the trajectory into multiple steps does not necessarily introduce added computational complexity. Our experiments consistently show improved generation performance under the same number of training batches and total number of model parameters using LFM. Additionally, using the 2D rose dataset, we demonstrate that the ODE trajectory for each LFM block (under OT interpolant) is fairly straight (see here). This is achieved by training each LFM block using a dependent coupling of $(p_{n-1}, p_n^*)$. Specifically, given the left end $x_l'\sim p_{n-1}$ and an independent $g \sim \mathcal{N}(0,I_d)$, the right end $x_r$ as defined in Eq. (6) depends on $x_l'$.

2. (Weakness) The paper presents some conceptual ambiguities regarding key terms. Continuous Normalizing Flows (CNF) and neural ODEs are two distinct models, but the paper incorrectly states that "CNF trains a neural ODE" (Line 106). In fact, FM is a method for training CNF models, not a model in itself, much like the proposed LFM method. The paper conflates these concepts in several places (e.g. Line 41, 51), which could cause confusion.

We appreciate the reviewer's comments and agree with their observations. We have revised our statements to address these points:

• Regarding "CNF trains a neural ODE": We acknowledge the ambiguity here, which may imply CNF is a method to train a neural ODE. To clarify, CNF uses a neural ODE as its underlying model, which is optimized by maximizing the data likelihood. We have revised this statement in the paper to reflect this more accurate description.
• Regarding FM and LFM: We agree that FM was originally introduced as "a simulation-free approach for training CNFs" rather than a model itself. In this context, our proposed LFM is best understood as an improved version of FM for training CNF models. This clarification helps to position both FM and LFM as methods for training CNFs, rather than standalone models.

3. (Weakness) In line 127, a reference is used as the subject of the sentence but is enclosed in parentheses. Please correct this formatting issue and ensure consistency throughout the paper.

We have fixed this issue and similar ones in the paper.

4. (Question) In the training process, at each step, the target distribution is generated from the sample of the previous step through the Ornstein-Uhlenbeck (OU) process. What is the theoretical basis for using the OU process in this context? This essentially defines checkpoints along the probability path. Would a different process yield similar results, or is there a specific reason the OU process was chosen?

Please refer to Section 4 for in-depth theoretical analyses on OU process benefits. The OU process ensures:

• Exponential convergence of the forward process in $\chi^2$ divergence (Proposition 4.1)
• Generation guarantee for distributions with regular data density (Theorem 4.2)

While alternative processes may be possible, their analysis is beyond this work's scope.

Thank you for your detailed rebuttal. I am generally satisfied with your responses to Questions 2, 3, and 5. Below, I provide feedback on other points:

Regarding Question 1 It would be beneficial to include additional comparisons to related works, such as Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport (Tong et al., 2023) and Simulation-free Schrödinger Bridges via Score and Flow Matching, which similarly focus on straightening vector fields to improve transport efficiency. Additionally, contrasting the trajectory visualizations of your method with FM models, such as the right part of Figure 1 in Tong et al. (2023), would provide more clarity and context.

Regarding Question 4 Thank you for directing me to the theoretical analysis in Section 4. While the assumptions in Proposition 4.1 and Theorem 4.2 are mathematically reasonable, I have concerns about their alignment with the experimental setting. For instance, in high-dimensional tasks like image generation, real-world data distributions are often multimodal or highly non-Gaussian, which may deviate from the theoretical assumptions. Further discussion or validation of these assumptions in such practical contexts would strengthen the paper.

Regarding Questions 6, 7 Thank you for the additional experimental results provided in the official comment. However, I did not find results that conclusively demonstrate the training efficiency advantage of your method. Specifically:

1. Your comparison uses a training scheme with advantages for LFM (e.g., batch size and number of training batches) against FM. This is not a fair comparison unless FM is also tested under optimal conditions or a comparable training scheme. Including FM results with its best (or near-best) FID scores would provide a more balanced comparison.
2. The rebuttal mentions the non-linear relationship between training time and model size, which is a valid point. However, direct experimental comparisons are necessary to substantiate LFM's training efficiency. For example, with same model size, reporting either the runtime required to achieve comparable FID scores or the FID scores achieved within comparable runtimes would better demonstrate the claimed efficiency advantages.

These points are not addressed in the rebuttal. Therefore, I maintain my score of 3.

Please see our first question in the common response regarding the comparison of LFM with FM, where the faster training of LFM is demonstrated by the fact that under identical model sizes, LFM reaches smaller FID than the baseline methods after we train them for the same number of batches.

1. (Weakness) More details should be provided, such as the number of function evaluations (NFEs) and the used ODE sampler for all methods in Table 1 and Table 2, to better demonstrate LFM's effectiveness.

We employ the Dormand-Prince-Shampine (dopri5) ODE sampler for LFM, with tolerances of 1e-5 for non-image examples (Table 1) and 1e-4 for 32x32 images (Table 2). NFE information for most methods in Table 1 is unavailable. However, we compare the NFE and FID performance of LFM against InterFlow and FM below, where we note that

• LFM achieves lower FID than both methods.
• While LFM's total NFEs across all local blocks exceed that of the global flow, each NFE uses a smaller model (1/4 size of global flow’s because we used 4 blocks).
• In the last column, adjusting the $\{ \gamma_n \}$ timestep schedule significantly reduces LFM's NFE with small impact on the FID (still better than FM’s FID).

Table: FID and NFE comparison of LFM vs. InterFlow/FM on Flowers 128x128, maintaining equal total model parameters (in millions) across all four LFM blocks and the global flow. All methods trained for 40K batches (batch size 40), using dopri5 ODE sampler with 1e-3 relative & absolute tolerance.

2. (Weakness) Although the tasks are diverse, the paper lacks a solid comparison with prior methods on some fundamental tasks. For instance, comparing LFM with Rectified Flow and OU diffusion on CIFAR-10 with the same amount of batches would offer a clearer understanding of LFM's training efficiency and distillation effects.

We refer the reviewer to the first question in the common response regarding the comparison of LFM with FM and InterFlow on CIFAR-10, Imagenet-32, and Flowers-128, where we demonstrated both faster training efficiency and improved distillation effects.

To further demonstrate the advantage on distillation, we additionally compare LFM with $K$-Rectified Flow (Liu et al., 2023) for $K=2, 3$ (1-Rectified Flow = FM). The Table below shows that LFM reaches better distillation effects than both FM and $K$-Rectified Flow for $K>1$.

Table: FID comparison of LFM and FM & $K$-Rectified Flow before and after distillation (into 4 NFEs) on Flowers 128x128. Note that the training data for the $K$-Rectified Flow is generated by the $(K-1)$-Rectified Flow based on (Liu et al., 2023), so that pre-distillation FID is worse as $K$ increases.

Ref:

Liu et al., 2023: Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. ICLR 2023

Please see our first question in the common response regarding the comparison of LFM with FM.

1. (Weakness) The proposed LFM seems to require more training time than the original FM because for training $v_n$ we need to inference from $v_1$ to $v_{n-1}$ for each sample.

We clarify that training the $n$-th block requires only the empirical push-forward distribution from the $(n-1)$-th block. Crucially, this distribution is computed only once, before training of the $n$-th block begins (see line 5, Algorithm 1). Thus, we avoid repeated inference through all $n-1$ blocks while training the $n$-th block.

2. (Weakness) The proposed method also need more inference time since we need to inference from

We note that throughout the entire experiment section, the total number of parameters over all $v_1,\ldots, v_N$ are the same as the global flow for InterFlow or FM. Please see the second question in the common response above regarding this. Thus, having to inference from $v_1$ to $v_N$ does not necessarily introduce additional inference time because each $v_i$ is proportionally smaller.

3. (Question) What does $g \perp x’_l$ mean in Eq. 6? and why do we need this constraint?

This notation means $g$ and $x’_l$ are independently sampled. We need this constraint so that $x_r$ would correctly follow the distribution $p_n^*$.