Flow matching achieves almost minimax optimal convergence

Thank you for your helpful comments. We provide our replies below.

Q1 (comparison between FM and DM): More precisely, Theorem 9 tells that the derived upper bound (almost) attains the universal minimax lower bound at $\kappa=1/2$ in terms of the convergence rate. For other $\kappa > 1/2$, the derived upper bound does not attain the universal lower bound, and there is still a possibility that the convergence rate for $\kappa > 1/2$ can be improved. To reflect this situation, we would improve the expression as “Both FM and DM can attain the same almost optimal minimax convergence rate for generalization error” (line 230)

Q2 (difference from Oko et al. 2023): As detailed in response to nHx4 (W1), there are two significant differences from Oko et al. 2023 in terms of proof techniques.

(1) To relate the Wasserstein distance and the squared difference of vector fields (scores in the case of DM), the Alekseev-Grobner Theorem is used for FM, while Girsanov’s Theorem is for DM. This causes a completely different argument to bound between the final ODE solution and the true probability $W_2(P_0, P_{T_0})$; the analysis of FM requires the uniform Lipschitz constant in Lemma 10, which is based on the form of $\sigma_t$ and $m_t$ considered in this paper and the assumption (A5).

(2) The extension of the paths with $\sigma_t$ and $m_t$ needs more elaborate arguments on the analysis. As a result, we can obtain the influence of the decreasing rate $\kappa$ in $\sigma_t \sim t^\kappa$ on the upper bound of the convergence rate.

Q3: Extension to the non-Gaussian case is an interesting problem, and in fact, we are working on this extension. In the current proof, special properties of Gaussian distribution help us make bounds at many places. Mathematically, it is challenging to extend it to general distributions.

Q4: We state at the beginning of Section 2 that we use $P_a$ for probability distribution and $p_a$ for its density function. We have rewritten the statement of Theorem 1 in the revision.

Q5: This assumption is not so restrictive. Because $\sigma_t$ is a decreasing function of $t$ and should approach 0 as $t\to 0$, it is natural to set it in the form $t^\kappa$ around 0. In practice, the OT flow uses $\kappa = 1$ and diffusion flow $\kappa = 1/2$.

Q6 (typo): We have fixed all the typos in reviewers’ comments.

Q7 (line 224; the relation between FM and DM) See our Global Responses.

Weakness: Thank you for your suggestion to include a figure for explanation. We agree that a visual explanation could improve accessibility, and we have included a figure to illustrate the basic idea of the time division in Appendix E and mentioned it in line 458 in the revision.

Q1 (probability estimation): FM methods do not estimate a density, but provide samples that approximate the true probability. We thus used the terminology probability estimation.

Q2: (i.i.d. samples) We agree. We have reflected all of this in the revision.

Q3, 4, 6 (typos): We have fixed as many typos as possible in addition to those you pointed out.

Q5 (support) The cube size does not affect the result. We have mentioned it in line 166 in the revision.

Thank you for your comments. We write our replies to them.

(1) Novelty: As detailed in our response to Reviewer nHx4 (W1), the dependence of convergence rates on the choice of $\sigma_t$ and $m_t$, and the bound around the final time point of ODE are significantly different from the existing literature and novel to the best of our knowledge.

(2) The setup covers various FM models, including the most popular OT-CFM (Tong et al 2023). Note that, as in the reply to Q3 of yJWQ, the way of constructing a coupling does not affect the analysis in the paper. We would greatly appreciate it if you could specify which aspects appear as toy models to address your concerns better.

Q1: As discussed in line 446, the constant $R_0$ can be taken so that $R_0 \geq (s+1)/min(\kappa,\bar{\kappa})$. However, we cannot know $s$ in many practical cases, so it is not easy to set $R_0$ as this minimum value.

Q2: $\sigma_{[\tau]}=1-\tau$ ($\kappa = 1$) is the most popular choice of FM, and $\sigma_{[\tau]}=(1-\tau)^{1/2}$ ($\kappa = 1/2$) is the diffusion path (Lipman 2023, e.g.). The choice in Theorem 1 covers practically popular ones. If my answer does not reply to your question, please ask further.

Q3: For $\kappa = 1/2$, the upper bound of Theorem 9 is of the same rate (up to an arbitrary positive $\delta$) as the lower bound of Wasserstein distance in the minimax sense, as shown in Proposition 2. So, the bound is tight for $\kappa = 1/2$. For other values of $\kappa$, there are no theoretical arguments for the tightness of the obtained convergence rate. As discussed in (2) above, the result holds for a wide range of FM, such as OT-CFM.

Typo: Thank you for pointing it out. We have fixed it in the revision.

Thank you for your helpful comments. We reply to your comments.

W1: Overlap with Oko et al 2023.

While we acknowledge some overlap in mathematical techniques with Oko et al. 2023, there are significant difference also as listed below.

(1) We have revealed the role of parameters $\sigma_t$ and $m_t$ on the upper bound of the convergence rates. No such results have been derived in the literature for FM and DM. To obtain the results, the mathematical techniques for the proof of main theorems (Appendix C) require significant refinements from Oko et al 2023.

(2) As detailed in the global comments, the generalization analysis of DM and FM have significant difference. Beyond this difference, you might notice that the bound $W_2(P_0, P_{T_0})$ requires completely novel arguments using the uniform Lipschitz constant in Lemma 10, which are based on the form of $\sigma_t$ and $m_t$ considered in this paper and the assumption (A5).

(3) Appendices A4 and A5 summarize known results (e.g., B-splines and Gaussian integral bounds) and do not intend novelty. We have explicitly stated this for clarity in the revision.

(4) We do not think $\sigma_t = \sqrt{t}$ is an empirically popular choice for FM unlike DM. In fact, the most popular choice of path construction for FM is $x_t = (1-t)x_1 + t x_0$, which corresponds to $\sigma_t = t$ and $m_t = 1-t$ ($\kappa = 1$).

W2: degree of growth of $\sigma_t$

Section 3.1 discusses how $\sigma_t$ acts as a smoothing parameter, analogous to the bandwidth in kernel density estimation. The early stopping time $T_0$ depends on $n$ ($T_0 = n^{-R_0}$ in the paper), and the growth rate $\kappa$ in $\sigma_t=t^\kappa$ controls the smoothing parameter at the stopping time $t=T_0$. It is well known that the smoothing parameter of a nonparametric estimator essentially affects on the convergence rate. This is why $\kappa$ is important for the convergence rate and the performance of the FM estimator.

Q1: other than Gaussian distribution

This is an interesting question, and in fact we are working on this extension. In the current proof, special properties of Gaussian distribution help us making bounds at many places. Mathematically, it is challenging to extend it to general distributions.

Q2: Extension to TV

As discussed in Global Responses, it is not straightforward to extend the (almost) minimax optimal convergence rate to the TV distance. This also illustrates the significant difference between the ODE and SDE for analyzing the generalization bounds.

Q3: Heuristics for FM such as OT-minibatch.

The current analysis does not depend on the joint distribution of the source $N(0,I_d)$ and the target $P_{true}$. Thus, for any algorithms to match samples in the two distributions, including OT-minibatch, the theoretical results hold. We have mentioned it at line 158 in the revision .