One-step Flow Matching Generators

Authors: Zemin Huang, Zhengyang Geng, Weijian Luo, Guo-Jun Qi
Primary area: generative models
Keywords: one-step generator, text-to-image generation, flow matching

Abstract:
In the realm of Artificial Intelligence Generated Content (AIGC), flow-matching models have emerged as a powerhouse, achieving success due to their robust theoretical underpinnings and solid ability for large-scale generative modeling. These models have demonstrated state-of-the-art performance, but their brilliance comes at a cost. The process of sampling from these models is notoriously demanding on computational resources, as it necessitates the use of multi-step numerical ordinary differential equations (ODEs). Against this backdrop, this paper presents a novel solution with theoretical guarantees in the form of Flow Generator Matching (FGM), an innovative approach designed to accelerate the sampling of flow-matching models into a one-step generation, while maintaining the original performance. On the CIFAR10 unconditional generation benchmark, our one-step FGM model achieves a new record Fréchet Inception Distance (FID) score of 3.08 among all flow-matching-based models, outperforming flow matching models that use 50 generation steps. 
Furthermore, we use the FGM to distill the Stable Diffusion 3, which is a leading text-to-image flow-matching model. The resulting model named the MM-DiT-FGM demonstrates outstanding industry-level performance as a novel transformer-based one-step text-to-image generator.  When evaluated on GenEval benchmark, MM-DiT-FGM has delivered remarkable generating qualities, rivaling other multi-step models in light of the efficiency of a single generation step. We will release our one-step FGM text-to-image model with this paper.

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

Reviews and discussion:
<thread>
creation date:
1730261947403

thread_type:
official_review

<thread_object>
creator: reviewer_bvru

<content>
summary:
The paper introduces Flow Generator Matching (FGM), a method for to distill a pretrained flow-matching model into a single-step model. Specifically, the paper proposes a data-free distillation approach to match the implicit vector field of the student model $v_{\theta, t}$ with the pretrained vector field u_t, where u_t can be regarded as the teacher. Experiments are carried out on CIFAR-10 and paired text-image dataset, which demonstrate the effectiveness.

strengths:
1. The paper proposes a novel method to distill a flow-matching model based on the property of its vector field from a probabilistic standpoint, which represents a good contribution.
2. Experiments show effectiveness of the appraoch.

weaknesses:
1. My biggest concern is the clarity of the paper. 
    - Algorithm 1 is not clear to me and hard to understand. What is the online flow model $v_\psi$? It gets introduced nowhere in the entire Section 4 except the Input field of Algorithm 1. Does it represent the implicit vector field? If yes, why shall we have another notation of $v_{sg[\theta], t}$ in 4.11 and 4.12? If not, then what does the online model do and how could we use $\theta$ to parameterize both the generator and the implicit vector field? Also in 4.12, there're both $u_t(x_t)$ and $u_t(x_t | x_0)$, yet $x_t$ is only sampled from $x_t | x_0$ ~ $q_t(x_t | x_0)$, then why do we have different notations here? This is the core of the paper which concretizes the theory into algorithm, and it shall be presented as clear as possible. 
    - in line 203, there's an abuse of notation where using $x_0 = g_\theta(z)$ instead of $x$ could make the notation more consistent with the latter integral.
2. The one-step results seem to have checkerboard artifact. When zoomed in Fig.1, most samples show the pattern. It also appear in the 2nd sample of Figure 3 while SD3-28 steps have much smoother texture. The quality degradation is still a concern. 
3. There's missed reference in Table 2 for StyleGAN2 + Smart. 
4.  How do we find the t* in 5.1? The ablation study is not presented.

questions:
Please refer to weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730616309185

thread_type:
official_review

<thread_object>
creator: reviewer_4qee

<content>
summary:
This paper introduces Flow Generator Matching, an innovative distillation framework for pre-trained flow-matching models. The framework is designed to accelerate sampling by constructing a one-step generator. To achieve this, the authors define a noise-to-image mapping $g_\theta(z)$, and optimize the distance between the vector field of the pre-trained flow model, and the vector field implicitly derived from the one-step generated images and the online flow model. By leveraging a theoretical approximation of this optimization using the flow product identity, the authors develop a fast and efficient one-step flow generator.

strengths:
- Theoretical justification: The proposed distillation formulation is both intuitive and novel in the context of flow matching. Although optimizing $L_{FM}$ directly with respect to $\theta$ is non-trivial, the authors introduce an innovative alternative whose gradient remains aligned with the original objective.

- Specialized distillation methods for flow matching: Given the widespread use of flow matching in various state-of-the-art large-scale generative models (e.g., Movie Gen), effective sampling acceleration is critical. This paper offers a promising approach to address this important contemporary challenge.

- Experiments: Comprehensive empirical evaluations, including unconditional, conditional, and text-to-image experiments, support the effectiveness of the method. The approach demonstrates competitive performance with other flow distillation techniques, achieving lower NFE.

weaknesses:
- Writing & Notations: The connection between the derivations and practical implementations is difficult to follow. For instance, while FGM relies on the online flow model parameterized by $\psi$, as outlined in Algorithm 1, $\psi$ is notably absent in Section 4, despite being a crucial component. Additionally, there is no explanation or justification for the phrase "stop the $\theta$-gradient" mentioned in Line 250, which is a critical implementation detail. From my understanding, FGM is conceptually similar to distribution matching distillation, where both the generator $\theta$ and an online critic $\psi$ are updated alternately. However, this understanding comes from Algorithm 1 rather than the main text. The authors should provide a clear and theoretically grounded explanation of "stop gradient" and revise the derivations to explicitly incorporate the online flow model $\psi$. Moreover, some important information are provided in appendix without pointers, e.g.   model parameterization in L952. Detailed pointers in the main paper would be appreciated.

- Novelty: As previously mentioned, FGM shares significant technical similarities with existing diffusion distillation methods. This overlap may diminish the novelty of the approach and, consequently, inherit some of the limitations of these existing methods, such as reliance on a potentially sub-optimal online flow model.

- Qualitative Results: Although FGM demonstrates better quantitative performance compared to other flow distillation methods, the generated images appear overly "synthetic" and color-saturated (e.g., the rightmost images in Lines 54 and 488), while i admit that this may vary based on human perception. This could be potentially attributed to the image-data-free property proposed as limitations by the authors.

questions:
Overall, the paper is well-structured and presents promising results. However, the derivations and notations could be refined to enhance clarity and facilitate faster understanding. Below are some questions for consideration:

- Ablation Study on Generator Initialization: The generator is initialized with pre-trained flow models, which may be crucial for warm-up and accelerating convergence, especially given that FGM does not utilize real-image datasets. This initialization might also mitigate the OOD gap for pre-trained flow models trained on real images. Conducting an ablation study or providing further analysis could improve understanding and highlight the significance of this choice.

- Bypassing Equation 4.11: Although Equations 4.11 and 4.12 both originate from the FGM loss, the authors opt not to use Equation 4.11, as mentioned in Line 368. This decision represents a significant bypass. Additional analysis, discussion, and empirical evidence are needed to justify this choice. Could the authors elaborate on why Equation 4.11 is deemed unnecessary in practice?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730648430983

thread_type:
official_review

<thread_object>
creator: reviewer_sdtk

<content>
summary:
This paper proposes extending score implicit matching from diffusion to flow matching, which is called FGM. The method seems to work well on text-to-image and unconditional generation.

strengths:
1. The paper is well-written and easy to follow.
2. The paper extends the score implicit matching to a flow matching generator.

weaknesses:
1. The novelty of the paper is limited. Since flow matching can be considered a special form of diffusion model, score-based diffusion and flow matching are equivalent. Therefore, rawly extending the implicit score matching for the flow matching model should work well.
2. The evaluation for the text-to-image section should include FID, Recall, ClipScore, Image Reward, PickScore, and AES score. These metrics are often used as standard metrics to evaluate generative models.
3. Since this method is closely related to implicit score matching, the author should include a background section for implicit score matching.

questions:
My biggest concern is the novelty and originality of this paper's idea. There is a list of works working on score implicit models, such as [1,2,3]. Extending the framework from diffusion to flow matching is not interesting and does not introduce something new. 

[1]: Diff-Instruct: A Universal Approach for Transferring Knowledge From Pre-trained Diffusion Models

[2]: One-Step Diffusion Distillation through Score Implicit Matching

[3]: Score identity Distillation: Exponentially Fast Distillation of Pretrained Diffusion Models for One-Step Generation
</content>
</thread_object>

</thread>
<thread>
creation date:
1730709395395

thread_type:
official_review

<thread_object>
creator: reviewer_jkbp

<content>
summary:
The paper introduces the **Flow Generator Matching (FGM) objective** for distilling a one-step generator from pretrained flow-matching models. It provides theoretical guarantees that the proposed objective yields the same gradient, with respect to the parameters of the one-step generator, as the standard flow-matching loss. The paper demonstrates competitive and, in some cases, superior results in both unconditional generation on CIFAR-10 and large-scale text-to-image generation, achieved by distilling Stable Diffusion 3 (SD3).

strengths:
1. The paper is well-written and easy to follow.
2. The proposed algorithm is elegantly designed, alternating the optimization of the parameters of the one-step generator, $\theta$, and an auxiliary flow model, $\psi$. The correctness and effectiveness of the algorithm are supported by Theorems 4.1 and 4.2.
3. The paper conducts thorough empirical evaluations of the proposed method across various tasks, including unconditional generation on CIFAR-10 and large-scale text-to-image generation, demonstrating superior results.

weaknesses:
1. The proposed method requires training an auxiliary flow model with parameters $\psi$, designed to generate the implicit flow determined by the one-step generator. Is this auxiliary model used after training?
2. The paper does not include a discussion of training efficiency or speed. One concern is that the proposed method might be slower compared to distillation approaches like Consistency Flow Matching (Yang, Ling, et al., "Consistency Flow Matching: Defining Straight Flows with Velocity Consistency"). Could you compare or discuss the training efficiency of the proposed methods with other common baselines?

questions:
Please see weaknesses part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732439603640

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Paper Revision

comment:
# Dear Reviewers and Area Chair,

We thank all reviewers for their useful feedback. We have made significant updates to our draft to improve the overall writing quality, provide more clarifications, and add additional experimental results. We have highlighted the changes in the draft in blue color. Here are the major changes:

1. As **Reviewer sdTK** and **Reviewer jkBP** wish, **we polish the presentation of the method**. We add a paragraph in Section 4.2 to clarify the difference between FGM and diffusion distillation by pointing out technical challenges. We also add a paragraph in Section 5.1 to compare the training efficiency of FGM and CFM, showing that FGM can significantly surpass CFM on both generation results and training efficiency. 

2. In response to **Reviewer 4qEE** 's suggestion. In Section 4.2, **we have included a comprehensive explanation of the stop $\theta$-gradient technique** utilized within our model. This addition aims to clarify its implementation details and the impact on gradient backpropagation

3. As suggested by **Reviewer BVRU**, **we have revised the notation in Algorithm 1** for improved clarity. Specifically, we replaced $v_{\text{sg}[\theta]}$  with the actual online flow model $v_\psi$ that we employ. Furthermore, we have provided additional explanations concerning the online flow model $v_\psi$ in Section 4.2 to enhance understanding.

4. Addressing the feedback from **Reviewer 4qEE** and **Reviewer BVRU**. **we have added detailed pointers in Section 5** that direct readers to the appendix for comprehensive training details, and **expanded our ablation studies**, which can be found in Appendix C. These studies focus on two primary areas: (C.1) generator initialization, and (C.2) the impact of including or excluding regression loss during training.

**Change 1** improved the overall presentation of the paper a lot.  **Change 2** aided readers in comprehensively understanding the purpose and functionality of the stop-gradient technique within our model. **Change 3** clarified the notation utilized in Algorithm 1, rendering it more accessible and easier for readers to follow. **Change 4** expanded our ablation studies, which substantially enrich our findings and provide deeper insights into the research.

We appreciate constructive suggestions from all reviewers that help strengthen our draft. 

Best,

Authors of submission #8666
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper suggest a distillation method based on the Flow Matching loss where the idea is to optimize a “generator” $g_\theta$ that maps noise to sample in one step by defining the student data distribution $p_{\theta,0}$ via $g_\theta$ and minimize the FM loss to match it to the teacher data distribution defined by the given pre-trained model. 

Concerns about this paper in its current state include its somewhat limited relative added contribution compared to recent implicit score 1-step distillation methods that take a similar approach but for score parameterization, the quality of the results presented, and lacking/unclear presentation. We still feel the paper is of value as a promising flow distillation method and encourage the authors to revise their paper and resubmit in the future.