On the Guidance of Flow Matching

Thank you for your detailed review and for acknowledging the contribution and potential impact of our work. We will address your concerns in the following:

Q1: Misleading expression of "guidance"

We apologize for the confusion. We will switch to 'energy guidance' to further enhance clarity.

Q2: The assumption of \mathcal{P}=1 seems to contradict the claim in the paper that the guidance is substantially different from diffusion guidance.

1. From the empirical perspective, we argue that $\mathcal{P}=1$ is a valid choice for dependent couplings in realistic datasets. As we show in the table below, the learned VF of the OT-CFM in CelebA and uncoupled CFM have small relative error at all flow time steps, so their guidance VFs are approximately the same, validating our approximation of $\mathcal{P}=1$.

2. Theoretically, with a slow-varying $J$, the approximate $\mathcal{P}=1$ holds for any coupling. A detailed discussion can be found in our response to reviewer YEWM (Q2).

3. We would like to emphasize that in addition to dependent couplings, our framework is also substantially different from diffusion guidance because (a) it extends to any source distribution and conditional probability paths as is also recognized by reviewer erey and reFm (b) a different theoretical framework is provided, as our derivations do not start from score-based models.

Q3. The Monte Carlo guidance has the limitation of lower sample efficiency, so is it practical?

1. We argue that $g_t^{\text{MC}}$ is practical (at least in some tasks) as it shows satisfactory performance in the widely used offline RL task of D4RL. The dimensionality ($23*20$) is already high enough in many practical generative modeling tasks, including generative modeling and molecular structure generation.
2. Besides, there are many variance-reducing MC techniques, such as importance sampling. Specifically, we can use another guided VF to sample from an alternative distribution $\tilde{p}$ such that $\frac{e^{-J(x_1)}}{\tilde{p}(z)}$ has lower variance. For the details of the new $g^{\text{MC-IS}}$, please refer to our response to reviewer YEWM (Q1). It should be noted that this estimation is still unbiased. We conduct experiments to validate the effectiveness of this method, and the results in the table below showcase a huge improvement comparable to the state-of-the-art methods on the image inverse problem, where the vanilla $g_t^{\text{MC}}$ fails.

Q4. Why do many methods perform poorly in toy examples?

Actually, the toy examples are not trivial: sampling from $\frac{1}{Z}p(x_1)e^{-J(x_1)}$ requires guidance theoretically exact. Otherwise, the samples will be distorted or exhibit mode collapse. Gradient-based guidance methods ($g^{\text{cov-A}}$, $g^{\text{cov-G}}$) rely on rough approximations; while approximate guidance ($g^{\text{sim-MC}}$, CEG) produces biased guidance vector fields, leading to distorted sample distribution.

Q5. Why do these methods exhibit significant improvements when applied to more complex problems?

Not all tasks require exact energy-guided sampling as in the toy example. There are tasks where sampling from a distorted, rather than exact, target distribution $\frac{1}{Z}p(x_1)e^{-J(x_1)}$ may be practically useful, as long as both $e^{-J(x_1)}$ and $p(x_1)$ are high. The image inverse problem in our experiment is one instance where the approximate methods demonstrate improved performance. On the contrary, the offline RL task is conditioned on different initial conditions, necessitating stable guidance, so theoretically exact energy guidance is more robust.

Q6. More explanation is required for classfier-free guidance in the current framework (End of section 3.5, line 347-349)

In our framework, directly following our definition of the general guidance VF, the classifier-free guidance is simply subtracting the original VF from the conditioned VF (which also extends previous CFG to dependent coupling and arbitrary source distributions). However, as you pointed out, our training-based guidance is actually the flow matching version of the diffusion classifier-based guidance, rather than the classifier-free guidance. We will revise this paragraph and move this to the discussion section, as well.

Q7. The calculation of $g^{\text{local}}$ is unclear.

You are correct. We will add the derivation of $g^{\text{local}}$ in the main text in the revised version.

Q8. Other typos.

Thank you. We will fix these. For 4, the conditional probability path is a design choice in flow matching, and we have thus assumed the new conditional probability path to be identical. We will make the assumption more explicit.

Thank you for acknowledging the theoretical and empirical soundness, the presentation, and ample details of our work, as well as the contributions and potential benefits for future works in the field. We would like to address your concerns in the following:

Q1: The correctness of guidance VF assumes the trained VF to be identical to the true VF.

Thank you for raising this question. We indeed made this assumption throughout the paper. We will make it explicit in the revised manuscript.

Q2: There are assumptions on the specific type of guidance VF and coupling, limiting the framework's generality.

Thank you for your comment, and we appreciate your recognition of the overall generality of our guidance. We will add a more detailed discussion on these assumptions we made in the revised version, but would like to make the following clarifications:

For the conditional VF assumption, it is indeed an interesting research question to explore other possible conditional probability paths that may have different advantages, e.g., enhancing "straightness" of the VF for accelerated sampling. Nevertheless, our assumption of conditional VF is natural, and it allows us to simplify the form of guidance, covering many existing guidance methods.

For the assumption of $\mathcal{P}=1$ in the dependent coupling case:

1. Empirically, we argue that $\mathcal{P}=1$ is a valid approximation for dependent couplings in realistic datasets such as CelebA $256\times256$ in our experiment. As we show in the table below, the learned VF for mini-batch OT CFM (batch size as large as 128) and uncoupled CFM are close at different time steps, with a relative error of $\sim 10^{-2}$, so their guidance VFs can be approximately the same, which validates our approximation of $\mathcal{P}=1$.

2. Theoretically, the $\mathcal{P}=1$ assumption holds for any dependent couplings when $J$ is slow-varying. A more detailed discussion can be found in our response to reviewer YEWM (Q2).

Q3: Inconsistency of the general guidance expression between Figure 1 and the main text.

Thank you for pointing this out. We will fix it in the revised manuscript.

Q4: Are MC and training-based guidance useful in practice? What do you think of the experiment results of these guidance in high dimensions?

1. We believe $g^{\text{MC}}$ is of practical use in high dimensions.

First, $g_t^{\text{MC}}$ is practical (at least in some tasks) as it shows satisfactory performance in the widely used offline RL task of D4RL, where the sample dimensionality is not very small ($23\times 20$ per sample).

Besides, there are many techniques that can be readily applied to enhance the efficiency of $g_t^{\text{MC}}$. For example, we can adopt importance sampling to reduce the variance of the MC estimation. Specifically, using

$$g_t^{\text{MC-IS}}(x_t) = \mathbb{E}_{x_1,x_0\sim \tilde{p}(z)} \left[\frac{{p}(z)}{\tilde{p}(z)} (\frac{e^{-J(x_1)}}{Z_t} - 1) \frac{p_t(x_t|z)}{p_t(x_t)} v_t (x_t|z) \right],$$ $$Z_t(x_t) = \mathbb{E}_{x_1,x_0\sim \tilde{p}(z)} \left[\frac{{p}(z)}{\tilde{p}(z)} e^{-J(x_1)} \frac{p_t(x_t|z)}{p_t(x_t)} \right],$$

we can select an alternative distribution $\tilde{p}$ such that $\frac{e^{-J(x_1)}}{\tilde{p}(z)}$ has lower variance (i.e., when $e^{-J(x_1)}$ is large, $\tilde{p}(z)$ is also large, and vice versa). This can be achieved by using another guided VF to sample from $\tilde{p}$ such as using $g^{\text{cov-A}}$ and then $\frac{p(z)}{\tilde p(z)}$ can be estimated using, for example, the Hutchinson trace estimator to preserve scalability [1]. It should be noted that this estimation is still unbiased. We conduct experiments to validate the effectiveness of this method, and the results in the table below showcase a huge improvement comparable to the state-of-the-art methods on the image inverse problem, where the vanilla $g_t^{\text{MC}}$ fails.

2. As for the training-based guidance, it has the potential to be widely applicable because it demonstrates the theoretical exactness in the synthetic dataset experiments. Although currently, its performance is restricted, possibly due to high variance induced by the dependency among multiple neural networks in training and inference, it can be potentially addressed by dynamically estimating $Z_t$ rather than training another NN to approximate it, or fine-tuning the guidance VF on the actual original VF.

[1] Lipman, et al., Flow Matching for Generative Modeling, ICLR 2023.

Thank you for acknowledging our contribution in the theoretical framework, the novelty of guidance methods, the soundness of our empirical validation, and the clarity of presentation. We will respond to your questions below:

Q1: How to further address the scenarios where the $\mathcal{P}\approx 1$ approximation does not apply?

First, we would like to emphasize that $\mathcal{P}=1$ is a valid approximation for dependent couplings in realistic datasets such as CelebA $256\times256$ in our experiment. As we show in the table below, the learned VF of mini-batch OT CFM (batch size as large as 128) and uncoupled CFM are close at different flow time steps, with a relative error of $\sim 10^{-2}$, so their guidance VFs can be approximately the same, which validates our approximation of $P=1$.

To produce exact guidance for dependent couplings, one approach is to dynamically adapt the guidance VF as you pointed out. Since $\mathcal{P}\neq 1$ changes high-dimensional integrals, it is costly to directly compute its influence. However, if we make assumptions on the form of $p(x_0|x_1)$ of the original VF, we can derive the corresponding $g^{\text{MC}}$ and approximate guidance. Meanwhile, the approximation error caused by $P=1$ can also be compensated by adapting the source distribution as we discussed in our response to YEWM (Q2). Therefore, we can also parameterize the source distribution and optimize it to recover the exactness of the energy guidance. We will add a discussion on these future directions in the revised manuscript.

Q2: Test more tasks to highlight real-world utility.

Thank you for your suggestions. The experiments in our paper include data modalities from time-series data to images, and are realistic and high-dimensional. Therefore, we believe our empirical evaluation can effectively reveal the utility of different guidance methods on different types of realistic generative modeling tasks. In addition, we provide an additional experiment on a class-label conditioned image generation task to increase the variety of guidance energy functions. With the trained CFM model in our image inverse problem experiment, we use another classifier on the gender to produce the objective function $J$, and consider two cases where either the male or female is set as the target. The table below shows the classification accuracy of the samples generated by different guidance methods:

These results reveal that on image generation tasks, the performance gap between $g^{\text{MC}}$ and $g^{\text{cov-A}}$ is narrowed when the objective function is less complicated.

Thank you for acknowledging the theoretical and empirical soundness of our work, as well as the potential impact in addressing the challenges in the generative modeling field. We answer your questions in the following:

Q1: How can the MC sampling efficiency be improved?

Many existing techniques can be readily applied to enhance the efficiency of $g_t^{\text{MC}}$. For example, we can adopt importance sampling to reduce the variance of the MC estimation. Specifically, using

$$g_t^{\text{MC-IS}}(x_t) = \mathbb{E}_{x_1,x_0\sim \tilde{p}(z)} \left[\frac{{p}(z)}{\tilde{p}(z)} (\frac{e^{-J(x_1)}}{Z_t} - 1) \frac{p_t(x_t|z)}{p_t(x_t)} v_t (x_t|z) \right],$$ $$Z_t(x_t) = \mathbb{E}_{x_1,x_0\sim \tilde{p}(z)} \left[\frac{{p}(z)}{\tilde{p}(z)} e^{-J(x_1)} \frac{p_t(x_t|z)}{p_t(x_t)} \right],$$

we can select an alternative distribution $\tilde{p}$ such that $\frac{e^{-J(x_1)}}{\tilde{p}(z)}$ has lower variance (i.e., when $e^{-J(x_1)}$ is large, $\tilde{p}(z)$ is also large, and vice versa). This can be achieved by using another guided VF to sample from $\tilde{p}$, such as using $g^{\text{cov-A}}$. Then $\frac{p(z)}{\tilde p(z)}$ can be estimated using, for example, the Hutchinson trace estimator to preserve scalability [1]. It should be noted that this estimation is still unbiased. We conduct experiments to validate the effectiveness of this method, and the results in the table below showcase a performance comparable to the state-of-the-art methods on the image inverse problem, where the vanilla $g_t^{\text{MC}}$ fails.

Q2: What are the limitations of our $\mathcal{P}=1$ approximation?

First, we would like to clarify our approximation. Our framework allows us to choose any $\pi'(x_0|x_1)$ (and hence $\mathcal{P}$) as long as the source distribution is consistent: $p(x_0) = \int\pi'(x_0|x_1)\frac{1}{Z}p(x_1)e^{-J(x_1)}dx_1$. In other words, the error induced by setting $\mathcal{P}=1$ can be characterized by the deviation in either the source distribution or the vector field. In the former case, we assume the guidance VF to be exact, i.e., $\pi'(x_0|x_1) = \pi(x_0|x_1)$. Here the error is induced by the fact that we should have sampled from $\int\pi'(x_0|x_1)\frac{1}{Z}p(x_1)e^{-J(x_1)}dx_1$, rather than the original $p(x_0)$. In the latter case, we assume the source distribution to be unchanged, i.e., we need $\pi'(x_0|x_1)=p(x_0)$ to make the source distributions compatible automatically. In this case, the error is caused by approximating $\mathcal{P}=\frac{p(x_0)}{\pi(x_0|x_1)}$ with $1$.

1. In the case of strongly dependent couplings, $\mathcal{P}\approx 1$ still holds as long as $J$ varies slowly. This is demonstrated by the small deviation in the error of the source distribution (assuming $\pi'(x_0|x_1)=\pi(x_0|x_1)$) as we discussed above. If J is always near its average value, the new source distribution $\int\pi(x_0|x_1)\frac{1}{Z}p(x_1)e^{-J(x_1)}dx_1$ is almost $\int\pi(x_0|x_1)p(x_1)dx_1 = p(x_0)$ that is the original source.
2. Nevertheless, when both the coupling is strong and $J$ varies intensively, a more complicated treatment is required for exact guidance. For example, we can try to sample from the new source distribution $\int\pi(x_0|x_1)\frac{1}{Z}p(x_1)e^{-J(x_1)}dx_1$. Although one may argue that this may be equally difficult as sampling from the target distribution $\frac{1}{Z}p(x_1)e^{-J(x_1)}$ exactly, it may be learned more easily as the target distribution is potentially smoothed after being convolved with the "kernel" $\pi(x_0|x_1)$. We will add this discussion to the revised manuscript.

[1] Lipman, et al., Flow Matching for Generative Modeling, ICLR 2023.