Correcting Flows with Marginal Matching

Authors: Meital Bojan, Nadav Bojan, Ev Zisselman, Aviv Tamar
Primary area: generative models
Keywords: flow matching, rectified flow, diffusion, generative models

Abstract:
Flow matching models, ODE-based generative models, generate samples by gradually morphing a simple source distribution into a target distribution. 
In practice, these models still fall short of perfectly replicating the target distribution, mainly due to imperfections of the learned mapping. Previous work mainly focus on alleviating discretization error, which rises from sampling a continuous trajectory with a finite number of steps. In this work we focus on prediction error, an error that is inherent in the model. 
Our main contribution is identifying a trajectory that complies with the imperfect flow model and leads exactly to the target distribution. Based on this finding, we propose Marginal Matching&mdash;a simple inference-time correction scheme to steer the generated samples in the direction of the data.
This scheme proves to reduce a bound on the distance between the data and the learned distribution, motivating two different implementations for the correction function. We show that our proposed method improves sample quality on CIFAR-10 and ImageNet-64,  with minimal overhead in computation time, or non at all when applying approximated correction.

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

Reviews and discussion:
<thread>
creation date:
1730143309448

thread_type:
official_review

<thread_object>
creator: reviewer_k9va

<content>
summary:
This paper proposes a mitigation strategy for prediction errors in flow-based generative models. While previous research has often focused on minimizing discretization errors (by optimizing the mapping between Gaussian noise and the target data distribution), this work is the first to specifically target the prediction errors in flow-matching models.

The authors propose leveraging time reversibility in flow matching models. By implementing a corrective mechanism that modifies the generated sample trajectories during inference, the model can reduce errors between the generated and target distributions. This correction, applied intermittently with the ODE solver steps, is theoretically shown to minimize the Wasserstein distance between the generated and target distributions, under certain assumptions.

The paper presents two implementations of this correction mechanism: one that uses a score model and one that is based on a classifier. Empirical results show that this method improves generation quality on CIFAR-10 and ImageNet-64, achieving better FID scores with fewer sampling steps, especially in the former. The contributions of the paper include a novel time-reversal framework, theoretical analysis, and practical implementations to enhance flow-matching model performance.

strengths:
- The paper presents an approach to improve flow models by enforcing that the forward and backward marginals need to match. This is a sensible thing to try.

weaknesses:
The presentation of the paper can be improved. Some specific points that need to be changed:
- The statement of Theorem 3.3 is too long. It should be broken up.
- The explanation on the score model approach is confusing. It is not explained in detail in Section 3.2.1: the sentence “We assume that the forward marginals are there noisy versions…” is not clear. Section A.3.2, which is where the approach should be explained, contains an obscure statement that does not clarify how the training is performed. Then, Section A.10.3 does contain the training loss for the score model, but the reason to introduce $\sigma$ is still unclear.
- In App. A.5, what does percentage mean?

The experiments presented show some improvement with respect to the baseline models, but some questions arise: 
- The improvement relative to the baseline models is stronger for CIFAR-10 than it is for ImageNet64. Performance on ImageNet64 is much more indicative of performance on models that practitioners care about. In fact, Table 2 shows that the baseline OT-CFM with 16 NFE obtains FID 15.11, while OT-CFM with Score at 16 NFE (6 of those corresponding to the score model) obtains either 15.57 or 15.69 depending on the variant. Even if the score model is cheaper to evaluate, it is hard to make a case for using OT-CFM Score given that the obtained FID is higher. 
- It is also surprising that the FID with 16 NFE is higher than with 12 NFE, for both score variants. This goes against the intuition developed in the paper.
Experiments on larger scale text-to-image models would help clarify the usefulness of the suggested method.

questions:
- Wouldn’t it make sense to fine-tune the FM vector field such that $\frac{1}{N}v_{\theta_{FT}}(t_n,x_{t_n}) = \frac{1}{N}v_{\theta}(t_n,x_{t_n}) + \alpha_n h_{\psi}(t_n, x_{t_n})$? That way, we don’t need to evaluate $h_{\psi}$ separately from the flow matching vector field, which means that we can arguably use the correction at every single time step without incurring additional compute costs. It seems that should work well in the cases where parallel sampling works well.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730664012793

thread_type:
official_review

<thread_object>
creator: reviewer_ekwc

<content>
summary:
The paper introduces a method to correct the path trajectories of continuous normalizing flow models trained with flow-matching. Instead of focusing on improving ODE solvers for potentially poorly fitted models, as much of the existing literature does, this work aims to train networks that directly guide the sample trajectory by correcting the model-fit failures at each step with a (residual) network $h_\theta$.

As the authors note, although the general idea has been previously explored for continuous normalizing flows [A], this paper differentiates itself by focusing on frameworks that train normalizing flows using conditional flow matching.

I have two main concerns:

1. **Theoretical Presentation**: Some theoretical sections are presented poorly and the notation could be improved. Portions of the theory seem inconsequential and do not tie back to the story (Separation of A and B in Theorem 3.3), and could be safely removed, moved to the appendix, or generally improved for easier readability. Above that, I do have some concerns about the assumptions for one of the theorems.

2. **Lack of Evidence in Large NFE Settings**: The paper initially distinguishes between improving flow matching models by $(i)$ improving model fit and $(ii)$ refining the ODE solver. However, the authors do not provide evidence that their method is beneficial in settings with a large NFE. Specifically, if the generative performance is poor due to model-fit issues, merely increasing NFEs should not significantly help; however, the results show that increasing NFE without any correction is still getting the same performance as doing correction but with lower NFE. Therefore, to convincingly advocate for their method, experiments in large NFE settings are essential.

The second concern is crucial for connecting the results back to the paper's motivation and introduction. Indeed, the paper even has a line (L142) mentioning that "while a larger $N$ reduces discretization error, it also leads to a greater accumulation of prediction error", so why do we not have comparisons in precisely these large NFE settings?  

My interpretation based on the experimental results is that adding the guidance $h$ seems to be more similar to distillation methods that allow the model to take larger steps, as opposed to inherently fixing model-fit errors. I also have some more detailed points and a suggested ablation which I will present in the "questions" section.

All in all, the paper does seem to introduce a nice method at its core, but without sufficient experimental evidence and improvements in writing, it is not ready for acceptance!

strengths:
1. The paper aims to address the problems associated with image generation in possibly the most performant variant of generative models. Any improvements to flow matching approaches are therefore highly useful!

2. The background, introduction, and discussion of related work are thorough and clear.

3. The main idea is clearly presented with a nice figure that helps understanding.

4. The theoretical contribution appears sound, though not all of them directly factor into the methodology.

5. I appreciate the amount of detail and extra context that is included in the Appendix for the ablations and comparisons.

weaknesses:
1. I find that separating Theorem 3.3 into parts A and B is tangential to the story and overly confusing. In reality, we do not have full control over our correction network and its Lipschitz constant. Therefore, we can never determine the best scheduling. This section seems like its being theoretical for its own sake! It might be clearer to simply present Lemma A.2 of the appendix in its most general form:
$$W_2(p^b_t, p^f_t) \le W_2(p^b_{t_0}, p^f_{t_0}) \cdot e^{L(t-t_0)}$$
and say that improving the Wasserstein distance on the RHS for $t_0$ can effectively bound the Wasserstein distance on the LHS, especially for $t$ that is sufficiently close to $t_0$. I don't think A, B, and examples 3.4 and 3.5 are particularly insightful when it is not directly factored into the decisions made in the experiments. The results in A and B can still be included, but in the appendix.

2. The parallel sampling section seems slightly oversold! To my understanding, while both forward passes can be done in parallel, it cannot be done in one batch because the forward call is on different methods. Could you please provide a time comparison between parallel and serial sampling on one experiment with the hardware that you have?
 
3. The statement of Lemma 3.6 seems to spill over to the rest of the main text and I generally do not agree with the base assumption that $p_t^f = p^b_{t, \sigma}$ which is the main driver for Lemma 3.6. Please let me know if I am misunderstanding this!

4. I don't find the comparison between this method and Dai and Wipf [B] appropriate! [B] trains a VAE on VAE to fix problems associated with the dimensionality mismatch between the data manifold and the manifold induced by the (first) VAE. That is not a concern in flow-matching and diffusion models as these models are known not to suffer from the manifold mismatch difficulties as much.

5. Although FIDs are still being widely used for evaluation, there have been clear flaws associated with them and the simplistic Inception network [C]. Please use DinoV2 Frechet Distances for the comparisons from [C], in addition to the widely used FID metric.

6. Please also provide evaluations "matching" the same NFEs in the corresponding non-corrected models.

### Minor points

1. I personally do not agree with the notation abuse of rewriting the conditional probability flow $p_t(x | z)$ as the marginal probability flow $p_t(x)$; it is highly confusing in my opinion.

2. Rather than introducing the new probability flows $\nu_t$ and $\mu_t$, in theorem 3.3, please consider using the same $p^b_t$ and $p^f_t$ for reduced notation overhead, and then restate the theorem in full formality for the appendix.

3. (nitpick) In Eq. (8), $t$ should be a sub-index of $u$.

questions:
1. Instead of correcting the marginals, what happens if you simply define a new network $\nu_\theta(t, x) := h_\theta(t, x) + v(t, x)$ and fine tune using the CFM loss? That seems like a more direct way to fix prediction error, so I'm curious if there is any gain from using the relatively more complicated algorithm you have suggested.


### References

[A] Chen Xu, Xiuyuan Cheng, and Yao Xie. Normalizing flow neural networks by jko scheme. Advances
in Neural Information Processing Systems, 36, 2024.

[B] Dai, Bin, and David Wipf. "Diagnosing and enhancing VAE models." arXiv preprint arXiv:1903.05789 (2019).

[C] Stein, George, et al. "Exposing flaws of generative model evaluation metrics and their unfair treatment of diffusion models." Advances in Neural Information Processing Systems 36 (2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730718564557

thread_type:
official_review

<thread_object>
creator: reviewer_isnq

<content>
summary:
This paper proposes the refinement of the flow-based/diffusion models with marginal matching, which delves into the gap between the forward and reverse process to mitigate the *prediction error* of the model, which is the gap between the optimal flow $u*$ and the estimated $v_\theta$. The principle behind the marginal matching is that since the model cannot be learned perfectly and the forward and backward marginal (path) must have some gap, there should be some corrector to move from the forward to the (correponding) backward marginal.
To implement the marginal matching, the paper piloted two models: score-based and robust classifier. Score-based marginal matching assumed the forward marginals to be the noisy version of the backward marginal, and the robust classifier first constructs the classifier discriminates the forward-marginal particle to the backward-marginal particle (using both data) and get the gradient to move from the forward to the backward marginal trajectory. 
With this corrector model with marginal matching, the generation quality from the image datasets are improved with a few steps.

strengths:
* The idea of marginal matching correction is sound, and can be ubiquitously applied to other generative model approaches.
* The convergence property from the marginal matching is concretely headlined with Wasserstein distance bound. 
* Moreover, the interpolation part which says that using the data and approximate source rather than the full trajectory for training resolves some of my interests, which is the multi-marginal matching from the partial trajectory rather than the full trajectory.

weaknesses:
* My main concern is that the idea of the marginal matching correction is not directly used in the score model, as the forward marginal is not directly connected to the backward marginal but is considered as the noisy version of the backward marginal. If so, (as I mentioned in the question) the corrector part is the external-model version of the existing work. As this paper implies that the score version of the corrector works better than the robust classifier version, it would be better if this issue is resolved.

questions:
* For my understanding, the corrector model $h_\psi$ constructed by the score model is the same or similar to the predictor-corrector proposed in [Song et al, 2021]? If it is so, the score model part of the marginal matching is the flow-based equivalent of the existing work with more theoretical background with Wasserstein distance bound and external corrector model.
* The result in Table 1 will be more concrete if the RK4 solver includes NFE=24, 28, and 32 (as a single step uses 4 NFE). I'm curious that the paper mentioned that the NFE of a single RK4 step is 2
* How is the corrector step $\alpha$ chosen? This seems that the choice of $\alpha$ heavily affects the performance of the corrector. Are the hyperparameters in Table 10-12 hand-crafted?

```
[Song et al., 2021] Score-Based Generative Modeling through Stochastic Differential Equations, ICLR 2021
```

---

I recommend borderline acceptance in the current stage of review: this can be varied in the discussion period, and additional question can be added.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730750575080

thread_type:
official_review

<thread_object>
creator: reviewer_ubfw

<content>
summary:
This paper proposes using corrector steps during the sampling of flow models. Two correction models are proposed: a score model, and a classifier model. If I understand correctly, in practice, the score model is trained with a constant noise instead of being time-dependent. The classifier model should act similarly to a discriminator, but empirical performances seem to favor the score model.

strengths:
- Discusses using additional inference-time compute which could be interesting.

weaknesses:
- The concrete connection between theoretical derivation and proposed algorithm is unclear to me.
- There is also missing information regarding the method.

questions:
- How does Algorithm 1 relate to the idea of the backward marginal?

- Why does Algorithm 1 add noise in line 4? Doesn't this add bias to the sampling when no correction is used?

- How do you choose beta and alpha?

- The proposal for the score model is to a time-dependent score. But in practice, only a constant noise is used. Why is that?

- If you used a time-dependent score and set beta = sqrt(2 * alpha), does your method recover Langevin dynamics as the corrector step?

- What is "Backward Marginals Interpolation"? How is this different from the standard "Score" correction model?

- Why can't this be used with diffusion model sampling (by just replacing line 6)? Intuitively, the choice of corrector steps should be independent to the predictor (ODE/SDE) steps.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732091749322

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We thank the reviewers for their valuable insights and suggestions. We have responded to each reviewer individually. In addition, we provide this general response to address points that were brought up by multiple reviewers. We incorporated all the reviewers suggestions in the revised version of the paper, the  changes and additions are highlighted in brown.

**Presentation:**

We modified the theoretical presentation in the revised paper of both Theorem 3.3 and Lemma 3.6, as suggested by reviewers ekWc and k9VA.
We added the lemma from the Appendix to the main paper and changed Theorem 3.3 and Lemma 3.6 to be informal in the main paper with the formal versions in the appendix. The informal version Theorem 3.3 does not contain parts A and B. 

**Hyper-parameters:**

Alpha and beta are hyper-parameters of the algorithm, tuned through grid search - full details of this search are provided in Appendix A.10.4 of the revised paper, as requested by reviewers ubfw and iSNQ.

*Noise $\beta$:* The configuration with the best FID was selected from a grid search over the interval $[0,0.1]$ with a step size of 0.01.

*Step size $\alpha$:* The configuration with the best FID was selected from a grid search over the interval $[0,2]$ with a step size of 0.05. For the classifier's final step, the grid search was conducted over the interval $[0,0.1]$ with a step size of 0.01.

The step-size heavily affects the performance but it is generally not too sensitive to it.  We tested the sensitivity of our correction models by adjusting their step sizes up and down by $10$%. For instance, where the original step size was 1.0, we tested with both 0.9 ($10$% decrease) and 1.1 ($10$% increase). Here are the results we observed:

| *Correction Model* | C-NFE ↓ | mean FID | std FID |
|-------------------|----------|-----------|----------|
| Classifier | 1 *($n=8$)* | 3.58 | 0.015 |
| | 5 *($n=6,7,8,9,10$)* | 3.49 | 0.011 |
| Score | 1 *($n=10$)* | 3.46 | 0.01 |
| | 3 *($n=0,5,10$)* | 3.40 | 0.017 |

The mean and std is computed over all the FID results from the different configurations obtained by the step-size adjustment described above. As shown in the table, the mean closely matches the result reported in the paper and exhibits minimal variance. The performance remained largely consistent across different step sizes.

**Connection between Theory and Practice:**

Following the questions of reviewers ubfw and ekWc we emphasize that the Theorem 3.3 as well as the examples (linear and multiplicative reductions) fit the idea of the paper and are important to our story as they motivate the use of different correction steps as illustrated in Figure 5. Figure 5 demonstrates that some correction steps exert more influence on the performance than others. This supports the findings of the theorem as it shows that depending on the error reduction, different correction steps influence the reduction to the final bound differently. While the exact error reduction is unknown in practice, this theorem explains why different time-step correction behave differently and motivates us to seek the influence of different time-step corrections, which can be determined empirically. Additionally, this theorem illustrates why the corrected inference algorithm is sound in principle.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper presents a method to reduce the prediction error in flow matching. The basic idea is to learn a corrector to reduce the gap between the marginal distributions associated with the forward and backward trajectories. Theoretical analysis is provided for the overall framework and two practical algorithms are presented. Several reviewers raise questions on the theoretical result and its connections to the algorithms. They also provide suggestions to improve the presentation. Some reviewers think the empirical results are not sufficient to demonstrate the advantages of the proposed method. There exist gaps between theory and algorithms. Some assumptions appear to be strong. For instance, the assumption in Section 3.2.1. The authors are encouraged to improve their results and presentation following the reviewers’s comments.