Multidimensional Trajectory Optimization for Flow and Diffusion

We revised the trajectory optimization section in the related works to reflect these distinctions more clearly. The revised content is as follows:

" $**Trajectory Optimizations in Flow and Diffusion**$
Various trajectory optimization approaches pre-define optimality without relying on final transportation quality. Approaches such as [5, 6] define straightness as the optimality criterion and optimize trajectories by maintaining the consistency of $(x_0, x_1)$ for training flow and diffusion models, aligning with an optimal transport perspective. Another example is [10], which defines optimality through a fixed sequence of diffusion steps aimed at reducing inference complexity rather than focusing on the quality of final samples. Additionally, [11] introduces neural flow models that implicitly set trajectory optimality within the diffusion process, aiming to generate high-quality samples without explicit trajectory adjustment post-training. There are also approaches that refine trajectories after training, such as [12], where optimality is defined by minimizing the trajectory length in the Wasserstein-2 metric, focusing on a shortest-distance criterion. Despite these diverse perspectives on trajectory optimality, there are two significant differences between these methods and ours. First, we calculate the optimality of the trajectory solely based on the final transportation quality, which is a crucial factor in generative modeling. Second, none of these methods achieve full flexibility of the trajectory on two fronts: multidimensionality and adaptability with respect to different inference trajectories. "

This revision ensures that the differences between our approach and existing methods are clearly articulated and the uniqueness of our contributions is highlighted. Thank you for your insightful comments that helped us clarify these points.

We once again sincerely thank to you for your thoughtful feedback and valuable insights. Your suggestions have greatly helped us to clarify and improve the presentation, scalability, and technical rigor of our work. We have revised the paper thoroughly to address all the comments and have uploaded a new version that reflects these changes.

If you have any further questions or comments, we would be happy to address them without hesitation. Thank you again for your constructive and encouraging feedback.

[1] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sher- jil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets, 2014.

[2] Axel Sauer, Katja Schwarz, and Andreas Geiger. Stylegan-xl: Scaling stylegan to large diverse datasets, 2022.

[3] Jae Hyun Lim and Jong Chul Ye. Geometric gan, 2017.

[4] Dongjun Kim, Chieh-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning probability flow ODE trajectory of diffusion, 2024.

[5] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations, 2018.

[6] 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, 2024.

[7] Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow, 2023.

[8] Bowen Zheng and Tianming Yang. Diffusion models are innate one-step generators, 2024.

[9] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sher- jil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets, 2014.

[10] Raghav Singhal, Mark Goldstein, and Rajesh Ranganath. Where to diffuse, how to diffuse, and how to get back: Automated learning for multivariate diffusions, 2023.

[11] Grigory Bartosh, Dmitry Vetrov, and Christian A. Naesseth. Neural flow diffusion models: Learnable forward process for improved diffusion modelling, 2024.

[12] Michael Samuel Albergo, Nicholas Matthew Boffi, Michael Lindsey, and Eric Vanden-Eijnden. Multimarginal generative modeling with stochastic interpolants, 2024

Regarding experiments on large and high-dimensional datasets like ImageNet, we completely agree with your point. We tried our best to validate our methodology's scalability and effectiveness but couldn't afford to train the EDM-based $H_\theta$ from scratch for ImageNet-64, which requires 32 A100 GPUs and several weeks of training for a fair comparison. We sincerely aim to experiment with our methods on large datasets like ImageNet and video datasets in future work. Thank you again for pointing out this important perspective.

[$**Question 1.**$] The proposed method generally performs worse than methods based on diffusion + distillation in Tables 3 and 4, despite having a 5x higher NFE than these baselines. Do the authors have an explanation for why this is the case?

[$**1. Reasons for the results in FFHQ and AFHQv2**$]

We apologize for not thoroughly discussing the results for FFHQ and AFHQv2. Thank you for pointing out this important aspect.

First, we want to emphasize that comparing the adversarial approach to MTO with diffusion distillation is not entirely fair, as these methodologies are fundamentally different and can even be complementary. However, due to the limited availability of comparable experiments from diffusion distillation and trajectory optimization on FFHQ and AFHQv2, it was necessary to compare our results with these methods.

Second, we aim to show the relationship between data dimensionality, model size for $\gamma_\phi$, and NFE for performance. By using the same model size for $\gamma_\phi$ (specific configurations are written in Appendix C.2) and the same NFE for all datasets (note that baseline EDM uses more than twice the NFE for FFHQ and AFHQv2 compared to CIFAR-10), we found that FFHQ and AFHQv2’s higher dimensionality (64$\times$64 resolution) compared to CIFAR-10 (32$\times$32 resolution) imposes additional burden on $\gamma_\phi$ to find adaptive trajectories, requiring more NFE to achieve comparable FID values.

These observations indicate that achieving significant performance improvements in FFHQ and AFHQv2 similar to CIFAR-10 would require either a larger model for $\gamma_\phi$ or more NFE during training. As noted in the $**Training efficiency and scalability issues**$ section, early convergence observed in FFHQ and AFHQv2 signals that a larger $\gamma_\phi$ model could improve performance further.

Future works could resolve these training efficiency and performance challenges in two ways:

1. Designing a more efficient $\gamma_\phi$ architecture to reduce computational costs (e.g., model size and NFE).
2. Developing more efficient training strategies to enhance performance.

Based on this, we added senteneces to the limitation section as below:

"Additionally, improving the design of $\gamma_\phi$ could enhance efficiency. As shown in Table 3 and Table 4, our method’s FID is higher than distillation methods for larger datasets. This may stem from using the same model size and NFE settings as in smaller datasets like CIFAR-10. Refining $\gamma_\phi$ could reduce model size and NFE while maintaining performance."

Thank you for pointing out these important points.

[$**Question 2.**$] There is some existing work on learning trajectories in diffusion models. e.g. Where to Diffuse... and Neural Flow Diffusion Models. What is the relationship between these methods (and other existing learnable diffusions) and this work?

[$**2. Comparing with existing trajectory optimization works**$]

We sincerely appreciate your effort in identifying relevant related works. We compare MTO and our adversarial approach for solving MTO with other research in terms of four key aspects:

(i) {Does $\phi$ exhibit multidimensionality? (multidimensional coefficient)

(ii) Is $\phi$ trajectory-dependent? (adaptive multidimensional coefficient as defined above)

(iii) Are $\theta$ and $\phi$ jointly learnable?

(iv) Is optimality defined by the final transportation quality?

WTH (Where to diffuse...) [10] defines optimality through a fixed sequence of diffusion steps aimed at reducing inference complexity rather than focusing on the quality of final samples. Additionally, NFDM (Neural Flow Diffusion Models) [11] introduces neural flow models that implicitly set trajectory optimality within the diffusion process, aiming to generate high-quality samples without explicit trajectory adjustment post-training.

While these methods satisfy (iii) and provide diverse perspectives on trajectory optimality, there are two significant differences between these methods and ours:

1. We calculate trajectory optimality solely based on the final transportation quality by end-to-end training (iv), which is crucial in generative modeling.

2. None of these methods achieve full trajectory flexibility regarding multidimensionality (i) and adaptability (ii) for different inference trajectories.

[$**Weakness 2.**$] One of the major benefits of flow and diffusion models is that they are simulation-free during training, making training generally straightforward to carry out. This paper moves in the opposite direction and now requires not only simulation from the model, but also a fairly complex adversarial setup. I would be interested in seeing an evaluation of the training cost of the proposed method, which I'd guess is significantly higher than other flow/diffusion methods.

[$**Weakness 3.**$] The datasets considered are all quite low-dimensional (2d synthetic datasets, CIFAR-10, and some 64x64 image datasets). This makes me wonder if the proposed method is scalable to higher-dimensional datasets.

[$**2 and 3. Illustrating the training cost of simulation training to highlight our efficiency and scalability**$]

After reading your comments, we realized that two points were insufficiently discussed:

1. Why we adopt the combination of simulation-free and simulation-based end-to-end methods.
2. Training cost, efficiency, and scalability issues.

We addressed these topics specifically as follows:

$**Reasons for combining simulation-free and simulation-based end-to-end methods**$

We tried to answer the main question:

“Given a differential equation solver with fixed configurations, which multidimensional trajectory yields optimal performance in terms of final transportation quality for a given starting point of the differential equation?”

This question highlights a trade-off inherent in diffusion models, where simulation-free objectives—while beneficial in reducing training costs—limit adaptability in trajectory optimization concerning output quality. This flexibility is retained in simulation-dynamics [5]. To enable trajectory optimization while maintaining simulation-free objectives for training cost efficiency, prior approaches have relied on pre-defined trajectory properties, such as straightness [6, 7], to minimize numerical error. However, such pre-defined properties for the trajectories diverge from true transportation optimality, as they fail to account for the sole measure of optimality in our perspective: the final quality of transportation.

We reintroduce trajectory adaptability by employing simulation dynamics combined with adversarial training [5], defining trajectory optimality solely based on the final generative output under fixed solver configurations. To effectively leverage the advantages of simulation-based objectives—which provide adaptability (in terms of inference trajectories) and flexibility (in terms of multidimensionality) of trajectories—while mitigating their inefficiencies in training, we use simulation-based objectives only for $\phi$ after pre-training $\theta$ with a simulation-free objective. This approach makes trajectory optimization feasible in terms of efficiency and scalability. Our experiments demonstrate that trajectories optimized through this approach significantly improve the performance of flow and diffusion models.

$**Training efficiency and scalability issues**$

We found that the training efficiency of our methodology is significantly better than other distillation methods, such as GDD [8]. We evaluated the number of images required for training and the training time, comparing these with GDD and GDD-I [8]. Remarkably, our approach for MTO requires far fewer images for training across all datasets compared to GDD-I, a method known for its efficiency in diffusion distillation:

CIFAR-10: 5000kimg (GDD-I) vs. 1382kimg (Ours)

FFHQ: 10000kimg (GDD-I) vs. 106kimg (Ours)

AFHQv2: 10000kimg (GDD-I) vs. 955kimg (Ours)

The training times for our method are 10, 2, and 6 hours for CIFAR-10, FFHQ, and AFHQv2, respectively, which are also comparatively low. Additionally, FID significantly decreases in the early stages of training, as illustrated in Figure 6 (we modified the x-axis from iteration number to kimg for improved readability). The primary computational cost arises from VRAM requirements; however, training remains practically feasible. Good performance is achieved with just 5 NFE, which is relatively low. All experiments for adversarial training were conducted on GPUs with 48GB of VRAM—less than the 80GB VRAM GPUs often used in related works.

These results underscore the effectiveness of simulation-based end-to-end optimality while highlighting the strength of combining simulation-free and simulation-based methodologies, leveraging the advantages of each while mitigating their limitations. Considering these aspects, we estimate that our adversarial approach for MTO is scalable to larger datasets while maintaining efficiency.

First of all, we really appreciate your insights and valuable feedback. We are especially grateful for recognizing the following strengths of our work, which we summarize below:

[$**STRENGTHS**$]

We thank Reviewer dLdR for encouraging feedback. Reviewer dLdR highlighted the novelty of having learnable trajectories, emphasizing how this idea introduces an interesting dimension to trajectory optimization. We are pleased that Reviewer dLdR recognized the value of our approach in improving generative model performance through learnable trajectories. Additionally, Reviewer dLdR appreciated our experimental results, which demonstrate that learnable trajectories can indeed lead to better performance. This acknowledgment of our experiments’ effectiveness reinforces the practical utility of our proposed methodology. We are glad that Reviewer dLdR found our contributions both innovative and impactful.

[$**WEAKNESSES**$]

[$**Weakness 1.**$] The paper is a little hard to read, and could generally benefit from additional exposition and improved clarity

[$**1. Paper clarity issues**$]

We totally agree with your comments. Sorry for the poor explanation. We revised all parts you have mentioned as below.

First, to clarify the boundary condition for the multidimensional coefficient, we added the sentence:

"Boundary conditions written above ensure that $x(t)$ becomes $x_0$ and $x_T$ for $t=0$ and $t=T$, which is a requirement for transportation. Values $k$ and $T$ for boundary conditions vary based on the task."

Second, for rigorous definition and better readability, we introduced a new definition called the adaptive multidimensional coefficient, which captures the concept of adaptability for various inference trajectories. The definition is written below:

$**Definition 2: Adaptive Multidimensional Coefficient**$
For $t \in [0, T]$ and inference trajectory $x_{\theta, \phi}(t)$, the adaptive multidimensional coefficient $\gamma_\phi(t, x_{\theta, \phi}(t)):[0, T] \times \mathbb{R}^d \rightarrow \mathbb{R}^{2 \times d}$ is parameterized by $\phi$. Boundary conditions follow Definition 1 (Multidimensional Coefficient), with $\gamma_\phi \in C^1([0, T], \mathbb{R}^{2 \times d})$.

This definition describes the trajectory-conditioned multidimensional coefficient. To reduce computational costs in calculating $\gamma_\phi$, we use only $t = T$ for $x_{\theta, \phi}(t)$ in $\gamma_\phi(t, x_{\theta, \phi}(t))$ rather than the inference trajectory at multiple time points. This approach allows us to compute $\gamma_\phi$ across the entire inference time schedule $\tau = \{ t_0, \ldots, t_N \}$ with a single function evaluation before initiating transportation. By using the adaptive multidimensional coefficient, we can address multidimensional trajectory optimization more rigorously.

Third, we unified all the notations for noise distribution and samples to $x_T \sim \rho_T$.

Lastly, we revised the adversarial training section with more details based on your comments:

"
Following GAN [1], we use an adversarial loss to minimize Equation 6. Specifically, we employ the StyleGAN-XL [2] discriminator for $D_\psi$ with hinge loss [3] for $\theta$, $\phi$, and $\psi$ (discriminator parameters), as used in [4]. In our method, $\theta$ and $\phi$ in $G_{\theta, \phi}$ aim to deceive $D_\psi$. The loss function for $\theta$ is:

$L_{\theta} = - E_{x_1 \sim \rho_1}[D_\psi \left( H_\theta(t, x(t), \gamma_\phi(t, z)) \right) ], \quad x(t) = x_0 + \gamma_{1, \phi}(t,z) \odot x_1, \quad z \sim \rho_T,$

where $L_\theta$ indicates that the gradient is only calculated with respect to $\theta$. Since $\theta$ only needs to handle elements in $\{\gamma_\phi(t, x_T) \mid t \in [0, T], x_T \sim \rho_T\}$ and not other elements in $\Gamma_h$, $\theta$ is trained exclusively on $\gamma_\phi$, helping reduce the load on $H_\theta$. The loss functions for $\phi$ and $\psi$ are:

$L_\phi = - E_{x_1 \sim \rho_1}[D_\psi \left( G(\tau, x_1, v_{\theta, \phi}) \right) ],$

$L_\psi = E_{x_0 \sim \rho_0}[\max(0, 1 - D_\psi(x_0))] + E_{x_1 \sim \rho_1}[\max(0, 1 + D_\psi(G(\tau, x_1, v_{\theta, \phi})))].$

Here, $L_\phi$ and $L_\psi$ indicate that gradients are calculated with respect to $\phi$ and $\psi$, respectively, as shown in Figure 3. With these loss functions, $\phi$ is optimized to find better trajectories, while $\theta$ adapts to $\gamma_\phi$, which is sparser than $\Gamma_h$.
"

With these clarifications, we hope that the paper's readability is significantly improved. Thank you for pointing out these specific issues.

4. $**Revision for our comment**$

Finally, we revisited our response to your initial review question: "In line 125, the authors use $\rho_1$ as the noise distribution for EDM, but in line 130, $\rho_T$ is used instead." Using $x_1 \sim \rho_1 = \mathcal{N}(0, I)$ and $\rho_T = \mathcal{N}(0, T^2I)$ is correct, as the former applies to the pre-training stage and the latter to the inference stage. We sample $x_0, x_1 \sim \rho_0, \rho_1$ to train the flow or diffusion model and sample $x_T \sim \rho_T$ for transportation. We apologize for the repeated revisions.

Overall, we thank you again for your beneficial feedback, which has helped us clarify our writing.

5. $**Additional experiments**$

We also conducted additional experiments to validate the empirical benefits of our methodology, which are now organized in Appendix E to further support the contributions of our work. We want to highlight three key points:

1. Pre-training with the multidimensional coefficient not only preserves but also enhances the performance of flow and diffusion models.
2. All MTO experiments in Appendix E were conducted by training only $\phi$, while keeping $\theta$ frozen.
3. Even with a small 4-layer discriminator network and Vanilla GAN Loss, we achieved comparable FID performance on SI.

Thank you for your attention to these updates.

First of all, we truly appreciate your concise feedback for clarifying and improving the writing of our work.

1. $**Explain our methodology precisely**$

We thank Reviewer dLdR for pointing out unclear and confusing parts in our paper that might have caused misunderstandings. Here, we aim to provide a precise explanation of our methodology.

The training is divided into two processes, with adversarial training applied only during the second process:

• $**1st process (Pre-training stage, non-adversarial)**$: We train $\theta$ using various randomly sampled coefficients from the hypothesis space $\Gamma_h$, which we specifically designed to exclude crude coefficients. This process employs the same loss function and training methodology used in standard diffusion models. After this pre-training stage, $H_\theta$ can handle denoising tasks with $\gamma \sim \Gamma_h$, preparing it for MTO.

• $**2nd process (MTO stage, adversarial)**$: During this stage, we train $G_{\theta, \phi}$ as a generator, composed of $H_\theta$ and $\gamma_\phi$, using a discriminator $D_\psi$. The adversarial training is conducted between {$\theta, \phi$} <-> {$\psi$}.

Based on this understanding, we clarified the points you mentioned as follows:

2. $**Clarifying the methodology, especially for adversarial training**$

We appreciate your feedback on the lack of clarity regarding the adversarial training section. To address this:

• We revised sentences like "Our approach employs simulation dynamics and adversarial training to optimize these inference trajectories" to "Our approach pre-trains flow and diffusion models with various coefficients sampled from a hypothesis space and subsequently optimizes inference trajectories through adversarial training of a generator comprising the flow or diffusion model and the parameterized coefficient" (lines 17–20, 212-214, 516–518).

• We refined the explanation: "We reintroduce trajectory adaptability by employing simulation dynamics combined with adversarial training (Goodfellow et al., 2014), defining trajectory optimality based solely on the final generative output..." to "We reintroduce trajectory adaptability by employing simulation dynamics combined with adversarial training, defining trajectory optimality based solely on the final generative output under fixed solver configurations. Specifically, first, we pre-train a diffusion model $H_\theta$ with randomly sampled multidimensional coefficients $\gamma$ from a well-designed hypothesis space. This pre-training enables the flow or diffusion model to handle various coefficients, preparing it for the trajectory optimization stage. Next, we introduce the parameterized adaptive multidimensional coefficient $\gamma_\phi$ to compose a flow or diffusion-based generator $G_{\theta, \phi}$, which produces $x_{\text{est}, \theta, \phi}$ through simulation dynamics. A discriminator $D_\psi$ evaluates the generated samples $x_{\text{est}, \theta, \phi}$ to optimize both $\theta$ and $\phi$, a process we term multidimensional trajectory optimization" (lines 54–62).

• To avoid confusion about the two-stage training process, we explicitly added the pre-training loss term in Section 4.3 (lines 284–289).

• We included a min-max equation from GANs (Equation (12)) to clarify the connection between MTO and the adversarial approach to MTO.

3. $**Move technical details to Appendix**$

• We agree that technical, implementation-level details should generally be moved to the appendix to focus on the main ideas in the paper. However, we believe Equations 9 and 10 should remain in the main text, as they are crucial for explaining our hypothesis space design. This decision is particularly relevant in addressing Reviewer xR7Z’s question: "In the pretraining stage, how exactly are randomly sampled trajectories obtained? Do you still use the proposed parameterized family of coefficients in this stage? If so, which parameters are set to random, and what distributions do they follow?" These details are critical for understanding the performance and stability benefits provided by excluding crude coefficients in our hypothesis space design. That said, we agreed to minimize technical explanations in the main text. The choice of $b_m(t)$ is now briefly explained in line 256, with detailed discussions moved to Appendix A. Other technical sections, such as "lines 179–200 (finite set of samples)" and the explanation of the Euler solver, have been condensed and moved to Appendix A.

where $L_\theta$ indicates that the gradient is only calculated with respect to $\theta$. Since $\theta$ only needs to handle elements in \{\gamma_\phi(t, x_T)} and not other elements in $\Gamma_h$, $\theta$ is trained exclusively on $\gamma_\phi$, helping reduce the load on $H_\theta$. The loss functions for $\phi$ and $\psi$ are:

$L_\phi = - E_{x_1 \sim \rho_1}[D_\psi \left( G(\tau, x_1, v_{\theta, \phi}) \right) ],$

$L_\psi = E_{x_0 \sim \rho_0}[\max(0, 1 - D_\psi(x_0))] + E_{x_1 \sim \rho_1}[\max(0, 1 + D_\psi(G(\tau, x_1, v_{\theta, \phi})))].$

Here, $L_\phi$ and $L_\psi$ indicate that gradients are calculated with respect to $\phi$ and $\psi$, respectively, as shown in Figure 3. With these loss functions, $\phi$ is optimized to find better trajectories, while $\theta$ adapts to $\gamma_\phi$, which is sparser than $\Gamma_h$. "

$[**Weakness 6.**]$ In Section 5.2, the authors mention that the NFE of the proposed method is 6 (+1 for $\gamma_\phi$). What exactly does this mean? Why “+1 for $\gamma_\phi$” rather than using learned $\gamma_\phi$ through the whole sampling?

$**[6. Ambiguity of the 5 (+1) notation]**$

There seems to be some misunderstanding in this part. We use learned $\gamma_\phi$ through the whole sampling. As the calculation of $\gamma_\phi$ involves function evaluation, we include this in the NFE count. To reduce ambiguity, we changed the notation from "5 (+1)" to "5 (+)" and added the clarification: "(+ for the calculation of $\gamma_\phi$, given that the network for $\gamma_\phi$ is smaller than the network for $H_\theta$)."

[$**Weakness 7.**$] In Table 2-4, when comparing with other distillation methods, it would be more fair to run other methods under the same NFE. For example, how do CD and CTM perform with 5 or 6 NFE?

$**[7. FID of CTM with 5 or 6 NFEs]**$

We totally agree with your point. For a fair comparison with other distillation methods in terms of NFE, we select CTM [4] as a representative due to its popularity, high performance, and the use of the same model architecture based on EDM and adversarial training. We then calculate the FID using 5 and 6 NFE.

CTM's FID for CIFAR-10 unconditional: 1.86 (NFE=5), 1.93 (NFE=6)

CTM's FID for CIFAR-10 conditional: 1.98 (NFE=5), 2.04 (NFE=6)

Ours achieves 1.69 (CIFAR-10 uncond) and 1.37 (SOTA on CIFAR-10 cond) with 5 NFE.

As shown above, increasing the NFE of CTM does not significantly decrease the FID, and the FID even increases. This indicates that the adversarial approach for MTO provides additional performance gains that cannot be achieved by distillation methods alone, even with increased computational cost.

[$**Weakness 8.**$] It would be great to add some comments in the caption of Figure 2 to help better explain the difference and illustrate the idea.

$**[8. Adding more explanation for the figure]**$

We completely agree with your suggestion. The caption has been updated to: "Crude coefficients: (a) Oscillatory behavior in $t$ due to high-frequency components; (b) High adjacent pixel differences in $d$. Refined coefficients: (c) Constrained multidimensionality for larger $t$ in pre-training; (d) Unconstrained multidimensionality for adversarial training." This should improve comprehension.

Once again, we sincerely appreciate your detailed and specific feedback. We have uploaded a revised version of the PDF reflecting your comments as thoroughly as possible. If you have any further questions, we are happy to address them without hesitation.

Thank you again.

[1] Bowen Zheng and Tianming Yang. Diffusion models are innate one-step generators, 2024.

[2] Axel Sauer, Katja Schwarz, and Andreas Geiger. Stylegan-xl: Scaling stylegan to large diverse datasets, 2022.

[3] Jae Hyun Lim and Jong Chul Ye. Geometric gan, 2017.

[4] Dongjun Kim, Chieh-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning probability flow ODE trajectory of diffusion, 2024.

First of all, we really appreciate your insights and valuable feedback. We are especially grateful for recognizing the following strengths of our work, which we summarize below:

[$**STRENGTHS**$]

We thank Reviewer xR7Z for insightful and encouraging comments! Reviewer xR7Z recognized the novelty of learning multidimensional diffusion/flow coefficients and appreciated the design of our parameterized family of coefficients. Reviewer xR7Z noted how our approach not only ensures boundary conditions but also enables efficient learning, as reflected in our empirical results. We are excited that Reviewer xR7Z sees potential for future follow-up work that explores the design space further. Additionally, Reviewer xR7Z found our experimental results promising. Reviewer xR7Z highlighted that our method consistently achieves lower Wasserstein distances in synthetic cases and performs on par with or slightly better than baseline methods on challenging real-world image generation tasks. We are also pleased that Reviewer xR7Z found our paper to be clearly written, despite some technical details being sparse due to space limitations.

Below, we address the weaknesses mentioned in your comments in detail:

[$**WEAKNESSES**$]

[$**Weakness 1.**$] For a better understanding of the design of the adversarial loss, I’d like to hear authors’ feedback on the following questions. Moreover, it would be great if the authors could improve the writing in Section 4.4.

$**[1-1. Ambiguity of the loss function notation]**$

You are absolutely right, and we acknowledge the vagueness in the notation. We have revised this and added further explanation, such as: "where $L_\theta$ indicates that the gradient is only calculated with respect to $\theta$," and "where $L_\phi$ and $L_\psi$ indicate that gradients are calculated with respect to $\phi$ and $\psi$, respectively."

$**[1-2. Clarifying the adversarial training section]**$

There seems to be some misunderstanding in this part. First, we apologize for the poor explanation. In our method, not only $\phi$ but also $\theta$ in $G_{\theta, \phi}$ aims to deceive $D_\psi$. $\theta$ and $\phi$ are on the same side to fool the discriminator. To improve clarity, we have added a sentence "In our method, $\theta$ and $\phi$ in $G_{\theta, \phi}$ aim to deceive $D_\psi$." to the paper explaining this point.

[$**Weakness 2.**$] While optimizing the coefficients $\gamma$ via the generator $G$, it seems that backpropagation through the ODE solver is inevitable. This is also shown in Fig. 3. Given the high computational cost, could you comment on the scalability of your method?

$**[2. Assessing the efficiency and scalability of our method]**$

Thanks to your insightful comments, we evaluated the number of images required for training and the training time, comparing these with GDD [1] and GDD-I [1] . Remarkably, our approach for MTO requires far fewer images for training across all datasets compared to GDD-I, a method known for its efficiency in diffusion distillation:

CIFAR-10: 5000kimg (GDD-I) vs. 1382kimg (Ours)

FFHQ: 10000kimg (GDD-I) vs. 106kimg (Ours)

AFHQv2: 10000kimg (GDD-I) vs. 955kimg (Ours)

The training times for our method are 10, 2, and 6 hours for CIFAR-10, FFHQ, and AFHQv2, respectively, which are also comparatively low. Additionally, FID significantly decreases in the early stages of training, as illustrated in Figure 6 (We modify the x-axis from iteration number to kimg for improved readability.) The primary computational cost arises from VRAM requirements; however, training remains practically feasible. Good performance is achieved with just 5 NFE, which is relatively low. All experiments for adversarial training were conducted on GPUs with 48GB of VRAM-less than the 80GB VRAM GPUs often used in related works. These results underscore the effectiveness of simulation-based end-to-end optimality while highlighting the strength of combining simulation-free and simulation-based methodologies, leveraging the advantages of each while mitigating their limitations. Considering these aspects, we estimate that our adversarial approach for MTO is scalable to larger datasets while maintaining efficiency. We have added a new paragraph like above in the experiments section to address these scalability issues. Thank you again for your comment.

[$**Weakness 3.**$] In the pretraining stage, how exactly are randomly sampled trajectories obtained? Do you still use the proposed parameterized family of coefficients $\gamma$ in this stage? If so, which parameters are set to random, and what distributions do they follow? How do you ensure the pre-training does not degrade performance much and effectively explores the parameter space?

$**[3-1. Clarifying the design choice of the coefficient's hypothesis space section]**$

Yes, we sample from the hypothesis space. We have revised numerous notations in Section 4.3 (Design Choice of the Coefficient's Hypothesis Space) to reduce ambiguity and enhance clarity. Specifically, the hypothesis space $\Gamma_h \subseteq \Gamma$ is defined as the set of $\gamma_\phi: [0, T] \times R^d \to R^{2 \times d}$, where:

$\gamma_\phi = [\gamma_{0, \phi}, \gamma_{1, \phi}], \quad \phi \in P,$

$\gamma_{0, \phi}(t, x_T) = T \frac{f_\phi(t, x_T)}{f_\phi(t, x_T) + g_\phi(t, x_T)}, \quad \gamma_{1, \phi}(t, x_T) = T \frac{g_\phi(t, x_T)}{f_\phi(t, x_T) + g_\phi(t, x_T)}.$

$P$ represents the parameter space from which $\phi$ is drawn, and it determines the specific forms of $\gamma_{0, \phi}$ and $\gamma_{1, \phi}$ within $\Gamma_h$. The definitions of $f_\phi$ and $g_\phi$ are as follows:

$f_\phi(t, x_T) = 1 - \frac{t}{T} + \left( \sum^M_{m=1} w^f_{m, \phi}(x_T) b_m(t) \right)^2, \quad g_\phi(t, x_T) = \frac{t}{T} + \left( \sum^M_{m=1} w^g_{m, \phi}(x_T) b_m(t) \right)^2,$

where $w_\phi(x_T) = s \ \text{LPF} \circ \tanh \left( U_\phi(x_T) \right), \ s \in \mathbb{R}.$

Random sampling from $\Gamma_h$ is performed as $w(u) = s \ \text{LPF} \circ u, \ u \sim \mathcal{N}(-1, 1) \in \mathbb{R}^{2 \times M \times d}.$

$**[3-2. Ensuring the design choice of hypothesis space exclude crude coefficients]**$

Thank you for pointing out this important aspect. We assume that trajectories including high frequencies in both $t$ and $d$ can be crude for transportation. To address this, we select sinusoidal functions with low frequencies and apply a Low-Pass Filter (LPF) across different dimensions $d$ to exclude high-frequency components.

Specifically, during the pre-training stages, we tested various hyperparameters ($s$, $M$, and LPF configurations) and observed that performance slightly improved when multidimensionality near $t = T$ was reduced. For this purpose, we set the convolution group size for LPF to 1, resulting in an LPF output shape of $[B, 1, \text{resolution}, \text{resolution}]$. This constrains $\gamma$ to exhibit multidimensionality for small $t$ and reduced multidimensionality for large $t$. Details regarding this configuration are specifically illustrated in Appendix A due to space constraints. The empirical results for this are presented in Table 5. As shown, EDM$-\gamma$'s FID (2.08, 1.81) is nearly identical to EDM$-\alpha$'s FID (1.97, 1.79), indicating that our design choice of $\Gamma_h$ is empirically efficient.

We acknowledge that this is not an optimal choice, and there remains significant room for improvement in the design, given that this is the beginning of the exploration for this topic.

[$**Weakness 4.**$] The authors do not explain the design choice of the discriminator $D_\psi$ in Section 4.4 when it was first introduced.

[$**Weakness 5.**$] Can you discuss the motivation for using hinge loss in the adversarial training in more detail? The authors just follow the previous work in GAN training without giving motivation.

$**[4 and 5. Clarifying the discriminator and loss function for adversarial training]**$

We employ the StyleGAN-XL [2] discriminator for $D_\psi$ with hinge loss [3], as used in CTM [4]. This choice is motivated by the near state-of-the-art performance of GAN-based methods utilizing the StyleGAN-XL architecture and loss function. Moreover, prior works on diffusion distillation, such as CTM, employed similar adversarial training setups, which informed our design. Reflecting all your feedback and comments, we revised the adversarial section with more details as shown below:

" Following GAN [5], we use an adversarial loss to minimize Equation 6. Specifically, we employ the StyleGAN-XL [6] discriminator for $D_\psi$ with hinge loss [7] for $\theta$, $\phi$, and $\psi$ (discriminator parameters), as used in [4]. In our method, $\theta$ and $\phi$ in $G_{\theta, \phi}$ aim to deceive $D_\psi$. The loss function for $\theta$ is:

$L_{\theta} = - E_{x_1 \sim \rho_1}[D_\psi \left( H_\theta(t, x(t), \gamma_\phi(t, z)) \right) ], \quad x(t) = x_0 + \gamma_{1, \phi}(t,z) \odot x_1, \quad z \sim \rho_T,$

First of all, thank you for your core and concise concerns, which are highly valuable for addressing the limitations of our paper and guiding follow-up research. We agree that the limitations you pointed out are important, and we have revised our limitations section to address these points.

1. $**One needs to train the model from scratch rather than leveraging existing models.**$

Yes, the coefficient labeling (as described in Section 4.3) makes the model structure different from existing well pre-trained models like EDM. This requirement can indeed make applying MTO cumbersome. Future work could mitigate this issue by replacing a few layers from a well pre-trained model and using it as initialization for pre-training.

2. $**The learned coefficients seem to be quite restrictive (if we learn the coefficients for the 5-step sampler, we cannot adapt them to a sampler ...**$

We agree with your point. Our current approach trades off flexibility for improved performance by restricting the inference sampling configurations. Future work could address this limitation by conditioning $\gamma_\phi$ on diverse sampling configurations, enabling better adaptability and efficiency.

We have revised our limitation section as below:

"First, since we use coefficient labeling (Section~4.3) for the diffusion model, the model structure differs from existing pre-trained flow and diffusion models. As a result, training models from scratch is required, which can be cumbersome. This issue could be mitigated in future works by replacing a few layers from well pre-trained model and using it as initialization for pre-training. Second, $\gamma_\phi$ is tied to the specific sampling configuration used during MTO, limiting its flexibility for inference under alternative configurations. Future work could address this by conditioning $\gamma_\phi$ on diverse sampling configurations, enabling better adaptability and efficiency. ...."

3. $**Some typos still exist in the newly revised part of your draft. For example, in Eq. (10),$x_1$on the right-hand side ...**$

Thank you for pointing out the typos we missed. We have revised Eq. (10) accordingly. Regarding $t \sim \mathcal{N}(-1.2, 1.2)$ in lines 3 and 9, this is actually correct given that $t \in [0, T]$ ($T = 80$ for EDM as in Equation 3), as stated in Definition 1 and 2. We will continue working diligently to ensure all typos are corrected in the final version.

We have revised our PDF and uploaded it to reflect your feedback. Thank you once again for your constructive feedback and detailed observations.

First of all, we really appreciate your insights and valuable feedback. We are especially grateful for recognizing the following strengths of our work, which we summarize below:

[$**STRENGTHS**$]

We thank Reviewer EetL for their positive and insightful comments! Reviewer EetL noted the extensive experiments conducted in our work, highlighting the use of state-of-the-art models. We are encouraged that Reviewer EetL appreciated the comprehensiveness of our experimental validation, as it underscores the robustness of our methodology. Moreover, Reviewer EetL mentioned that our method is easy to understand and that the writing is clear. We are delighted that our efforts to present the ideas and technical details in an accessible manner were recognized. This feedback inspires us to continue ensuring clarity and accessibility in our work.

[$**WEAKNESSES**$]

[$**Weakness 1.**$] Flow Matching Results: If possible, please provide results on the image dataset based on the flow-matching model.

[$**1. Adversarial approach to MTO for Flow Matching**$]

It would be valuable to show empirical experiments with the FM (Flow Matching) [1] framework given its popularity. However, FM has limitations for performing MTO as it directly estimates the entire vector field without disentangling the coefficient and target distributions like EDM. In equation form, flow matching's objective is:

$v(t, x(t)) \approx \hat{v}_\theta(t, x(t))$,

while diffusion predicts the vector field by estimating noise and target distribution's samples $x_{0, \theta}$ and $x_{1, \theta}$ as follows:

$v(t, x(t)) \approx \alpha_0^{'}(t)x_{0, \theta} + \alpha_1^{'}(t)x_{1, \theta}$.

This difference makes it challenging to solve MTO in FM, as we must couple the vector field with $\theta$ and $\phi$, which is not feasible in FM. However, we can use SI (Stochastic Interpolant) [2] for MTO, as shown in our 2-dimensional experiments, given that it estimates $x_{0, \theta}$ and $x_{1, \theta}$ rather than directly modeling the vector field. SI is known for generalizing the flow and diffusion methodologies, which is why we employed SI in the 2-dimensional experiment to demonstrate the performance of a generalized FM methodology.

[$**Weakness 2.**$] Image Dataset Experiment with $N=5$: In the image dataset experiment, $N=5$ was selected, seemingly for efficiency in image generation (simulating Equation (13)). Is this correct? Can a different sampler or time scheduler be used to draw samples after training?

[$**Weakness 5.**$] Optimization of the $\gamma$ Model: The $\gamma$ model is optimized only at specific time schedules (used in Equation (13)) during the second-stage training. Does this imply that the inference schedule used during training must match the schedule used during sampling?

[$**2 and 5. Answering questions for inference by using learned coefficient only for fixed sampling configurations**$]

Thank you for pointing out these important properties of our method.

First, you are correct. The reason for using NFE=5 in the image generation experiment was to prioritize training efficiency. We can employ various configurations by trading off computational cost and image quality.

Second, since our training scheme depends on the sampling configuration, the inference schedule used during training must match the schedule used during sampling. Although we can utilize a different inference schedule, as $\gamma_\phi$ provides continuous coefficients, the outputs are not optimized for those schedules since $\gamma_\phi$ was not exposed to them during training.

The problem we aim to solve is:

“Given a differential equation solver $**with fixed configurations**$, which multidimensional trajectory yields optimal performance in terms of final transportation quality for a given starting point of the differential equation?”

The additional performance improvements gained are derived from simulation-based objectives, which require fixed sampling strategies. We acknowledge that this can be a constraint for flexible inference and have revised the limitations section to address this topic as below:

"solving MTO under fixed inference schedules may limit flexibility. Future research could explore general sampling configurations to enhance adaptability and efficiency."

Thank you for your constructive feedback again.

[$**Weakness 3.**$] Ablation Study Findings: The ablation study shows that EDM-$\gamma$+ Adv achieves comparable performance. This adversarial objective was previously proposed and appears to be the critical component, rather than the proposed trajectory optimization coefficients.

[$**3. Question about the practical benefits of optimizing$\gamma$by observing ablation study**$]

We agree with your point. It is true that adversarial training of $\theta$ is crucial for performance improvements. However, we emphasize that optimizing only $\gamma_\phi$ with a simulation-based end-to-end objective alone is effective and novel, as shown in the 2-dimensional experiment. These experiments validate the existence of performance gains achievable through optimizing only $\phi$. To clarify this, we added an explanation in 2-dimensional experiment section.

Additionally, we highlight our training methodology by comparing it with CTM [3], which uses adversarial training for $\theta$ but not for $\phi$. For this comparison, we calculated FID using 5 and 6 NFE:

CTM's FID for CIFAR-10 unconditional: 1.86 (NFE=5), 1.93 (NFE=6).
CTM's FID for CIFAR-10 conditional: 1.98 (NFE=5), 2.04 (NFE=6).
Ours achieves 1.69 (CIFAR-10 uncond) and 1.37 (SOTA on CIFAR-10 cond) with 5 NFE.

As shown above, increasing the NFE of CTM does not significantly reduce FID, and it even increases slightly. This demonstrates that the adversarial approach for MTO provides performance gains that cannot be achieved by distillation methods alone, even with increased computational cost. These performance improvements for near-SOTA and SOTA results stem from our unique training methodology focusing on jointly training $\theta$ and $\phi$.

We also highlight that training $\theta$ and $\phi$ using simulation-based end-to-end optimality is typically inefficient, but our methodology addresses this issue. This point is elaborated further in [$**6. Training efficiency of simulation-dynamics**$].

By these additional experiments and observations, we emphasize two key aspects:

1. The effectiveness of flexibility and adaptability achieved through the adaptive multidimensional coefficient $\gamma_\phi$.
2. The practical benefits of combining simulation-based end-to-end optimality with simulation-free methodologies, leveraging the advantages of each while mitigating their limitations.

These comparisons and insights underscore the practical benefits of adversarial approach to MTO. Thank you again for pointing out this critical aspect.

[$**Weakness 4.**$] State-dependent Multidimensional Coefficients: Can the multidimensional coefficients also depend on the state?

[$**4. State-dependent adaptive multidimensional coefficient**$]

Yes, you are absolutely correct. After reading your comments, we realized it would be beneficial to define state-dependent multidimensional coefficients separately and clarify our choice. For this purpose, we introduced a new definition called the adaptive multidimensional coefficient, which captures the concept of adaptability for various inference trajectories. The definition is as follows:

$**Definition 2: Adaptive Multidimensional Coefficient**$
For $t \in [0, T]$ and inference trajectory $x_{\theta, \phi}(t)$, the adaptive multidimensional coefficient $\gamma_\phi(t, x_{\theta, \phi}(t)):[0, T] \times \mathbb{R}^d \rightarrow \mathbb{R}^{2 \times d}$ is parameterized by $\phi$. Boundary conditions follow Definition 1 (Multidimensional Coefficient), with $\gamma_\phi \in C^1([0, T], \mathbb{R}^{2 \times d})$.

This definition describes the trajectory-conditioned multidimensional coefficient. To reduce computational costs in calculating $\gamma_\phi$, we use only $t = T$ for $x_{\theta, \phi}(t)$ in $\gamma_\phi(t, x_{\theta, \phi}(t))$ rather than the inference trajectory at multiple time points. This approach allows us to compute $\gamma_\phi$ across the entire inference time schedule $\tau = \{ t_0, \ldots, t_N \}$ with a single function evaluation before initiating transportation. By using the adaptive multidimensional coefficient, we can address multidimensional trajectory optimization more rigorously.

Thank you for your insightful suggestion, which has helped us better articulate this important concept.

[$**Weakness 6.**$] Training Efficiency in Simulation Dynamics: Please include a discussion on the training efficiency.

[$**6. Training efficiency of simulation-dynamics**$]

We found that the training efficiency of our methodology is significantly better than other distillation methods, such as GDD [4]. We evaluated the number of images required for training and the training time, comparing these with GDD and GDD-I [4]. Remarkably, our approach for MTO requires far fewer images for training across all datasets compared to GDD-I, a method known for its efficiency in diffusion distillation:

CIFAR-10: 5000kimg (GDD-I) vs. 1382kimg (Ours)
FFHQ: 10000kimg (GDD-I) vs. 106kimg (Ours)
AFHQv2: 10000kimg (GDD-I) vs. 955kimg (Ours)

The training times for our method are 10, 2, and 6 hours for CIFAR-10, FFHQ, and AFHQv2, respectively, which are also comparatively low. Additionally, FID significantly decreases in the early stages of training, as illustrated in Figure 6 (we modified the x-axis from iteration number to kimg for improved readability). The primary computational cost arises from VRAM requirements; however, training remains practically feasible. Good performance is achieved with just 5 NFE, which is relatively low. All experiments for adversarial training were conducted on GPUs with 48GB of VRAM—less than the 80GB VRAM GPUs often used in related works.

These results underscore the effectiveness of simulation-based end-to-end optimality while highlighting the strength of combining simulation-free and simulation-based methodologies, leveraging the advantages of each while mitigating their limitations. Considering these aspects, we estimate that our adversarial approach for MTO is scalable to larger datasets while maintaining efficiency.

[$**Weakness 7.**$] Writing adjustments: Verify the correctness of $\gamma_0(T)$ on the right-hand side of Equation (8).

[$**7. Writing adjustments for boundary conditions**$]

Yes, $\gamma_0(T)$ on the right-hand side of Equation (8) is correct. For cases like $T=1$, $k$ should be $k=0$ to satisfy the boundary conditions. Boundary conditions like this ensure that $x(t)$ becomes $x_0$ and $x_T$ for $t=0$ and $t=T$, which is a requirement for transportation. The values of $k$ and $T$ for boundary conditions vary based on the task. To avoid misunderstandings, we have added an additional explanation of the boundary conditions for better comprehension. Thank you for your observation.

We sincerely thank you once again for your thoughtful feedback and valuable insights. Your suggestions have been instrumental in helping us clarify and enhance the presentation, scalability, and technical rigor of our work. We have carefully revised the paper to address all the comments and uploaded a new version that incorporates these changes.

If you have any additional questions or further comments, we would be more than happy to address them. Thank you once again for your constructive and encouraging feedback.

[1] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, Matt Le, Flow Matching for Generative Modeling, 2023.

[2] Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions, 2023

[3] Dongjun Kim, Chieh-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion, 2024.

[4] Bowen Zheng and Tianming Yang. Diffusion models are innate one-step generators, 2024.

$**Q. As reviewer ZX4x points out, the biggest technical concern with this paper is that the major performance gains are rooted ...**$

We acknowledge that the loss function itself is not highly novel. However, $**our approach is non-trivial:**$ using distinct objectives for $\theta$ and $\phi$, achieving remarkable efficiency and performance. Also, we would like to emphasize the contribution of the additional $\phi$ parameterization and its optimization, which is supported by extra experimental results in the supplementary materials.

We want to highlight that all the experiments in our supplementary materials for MTO are $**obtained by training only$\phi$, while keeping$\theta$frozen**$. As shown in the supplementary materials, particularly in the SI results on CIFAR-10 and ImageNet-32, we achieved comparably good FID values using only a few steps. These additional experimental results demonstrate that SI and MTO are a good combination, potentially sparking new interest in SI methodology. While these FID values are not close to those of current SOTA methodologies, this is due to the limitations of FM and SI and optimizing only $\phi$ rather than limitations of MTO itself, given that we demonstrated SOTA performance (CIFAR-10 conditional) using EDM in the main paper.

Overall, the results in the supplementary materials show that we achieved on-par performance by optimizing only $\phi$ with an adversarial objective across various methodologies. This underscores the novelty and effectiveness of our approach in leveraging $\phi$ optimization for trajectory adaptability and improved generative performance.

Thank you for your valuable responses.

$**Q. The current evidence does not support the claim that$\gamma_\phi$learned in this manner is empirically or theoretically beneficial for real-world ...**$

We have uploaded additional experimental results to support the usage of $\gamma_\phi$ determined by the noise vector $x_T$. Specifically, we experimented with three different inputs for $\gamma_\phi$:

1. a vector filled with $1$
2. a random vector $z$ sampled from $\rho_T$ but unrelated to $x_T$
3. $x_T$.

As shown in Table 5 of the supplementary material, using $x_T$ as the input for $\gamma_\phi$ consistently achieves better performance across three different methodologies. This provides empirical evidence supporting the advantage of the "Adaptive Multidimensional Coefficient" over using the same coefficient for all different $x_T$.

Thank you for pointing out the missing logic in our paper, which we have now clarified.

First of all, thank you for your valuable questions. We have uploaded the extra experimental results you requested in the supplementary materials.

$**Q1-1. I still don't understand why the flow matching model is considered harder to tune using your method, while the stochastic interpolant is ...**$

We first want to point out that for MTO, it is essential to know each $\hat{x}_0$ and $\hat{x}_1$ separately. The reason why it is challenging to perform MTO on FM lies in the target value of the FM model. All three methods (SI, EDM, and Flow Matching) target different values, as shown below:

1. Flow Matching: $H_\theta(t, x(t)) = \hat{v} \approx v = \dot{\gamma}_0 x_0 - \dot{\gamma}_1 x_T$
2. SI: $H_\theta(t, x(t)) = [\hat{x}_0, \hat{x}_1] \approx [x_0, x_1]$
3. EDM (Diffusion): $H_\theta(t, x(t)) = \hat{x}_0 \approx x_0$

For SI and EDM, the model's target value is disentangled from its time schedules. Thus, we can directly compute $[\hat{x}_0, \hat{x}_1]$ (for EDM, $x_1 = \frac{x(t) - \hat{x}_0}{\gamma_1}$) without dependency on the coefficients $\gamma$. This makes MTO feasible, as we can explicitly control $\gamma$ using $\phi$ during the sampling stage.

However, in FM, the model's target value is entangled with the coefficients $\gamma$, as seen above. To calculate $\hat{x}_0$ and $\hat{x}_1$ for MTO in FM, unlike SI or EDM, solving the following simultaneous equations is required, which is computationally onerous:

$H_\theta(t, x(t)) = \dot{\gamma}_0 x_0 + \dot{\gamma}_1 x_1$

$x(t) = \gamma_0 x_0 + \gamma_1 x_1$

Solving these simultaneous equations requires the continuous derivative values of coefficients $\dot{\gamma}$. However, using both $\dot{\gamma}$ and $\gamma$ values is practically inefficient and can confuse $\gamma_\phi$, as numerical errors are inevitable when sampling with a limited number of steps. In other words:

$\dot{\gamma}_\phi(t_i) \neq \Delta\gamma(t_i) = \gamma(t_j) - \gamma(t_i) \quad (j = i+1)$

Given the above equations, we only use $\gamma$ and not $\dot{\gamma}$ for SI and EDM (as detailed in Appendix B). However, to achieve this in FM, onerous computation is required to calculate $\Delta t_i v_{\theta, \phi}$ while avoiding the use of $\dot{\gamma}$. Nevertheless, $we performed MTO in FM on CIFAR-10.$ You can see these results in supplementary material Tables 3, 4, and 5. As shown, the performance of FM is consistently worse than SI. This aligns with the reasoning provided above:

1. FM's model output is not disentangled from $\dot{\gamma}_\phi$.
2. Onerous computational steps are required for MTO in FM to calculate $\Delta t_i v_{\theta, \phi}$, which is computationally inefficient and can accumulate numerical errors.

$**Q1-2. How can your methods be adapted to SI-based image generative models?**$

It is entirely feasible to train SI-based image generative models. The only difference between SI and EDM lies in their target values and parameterization. We performed image generation experiments with SI, but it performed worse than EDM (as shown in supplementary material Table 3). As EDM is known for its high performance due to its well-tuned parameterization, we focused on experimenting with EDM to validate the empirical performance of MTO. We believe this sufficiently demonstrates the validity of our methodology by achieving on-par performance compared to other few-step generation methods.

$**Q5: When referring to "fixed configurations," what specific configurations are being discussed? Are they related to sampler hyperparameters ...**$

Yes, specifically, the time schedule $\tau = \{t_0, \dots, t_N\}$ for inference and the solving methodology (e.g., Euler, 2nd-order Heun) are the fixed configurations being discussed. This can be freely adjusted, as shown in Table 4 of the supplementary materials. We experimented with different NFEs for SI, FM, and DDPM.