Paper: Bellman Optimal Stepsize Straightening of Flow-Matching Models (ICLR-2024)

Reviewer comment (Presentation):
Although the main idea of this paper is interesting, this paper gives me the initial impression of being incomplete, with this incompleteness of the presentation of both the methodology and experimental sections.

Author response:
We appreciate the reviewer’s comment and have added various elaborations in the main text. Please refer to the improved version of our submission with added elaborations, new results with additional datasets.

Reviewer comment (Reproducibility):
For method Sec.3.3, the authors did not provide enough details about how to address the dynamic programming (DP) problem in practice. From the experiments, different datasets should share different sampling schedules. Moreover, how many samples are used to calculate the optimal scheduler? What will the optimal scheduler be like, if using different initial noises? The analysis of the generalization of the optimal scheduler selected by DP is very important.

Author response:
We have added a section in the Appendix B with a pseudocode that provides details on our Dynamic Programming algorithm. The time complexity of this algorithm is $O((K^{max})^2 \times K)$. The choice of $K^{max}$ is crucial, aiming for the Euler sampling method with $K^{max}$ stepsizes to precisely replicate the trajectory of the Ordinary Differential Equation (ODE). Typically, $K^{max}$ falls within the range of 100 to 1000, ensuring accuracy. Given this range for $K^{max}$ and the variable $K$ ranging from 2 to 1000, the algorithm executes within seconds in all scenarios. Moreover, our solved problem is fast sampling in flow-matching models which focuses on relatively small values of K (<20), so the intensity of the dynamic programming algorithm is not a worrying problem even in the case that $K^{max} > 1000$.

We have added a section in the Appendix A on experimental details. In short, we used 100 samples for calculating the Bellman steps for all dataset, as we did not observe any improvements in FID when using more than that number of samples for calculating optimal step-size scheduler. Moreover, we aim to use as low computational resources as possible to keep it inline with the goal of our work. 
We agree with the reviewer that an analysis of the generalization of the optimal scheduler/sampler with DP steps is very important. However, we think it requires setting up plenty of additional definitions and assumptions related to learning theory (e.g. neural network class of the velocity networks, assumption about data and noise distributions, smoothness of the velocity functions, etc.) and therefore we consider this beyond the scope of the current work. However, we find this a very interesting direction to explore in the future work, and we thank the reviewer for helpful comment.

Reviewer comment (Novelty):
Missing necessary discussions about diffusion sampling schedules in related works (such as [1]). Moreover, the related works are too simple, and only include some matching-flow methods.

Author response:
We thank the reviewer for pointing this out, and we have updated the related work sessions accordingly. In short, the focus of our work is on improving the ODE sampler of flow-matching models. While flow matching is a special case of diffusion probability ODE flow, flow matching having different formulations for its probability path (linear interpolation $X_t = (1-t)X_0 + t X_1$ makes it possible for redressing/reflowing to work. The principle of measuring the straightness of the trajectory curve is not applicable to general affine formulation from diffusion probability ODE. Nevertheless, we think that extending the philosophy of calculating optimal step sizes/noise schedulers with dynamic programming is an interesting direction, and we will leave it for future works.

Reviewer comment (Experiments):
Insufficient qualitative comparisons: the authors only provide qualitative comparisons of human faces.

Author response:
We have added more detailed qualitative comparisons in the main text as suggested by the reviewer. Please refer to [Figure 5].

Reviewer comment (Reproducibility):
No details about the velocity network re-alignment training are provided.

Author response:
We are unsure if we understood correctly that the reviewer asked for details and training objective of the straightening velocity network. If this is the case, then the reviewer can refer to Section 4. On the other hand, if instead the reviewer wanted to know more about the details of velocity network retraining (GPU hours, training iterations, etc.), please kindly refer to our Appendix A (Experimental details). Below is a copy of the paragraph. 
 
The pre-trained models are finetuned by three methods including Uniform-Reflow, Distill-k Reflow, and BOSS  in 12000 iterations. An iteration is the passing and backpropagation process for a batch including 15 samples. Due to the similar cost of training between finetuning methods, we report the average GPU hours consumed on each pre-trained model up to 12000 iterations, using NVIDIA RTX A5000.

* CIFAR10: 3.56 training hours.
* CelebA-HQ: 10.35 training hours.
* LSUN-CHURCH: 13.43 training hours.
* LSUN-BEDROOM: 14.23 training hours.
 * AFHQ-CAT: 9.30 training hours.

Reviewer comment (Experiments):
The baseline is too weak, as it only compares quantitatively with the fixed-step size ODE solver euler, while it doesn't compare with adaptive step size ODE solvers such as dopri5 or rk45. This comparison should be included in Table 1.

Author response:
We thank the reviewer for the suggestion. We have added more baseline methods and additional datasets to showcase the effectiveness. The quantitative results can be found in [Table 1 & 2], Figure 4 of our revision of the paper, and the qualitative results can be found in Figure 5 of our paper.

Reviewer comment (Experiments):
The optimization needs to be conducted on a case-by-case basis.

Author response:
We agree that our method requires solving additional dynamic programming problem to find sampling step sizes, compared to the many other training-free diffusion/flow matching samplers that have closed-form formulation. However, we want to point out that in most cases, there are no guarantee  that these closed-form formulas are optimal, as they are mostly on heuristic rules.

Reviewer comment (Evaluation):
The use of the dynamic programming algorithm using Gorubi may result in slow performance, rendering the method impractical.

Author response:
We agree that our method requires solving additional dynamic programming problem to find sampling step sizes, compared to the many other training-free diffusion/flow matching samplers that have closed-form formulation. However, we want to point out that in most cases, there are no guarantee  that these closed-form formulas are optimal, as they are mostly on heuristic rules.

Reviewer comment (Evaluation):
Optimization time is also not discussed in this paper.

Author response:
Indeed this is missing from the previous version of our work. We have added a paragraph in the Appendix A & B regarding CPU+GPU hours of calculating Bellman steps and GPU hours of retraining/redressing the pretrained velocity networks.