Local Flow Matching Generative Models

Authors: Chen Xu, Xiuyuan Cheng, Yao Xie
Primary area: generative models
Keywords: flow-based generative modeling, stepwise training, model distillation

Abstract:
Flow Matching (FM) is a simulation-free method for learning a continuous and invertible flow to interpolate between two distributions, and in particular to generate data from noise in generative modeling. In this paper, we introduce Local Flow Matching ($\texttt{LFM}$), which consecutively learns a sequence of FM sub-models and each matches a diffusion process up to the time of the step size in the data-to-noise direction. In each step, the two distributions to be interpolated by the sub-model are closer to each other than data vs. noise, and this enables the use of smaller models with faster training. The stepwise structure of $\texttt{LFM}$ is natural to be distilled and different distillation techniques can be adopted to speed up generation. Theoretically, we prove a generation guarantee of the proposed flow model in terms of the $\chi^2$-divergence between the generated and true data distributions. In experiments, we demonstrate the improved training efficiency and competitive generative performance of $\texttt{LFM}$ compared to FM on the unconditional generation of tabular data and image datasets, and also on the conditional generation of robotic manipulation policies.

Decision: reject
Reviewer scores: 5, 3, 3, 6

Reviews and discussion:
<thread>
creation date:
1729500052108

thread_type:
official_review

<thread_object>
creator: reviewer_snd7

<content>
summary:
The paper proposes a novel generative model, Local Flow Matching (LFM), which segments Flow Matching along the time dimension. The transition between two distributions within each time segment is modeled using a smaller-scale sub-model. This generative model offers the advantages of having smaller sub-models, faster convergence during training, and greater convenience for distillation.

strengths:
The paper provides a detailed theoretical proof that the data distribution generated by LFM is close to the true data distribution, demonstrating the exponential convergence of the forward process and the generation guarantee of the backward process. In addition, the paper conducts extensive experiments to validate the effectiveness of LFM, including generation on two-dimensional toy data, tabular data generation, image generation, and robotic manipulation. Moreover, the paper is well-written.

weaknesses:
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.

questions:
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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1729949331955

thread_type:
official_review

<thread_object>
creator: reviewer_vdpz

<content>
summary:
The paper introduces Local Flow Matching (LFM), a generative approach that builds upon the existing Flow Matching (FM) framework. LFM improves upon FM by dividing the global flow into multiple local flow sub-models. Each sub-model is trained to match a diffusion process between closer intermediate distributions, which reduces the model size and accelerates training.

strengths:
1.	The paper offers a fresh approach in the field of generative modeling, combining ideas from diffusion processes and FM. The idea of breaking down a single large flow into several smaller steps (local flows) is natural.
2.	The paper provides solid theoretical guarantees, specifically proving a generation guarantee in terms of the $\chi^2$-divergence between the generated and true data distributions. The experiments are comprehensive, covering a range of tasks from image generation to robotic manipulation, and the results demonstrate the performance of LFM.

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

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

3.	The paper highlights improved training efficiency through the use of smaller local models. However, the exact speed improvements are not always well-quantified. For example, while the paper states that LFM leads to faster convergence and reduced memory usage, these claims could be supported by more detailed analysis. Providing concrete timing metrics (e.g., runtime comparisons, memory benchmarks across different hardware) would bolster the argument for LFM's practical benefits. More information can be found in Questions part.

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

questions:
1.	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?

2.	In Appendix B1, how were the stepwise parameters $\rho$ and $c$ selected? Were they chosen solely based on extensive experimentation and selecting the ones that gave the best results, or was there a more principled approach to their selection?

3.	In the image generation tasks (Table 2), it would be beneficial to include comparisons with Flow Matching (FM) [1], rather than only with InterFlow, as Local Flow Matching (LFM) is based on FM. This would better illustrate the claimed improved training efficiency. Additionally, either in the main text or in the appendix, it would be helpful to provide a comparison of the total number of parameters for each method across all experiments, rather than only ensuring the same training scheme.

[1] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, 2023.
</content>
</thread_object>

</thread>
<thread>
creation date:
1729999410050

thread_type:
official_review

<thread_object>
creator: reviewer_k69v

<content>
summary:
This paper presents Local Flow Matching (LFM), an enhancement of Flow Matching (FM) for faster, efficient generative modeling. LFM divides the process into sequential FM sub-models, each bridging distributions that are progressively closer from data to noise, allowing smaller models and quicker training. This stepwise approach also enables the use of distillation techniques to further accelerate data generation. Theoretically, LFM provides a generation guarantee based on $\chi^2$-divergence. Experiments show LFM’s improved training efficiency and strong performance in generating tabular data, images, and robotic manipulation policies.

strengths:
The paper explores integrating flow-matching submodels into diffusion processes to enable faster and more efficient learning and inference. The approach is novel and is supported by theoretical guarantees of generation. Experimentally, the method can be applied to various tasks, including image generation, tabular data generation, and robotic manipulation.

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

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

3. The advantage of LFM on fast training has not been demonstrated well except Table 2.

questions:
1. Can you clarify more on how the method benefits fast training based on the experimental results?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730691242561

thread_type:
official_review

<thread_object>
creator: reviewer_9uhe

<content>
summary:
This paper proposes to divide the original Flow Matching (FM) into a sequence of parts, called Local Flow Matching (LFM). The authors provide theoretical analysis of the convergence of the forward process, and the theoretical guarantee of the $\chi^2$ divergence between the real data and generated data distributions. The authors evaluated the proposed method on different tasks, tabular data generation, image generation, robotic manipulation and so on.

strengths:
The authors proposed the Local Flow Matching (LFM), which divides the Flow Matching into a sequence of local parts. The approach is intuitively correct. The authors provided theoretical analysis of the convergence of the forward process, showing that the $\chi^2$ divergence between the noise distribution and noised real data distribution decreases exponentially w.r.t. the number of sub-flows. Furthermore, the authors presented the theoretical results showing the $\chi^2$ divergence between the real distribution and the generated distribution is close w.r.t. the flow training loss $\epsilon$.

weaknesses:
From the experiments, it turns out that the proposed LFM is not working as well as the original Flow Matching (FM) (Lipman et al., 2023). The proposed method achieved an FID of 8.45 on the CIFAR-10 dataset whereas the FM achieved a lower FID of 6.35 on the same dataset. On ImageNet 32x32 dataset, the proposed method achieved an FID of 7.0, but the FM achieved a lower FID of 5.02. The authors should compare with FM on the CIFAR-10, ImageNet 32x32, the Oxford Flowers and the LSUN Church datasets. 

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. 

The proposed method also need more inference time since we need to inference from $v_1$ to $v_N$.

questions:
What does $g \bot x'_l$ mean in Eq. 6? and why do we need this constraint?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732504962384

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Common response to all reviewers

comment:
We appreciate the comments and insights from all reviewers. We have revised the paper with revised parts highlighted in blue. Below, we first address common questions and then address individual ones.

1. **Comparison of LFM with FM (Lipman et al., 2023).**
   
   **Ref:**
   
   **Lipman et al. 2023: Flow Matching for Generative Modeling. ICLR 2023.**

We emphasize that in the submitted version, LFM already demonstrated superior performance to FM on non-image tasks, showing improved training efficiency with equivalent model size and training epochs: 
- Figure 2: LFM achieved better 2D data generation, both qualitatively and quantitatively (lower NLL).
- Table 4: For robotic manipulation, LFM attained higher success rates than FM on Can, Transport, and Toolhang tasks.	

To strengthen our argument, we conducted additional comparisons between LFM and FM on CIFAR-10, Imagenet-32, and Flowers-128 image tasks. For fairness, we maintained equal total model parameters across all LFM blocks and FM's global flow. The updated Table 2 is included in the revision and copied below for convenience. We highlight that with identical training batch count and size, LFM achieves lower FID values than FM, indicating faster training efficiency. 

Table: FID (lower is better) comparison of LFM against InterFlow and FM under same model sizes. FIDs with the symbol * are quoted from the original publication. Note that InterFlow uses the Trig interpolant and FM uses the OT interpolant to train the global flow.

| | | CIFAR-10| Imagenet-32|Flowers-128|
|-|-|-|-|-|
| LFM (Trig interpolant)| FID| **8.45**        | **7.00**              | **59.7**           |
|                           | Batch size| **200**             | **256**                   | 40                 |
|                           | # of batches| **5 × 10^4**        | **2 × 10^5**              | 4 × 10^4           |
| InterFlow| FID|  10.27*      | 8.49*             | 65.9           |
|                           |Batch size| 400             | 512                   | 40                 |
|                           |# of batches|5 × 10^5        | 6 × 10^5              | 4 × 10^4           |
||||||
| LFM (OT interpolant)|FID| **8.55**        | **7.20**              | **55.7**           |
|                           |Batch size| 200             | 256                   | 40                 |
|                           |# of batches| 5 × 10^4        | 2 × 10^5              | 4 × 10^4           |
|FM      |FID| 12.30|  7.51              | 70.8           |
|                           |Batch size| 200             | 256                   | 40                 |
|                           |# of batches| 5 × 10^4        | 2 × 10^5              | 4 × 10^4           |

Ref:

Albergo & Vanden-Eijnden, 2023: Building normalizing flows with stochastic interpolants. ICLR 2023.


2. **Model size of LFM against global flows. Are they comparable?**

We emphasize that throughout the experiment section, we maintained equal model parameters between LFM and the global flow. Specifically, for a global flow with $M$ parameters and $K$ LFM local flows, each local flow has $M/K$ parameters. The value of $M$ is detailed in Tables A.1 - A.3.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732674051722

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Awaiting rebuttal response

comment:
Dear Reviewers,

We sincerely appreciate your constructive feedback and thoughtful questions, which have greatly enriched our work. We wanted to kindly follow up regarding the response to our rebuttal.

Thank you again for your time and effort in reviewing our submission.

Best regards,
The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper suggests training a sequence of flow models by breaking the interval of interest to several parts and training each part using Flow Matching to match marginal probabilities defined by an OU process, in particular Variance Preserving (VP). 

It is not clear whether breaking the generation problem to a collection of smaller problems (here with VP marginal probability paths) indeed simplifies the problem and/or provides practical benefits, e.g., compared to training a flow directly using this VP marginal probability path over the full interval. The main difference is that the current approach in practice adds data couplings at each interval compared to the global coupling used for the original problem. The experimental part falls short in showing the suggested method has a strong practical benefit in terms of performance, faster training and/or sampling complexity. Furthermore, some implementation details are missing in the experimental part.