Correcting Flows with Marginal Matching

Larger Scale Text-to-Image Models:

We agree that results on large scale text-to-image models are very impactful, but we simply do not have the (very extensive) compute resources to carry out such experiments. We kindly refer the reviewer to some recent studies on flow matching such as [A,B], which were all very impactful, yet also did not include large scale text-to-image experiments, and hope that the reviewer agrees that our work can be evaluated within the (not so small) compute budget that we have.

"Parallel" Network:

This is a great idea! It would leverage parallelization more effectively, though it requires additional flow model training. As this approach would bake in a specific correction configuration, adapting to new configurations would necessitate further fine-tuning. It would be interesting to investigate it as a future work.

References:

[A] Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative mod- els with minibatch optimal transport. Transactions on Machine Learning Research, 2024.

[B] 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.

Thank you for your review and your suggestion for improving the presentation!

Theorem 3.3:

Thank you for your valuable suggestions, we changed the presentation of the Theorem in the revised version. Please refer to the general comment for additional details.

Score Model:

We are sorry for the confusion, we revised the score model section in Sections 3.2.1 and A.3.2 of the paper with further explanations about the approach including a pseudo-code, (highlighted in brown).

Score models are a type of generative models that work by gradually cleaning up random noise into meaningful data through a process guided by the score function - the gradient of the log probability density. These models use varying noise levels during generation: they start with very noisy data where the score function provides stable but coarse guidance, then progressively reduce the noise while using increasingly refined score functions to sculpt the final output. Through this iterative denoising process, random noise is transformed into samples from the target distribution.

By assuming the forward marginals are a noisy version of the backward marginals (with a small $\sigma$), a score model can be used to transverse between them.

Algorithm: Score Model Training

Input: flow model $v_\theta$, score model $h_\psi$, $\sigma$, backward marginals $\{p_{t_n}^b(x)\}_{n=0}^N$

Repeat:

1. $\epsilon \sim \mathcal{N}(0,I)$
2. $n\sim U(\{0,1,..,N\})$
3. $x_{t_n} \sim p_{t_n}^b(x)$
4. Take gradient descent step on
5. $\nabla_\psi \| h_\psi(t_n, x_{t_n} + \sigma\cdot \epsilon) + \epsilon\|^2$

Until: converged

Appendix A.5:

Figure 7 shows a histogram of the differences in VGG features. These features are computed for uncorrected samples (flow only) and corrected samples (flow+correction models) starting from identical Gaussian noises. The $L_2$ difference between the features of the uncorrected and corrected samples is then calculated. The percentage indicates the proportion of samples with the corresponding VGG features difference. The plot demonstrates that most samples' features remain very similar, with minimal changes. The correction models only significantly modify outliers. This aligns with the observation in Figure 1 (toy problem) that correction is only needed for the symmetric difference between the Gaussian noise and $p^b_0$, which corresponds to these outliers. We have incorporated this clarification into the figure caption in the revised paper (in brown).

ImageNet-64 Results:

We had an error in the NFE count, it is fixed in the revised version, we apologize for the confusion. The RK4 scores are counted in multiples of 4. Sampling the OT-CFM flow model with 8-RK4 steps (32 NFE) gives a score of $15.11$, 7-RK4 steps (28 NFE) gives a score of $15.71$ and 24-RK4 steps (24 NFE) gives a score of $16.60$. In contrast, sampling the OT-CFM flow model with 5-RK4 steps and 6 C-NFE (for a total of 26 NFE) results in a score of $15.57$ or $15.69$ (depending whether is was approximated). In any case, our work focuses on the lower-NFE regime where it shows the most improvement, and is the more interesting challenge.

Additionally, the CIFAR-10 models were easier to train as it is a smaller dataset. We believe that with more compute the ImageNet-64 models could perform better.

It is also surprising that the FID with 16 NFE is higher than with 12 NFE - some correction steps are more beneficial than others, as seen in Fig.5. In general, taking correction steps reduces the FID score compared to not taking them. It is interesting that omitting some correction steps yields a slightly superior result. We assume this happens due to errors in the correction model accumulated over the trajectory. Designing a more sophisticated correction model might be able to resolve these issues.

We would like to thank the reviewers for their detailed and constructive feedback.

Presentation:

Thank you for your presentation suggestions, we incorporated them in the revised paper. Please refer to the general comment for additional details.

Theorem 3.3:

Thank you for this point. The relationship between the theory and algorithm is explained in detail in the general comment.

Larger NFEs:

Thank you for raising this point, we apologize for the confusion. In the sentence ”while a larger reduces discretization error, it also leads to a greater accumulation of prediction error” - by greater accumulation we meant that the sum of the prediction error is over more steps, and not necessarily that the prediction error is greater as it normalized by the number of steps, but to the best of our knowledge there is no way of measuring it. We agree that this is not phrased well and confusing and we removed it from the revised paper.

Generally, achieving low FID with a low number of NFEs is always preferred over a high number of NFEs. In our case, sampling with a large number of NFEs will require training and evaluation over more marginals up to approximations - marginals interpolation or using only a subset of the marginals. This makes the training and evaluation more difficult and therefore we did not add it to the paper. While this is worthy of future work, it is beyond the scope of this paper

However, we understand that it is one of the reviewer's main concerns and in order to show evidence in the large NFEs settings within the limited discussion period time we made an approximated evaluation. With the trained correction models (score and classifier) over 10-RK4 steps (40 NFEs) employed in the paper, we evaluated a setting where the base flow model (OT-CFM [A]) is sampled with 50-RK4 steps (200 NFEs), which is a substantially larger number of NFEs. This measures the performance in large NFE setting, but also generalization of the correction models, as they were not trained over these marginals which are more refined, as the ODE-solver is integrated over a larger number of steps. The FID score of 50-RK4 steps of OT-CFM is 3.67. The evaluation of correction step-sizes of the classifier and score correction models are a normalized version (scaled down) of the ones used in the paper, as the correction is more subtle. In the table below are the results.

The results show that it always beneficial to use correction. However, adding more correction steps for the classifier model was not beneficial, we believe this is due to the generalization nature of this experiment. It would be interesting to explore the generalization capacities of the correction models in future work.

Parallel Sampling:

The reviewer raises a fair point. You are correct that the forward passes for the two models cannot be executed in a single batch, since they involve different method calls. The parallel sampling is parallel since no model is dependent on the output of the other model, and therefore it is possible to parallelize.

Cuda has a queue, where the CPU sends tasks to be run on the GPU. The GPU may execute tasks in parallel that are independent of one another, such as our parallel sampling. The CPU may wait when the queue is full or during synchronization events, (i.e. item(), synchronize(), etc). Given a powerful enough Cuda machine, the models could be run in parallel on different GPUs or using advanced parallelism techniques. Parallel calculation on separate devices could be useful when the correction calculation takes more time than the overhead of moving between devices, as is the case for large models.

The extent of parallel processing viability will depend on the specific hardware and infrastructure available. However, on our machines, it is not possible due to the hardware limitations. Nevertheless, we want to emphasize that this approach is parallelizable in principle. We will add this explanation in the Appendix (A.3.1 Parallel Sampling) of the revised paper (in brown) and tone down the presentation.

As the rebuttal deadline is approaching, please acknowledge the response from the authors and reply if needed.

Thank you very much for the response!

1. I do like the theoretical section much more now. Thank you very much for the edits!
2. I especially appreciate you implementing the vector field correction approach. Please add that to the final version of the paper; I think it is a very relevant ablation. I would even recommend tuning that approach more.
3. I look forward to the DinoV2 results in the final iteration of the paper. The paper that originally proposed this showed that diffusion models don't get good FID scores sometimes compared to GANs, even though they obviously outperform them from a qualitative perspective. One thing I noticed from all of your results is that your improvements are really marginal, and as a practitioner, it might not be worth all the additional effort to fine-tune or train my model using your approach. The DinoV2 results might even help your method to stand out more if it is truly improving generation quality.
4. Finally, and this is perhaps a more philosophical concern, but the line between ODE discretization error and model-fit error is much blurrier than you explain in the paper. You can, in principle, many times sidestep the model-fit error if you increase the resolution of your ODE solver. Let me give you a hypothetical example: let's say your model learns the true vector field with a an error that comes from a zero-center distribution. Now, if you increase the discretization steps, the total error rate converges to 0, just as a result of the law of large numbers. I'm obviously not suggesting this case happens in practice, but I think the arguments made between the "model-fit" vs. "discretization" error should be presented in a more convincing manner.

Overall, I would like to increase my score for all the effort the authors put into the rebuttal! But I still have some reservations.

In any case, looking forward to seeing the final version of your paper :)

We thank the reviewer for the valuable input. We address each point separately:

Score Model:

The corrected inference algorithm is a predictor-corrector algorithm, however the score model and the sampling is different from the one used in [Song et al, 2021]. Below are the main differences:

• Model Training: The score model used in [Song et al, 2021] is trained to predict the score function of the data distribution with varying degrees of Gaussian noise (added by the SDE/ODE). In contrast, our score model learns to fit the backward marginals distributions trajectory obtained by integrating over the flow model in reverse. We added a pseudo-code for training the correction score model in the revised paper under the Score Model section (3.2.1).

• Time-Dependency: In [Song et al, 2021] the time $t_n$ refers to different noise levels, where in our model the noise level is constant and it refers to the time of the $n^{th}$ step of the flow model's ODE-solver where the single correction step is applied.

• Predictor-Corrector Distribution: In [Song et al, 2021] the same model is used for the predictor and the corrector sampling in order to reduce the error in the approximation of the distribution this model represents. In our work the predictor is a flow model and the corrector can be a score model. The correction and the flow model approximate different distributions. The correction model is trained after the flow model with the aim of covering its prediction error to gain superior performance (see Appendix A.10.3).

• Corrector Sampling: In [Song et al, 2021] the corrector uses Langevin dynamics with update step: $x_{t_n}+\alpha_n h_\psi(t_n,x_{t_n})+ \sqrt{2\alpha_n} \epsilon$. Conversely, in our algorithm the correction step is: $x_{t_n}+\alpha_n h_\psi(t_n,x_{t_n}+\beta_n \epsilon)$. Sampling flow models is done without intermediate noise, adding noise to this process is not advisable. In practice, we experimented with Langevin dynamics and noticed that it yields inferior results. Additionally, only one correction step is taken at each time-step, which is very different from how Langevin Dynamics is used.

The marginal matching is always used as the score model is trained to fit the backward marginals distributions. The score model directs the samples to the matching time backward marginals, thereby applying it results in a better alignment. This is the reason we expect the score model to improve the performance.

Table 1:

You are correct, thank you for pointing it out and we apologize for the confusion, it is fixed in the revised version. The additional results are added to an extended table in the appendix due to lack of space, a reference to this table is added in the caption of Table 1 in the revised version. For additional information please refer to the general comment.

Corrector-Step:

Thank you for raising this point. The hyper-parameter $\alpha$ was explored through a grid search. Please refer to the general comment for additional details.

Thank you for your feedback and time. Here are our responses:

Connection between Theory and Algorithm:

Thank you for raising this point, for an elaboration on the connection between the theory and the proposed algorithm, please refer to the general comment.

Algorithm 1:

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

The backward marginals are manifested in the training process of $h_\psi$, the correction model, which learns to align the forward with the backward marginals. We added a pseudo-code for training the correction score model in the revised paper under the Score Model section (3.2.1), for more details about training the correction models see A.3.2 and A.3.3. Algorithm 1 is the corrected inference time algorithm and assumes it has access to $h_\psi$. On lines $5$ and $8$ in the algorithm correction steps are taken.

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

Thank you for this question! You are correct and we did not add any noise to the sample when no correction steps were taken. We fixed the pseudo-code in the revised version, the changes are highlighted in brown, we apologize for the confusion.

How do you choose beta and alpha?

Thank you for raising this question. Alpha and beta are hyper-parameters that were explored with hyper-parameter tuning - using a grid search of the values. Please refer to the general comment for additional details.

Score Model:

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

This is a great question. The time-dependency $t_n$ of the score model $s_\psi(t_n, x_{t_n})$ refers to the $n^{th}$ step of the flow model's ODE-solver where the correction $s_\psi$ is applied. This is in contrast to the time-dependency in the original score model which refers to the noise level. The score model is trained with a constant value of noise with time dependency over the ODE-solver time-steps.

In Appendix A.3.2 (line 1126) we elaborate on the choice of a score model with a constant noise level. We emphasize that while more sophisticated approaches like annealing score models with time-varying $\sigma$ exist, they were unsuitable for our needs. This is because our inference process that starts at a forward marginal or an intermediate point where we lack information about the current and next noise levels. Therefore, we cannot use these models effectively, as they assume knowledge of the noise levels.

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

It would not be the same, assuming $\beta_n=\sqrt{2\alpha_n}$. In Langevin Dynamics the update step is: $x_{t_n}+\alpha_n h_\psi(t_n,x_{t_n})+ \sqrt{2\alpha_n} \epsilon$, where in our algorithm the update step is $x_{t_n}+\alpha_n h_\psi(t_n,x_{t_n}+\sqrt{2\alpha_n} \epsilon)$. As you mentioned before, adding noise to the sampling of the vector field adds bias as sampling flow models is done without intermediate noise. In practice, we experimented with Langevin dynamics and noticed that it yields inferior results. Additionally, only one correction step is taken at each time-step, which is very different from how Langevin Dynamics is used.

Backward Marginals Interpolation:

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

In order to keep the training simulation-free all the marginals are stored ahead of time. The ”Backward Marginals Interpolation” is an approximation method for the backward trajectory that saves storage. In the standard training of score correction model, it trains over the exact backward marginals computed with the ODE-solver. In contrast, the ”Backward Marginals Interpolation” version trains over an approximation of that trajectory. This approximation is obtained by interpolating the data $p^b_1$ and $p^b_0$ (approximate source distribution). In the standard case $p^b_t$ is calculated with an ODE-solver, however in the ”Backward Marginals Interpolation” case $p^b_t$ is an interpolation with coefficient $t$ (line 420 in the paper).

Diffusion Model Sampling

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.

This is an interesting question, and we briefly mentioned it in the Conclusions and Limitations section. In general, diffusion models are "approximately invertible" where flow models can be assumed to represent an invertible vector field as they are generally straighter, and our method depends on it. It would be an interesting direction to explore the extension of our method to these models.

Extended Table 1:

We added additional results to the "Flow models RK4 Results" table in Appendix A.6.6 and reference it from the caption of Table 1 in the main paper, as requested by reviewers iSNQ and ekWc. These are the results:

Correction steps improve performance with minimal additional NFEs, we outperform the baselines in the small NFE regime. With 43 (score)/46 (classifier) NFEs, our correction models achieve the same or better FID scores than flow models using 60 NFEs, representing a 23-28% reduction. Even a single correction step significantly reduces the FID.

References:

[A] Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024.