Paper: Generative Modeling with Phase Stochastic Bridge (ICLR-2024)

Reviewer comment (Presentation):
The paper could provide more insights into the potential applications of the AGM framework beyond image generation.

Author response:
#### **1. potential applications of the AGM framework beyond image generation.**
This is a great suggestion and we do believe that this framework might be particularly useful for trajectory inference tasks in which the Newtonian dynamics (phase dynamics) constraint is imposed, for example, single cell trajectory inference [2] or molecular dynamics  [3] simulation. Combining our framework with the approach described in reference [1] would not only be interesting but could also provide a scientifically robust foundation for addressing these trajectory inference challenges. 

Regrettably, owing to the authors' limited expertise in these particular domains, they cannot make further concrete assertions regarding the appropriateness of AGM within these domains.

[1] Liu, Guan-Horng, et al. "Generalized Schr\" odinger Bridge Matching." arXiv preprint arXiv:2310.02233 (2023).

[2] Tianrong Chan et al . 'Deep Multi Marginal Momentum Schrodinger Bridge'

[3] Holdijk, Lars, et al. "Stochastic Optimal Control for Collective Variable Free Sampling of Molecular Transition Paths." Thirty-seventh Conference on Neural Information Processing Systems. 2023.

Reviewer comment (Experiments):
Can the AGM framework be applied to other domains beyond image generation, such as natural language processing or time series data?

Author response:
#### **1. Can the AGM framework be applied to other domains beyond image generation, such as natural language processing or time series data?**
Similar to the diffusion model [1,2] for discretized variables in language generation tasks, it is possible to utilize AGM for language generation. In this case, the transition matrix $\Phi(\cdot,\cdot)$ will be matrix-lized. We leave further explorations in sequence modeling as future work which we believe is an exciting direction.

[1]Zhang, Yizhe, et al. "PLANNER: Generating Diversified Paragraph via Latent Language Diffusion Model." arXiv preprint arXiv:2306.02531 (2023).

[2] Austin, Jacob, et al. "Structured denoising diffusion models in discrete state-spaces." Ad

Reviewer comment (Writing):
The exposition is rather dense and, as a consequence, the paper is somewhat difficult to read. This is a pity since the underlying concept are rather intuitive and can be understood by a wide audience.

Author response:
We apologize for the difficulties caused by the dense exposition. We add one more gentle introduction of SOC in the appendix which hopefully can increase the readability.

Reviewer comment (Reproducibility):
I find the pseudocode in Algorithm 1 and 2 to be rather uninformative. A good pseudo-code should allow the reader to implement the algorithm almost without referring to the rest of the paper. In this case, the most important parts of the code (e.g. the form of the loss) are omitted. Could you update it to make it more self-contained?

Author response:
Thanks for this valuable suggestion. In the revised version, we have included an enhanced pseudocode that provides more comprehensive information.

Reviewer comment (Evaluation):
The idea of the initial velocity conditioning is interesting, but it is difficult to evaluate its potential without quantitative results and comparisons. Intuitively, it seems to me that it will likely lead to a substantial drop in diversity. Can you report the FID for the conditional sampler?

Author response:
The velocity conditioning experiment was designed to qualitatively showcase the properties of the velocity space, but the reviewer touched upon a subtle but very interesting complication (thank you!). It is indeed leading to a substantial drop in diversity. We found that it corresponds to the value of the hyperparameter $\xi$. We provide an ablation study in the Appendix F. When $\xi=0$, the trajectory degenerates to the unconditional case which leads to the highest diversity with the lowest faithfulness of reference data point. When $\xi$ is large, the diversity drops dramatically but the faithfulness increases. To achieve a balance between faithfulness and diversity, one needs to tune the hyperparameter $\xi$.

Reviewer comment (Experiments):
In Table 4 which shows experiment results on ImageNet 64, DDPM uses a stochastic sample while the other approaches including FM-OT, MFM, and AGM-ODE all use deterministic samplers. This may not be a fair comparison because even for the same type of diffusion model, the sample quality and sampling efficiency could be very different with different types of sampler and deterministic samplers based on ODE numerical methods generally take less number of steps than stochastic sampler. I would suggest the author include both AGM-SDE and AGM-ODE results under different number of evaluations.

Author response:
We deeply thank the reviewer for all the comments. The summary is accurate and the questions are interesting and helpful.

Please kindly see our itemized replies below to address the reviewer’s concerns.
#### **1. Unfair comparison with DDPM and FID evaluation for different NFE**
We totally agree with the reviewer's suggestion and we have removed DDPM from the table. 

For AGM, We reported the performance of our model at 0.8M training iteration which turns out the model is under-trained. After the submission, we kept training the model for another 0.8 iterations and now the performance has increased dramatically. We achieved FID 10.97 at 40 NFE compared with SoTA model [1] (to the best of our knowledge) which achieved 11.82 at 132 NFE. We updated the table with varying NFE as reviewer suggested. 

Regrettably, in our model, one can only opt to train either AGM-SDE or AGM-ODE, and it is not possible to switch between the two interchangeably like a diffusion model. We acknowledge that this limitation is one of the drawbacks of our approach. Moreover, we deeply apologize for not being able to train an AGM-SDE from scratch during the rebuttal phase, primarily due to the significant training complexity associated with ImageNet. Given our current computational resources (8 x Nvidia A100), training an AGM-SDE from scratch would require approximately 8 weeks [1], for which we sincerely apologize.

[1] Karras et al. "EDM"

Reviewer comment (Writing):
The clarity of the paper will be better if the conditions of the Lemmas and Propositions written in this paper is stated more concretely and with full notations, especially in the appendix.

Author response:
Thank you for raising this issue. we have made revisions to the notations and refined the statements of the lemmas and propositions (Lemma 7 and Section D.10 in appendix, Proposition 5 in main paper).

Reviewer comment (Writing):
In the Probabilistic ODE part of (7), an additional notation rather than \( g(t) \) is recommended to be used, like \( g_t \to h_t \) in the matrix notation and \( g(t) = h_t \) for BM-SDE part. Because the notation \( g(t) \) or \( g_t \) is abused, it can be misleading that the score term of the probabilistic ODE is neglected.

Author response:
Thank you for careful reading! We have fixed this issue in the revision!

Reviewer comment (Theory):
The SOC theorem is only used limitedly; the regularization in terms of \( \int ||a_t||^2 \) is ignored and this can threaten the stability of the acceleration space, even though this is not directly revealed in the paper.

Author response:
We would like to respectfully stress that the regularization of the acceleration $||a_t||_2^2$ is not ignored in our framework (see, eg, Equation 5).

To be more precise, as $r$ approaches positive infinity, there still exist multiple possible stochastic processes which start from $x_0$ at $t=0$ and converge to $x_1$ at $t=1$ and all of them are minimizers of the objective function **without** regularization. Among these processes, our preferred stochastic process should also minimize the control effort $\int_{0}^{1}||a_t||_2^2 dt$ (regularization). This solution, which aligns with the unique solution presented in Definition 2 (In general, the solution of SOC is not unique. But in our simple Linear Quadratic case, the solution is unique, see Theorem 6.1 in [this textbook](https://www.control.utoronto.ca/people/profs/kwong/ece410/2008/notes/chap6-08.pdf)), is what we have employed to construct our methodology.

To address any confusion that may have arisen, we have provided additional clarification in Appendix C.1 in the revision. This supplementary section aims to resolve any ambiguities and enhance the understanding of the subject matter.

Reviewer comment (Experiments):
Whereas the theoretical background is sound and the improved performance is guaranteed, the hyperparameters such as the diffusion coefficients and the SDEs are not optimized, which causes its lacking performance compared to EDM (look at Figure 5). However, this is expected to be enlightened with further works.

Author response:
The experiment we evaluate on is the **unconditional** imagenet-64 and we do not have computation resources to conduct the training on this dataset for all baselines. As a result, We report the performance of baseline models from a recent paper (see Table.8 from [1]).

For AGM, We reported the performance of our model at 0.8M training iteration which turns out that the model was under-trained. After the submission, we kept training the model for another 0.8M iterations and the performance increased dramatically. We achieved FID 10.97 at 40 NFE compared with SoTA model [1] (to the best of our knowledge) which achieves 11.82 at 132 NFE. 

We kindly invite the reviewer to suggest any other baselines we are missing, and we are more than happy to incorporate them into our paper.

[1] Pooladian, Aram-Alexandre, et al. "Multisample flow matching: Straightening flows with minibatch couplings." arXiv preprint arXiv:2304.14772 (2023).

Reviewer comment (Presentation):
It will be helpful if the acceleration coefficient \( a_t \) for image datasets is depicted, as the momentum of how the image data is being generated in Figure 1 or Figure 2. It is expected that the acceleration coefficients show similar semantic features like \( x_t \) and \( v_t \), but have varying scales.

Author response:
We have added the plot of position, and velocity together with acceleration in the appendix G in the revision. Upon reviewing Appendix G, it becomes evident that the trajectories of the variables remain consistent over varying NFEs, despite the trajectories of any variable not being perfectly linear in the infinite timestep limit.

Reviewer comment (Writing):
Can you provide elementary introduction of the stochastic optimal control? While this paper works only the simple case of the SOC (no regularization case), introducing some details, or at least some introductory materials will help the readers to follow up the backgrounds.

Author response:
Thanks for the valuable suggestion. We have provided a gentle introduction of stochastic optimal control in the appendix C.

Reviewer comment (Experiments):
I guess that the ImageNet64 performance is not yet optimized: the generative performance of the SoTA models are expected to be much better than the paper have proposed. I think at least the performance should be compared with CLD-SGM from the same architecture.

Author response:
The experiment we evaluate on is the **unconditional** imagenet-64 and we do not have computation resources to conduct the training on this dataset for all baselines. As a result, We report the performance of baseline models from a recent paper (see Table.8 from [1]).

For AGM, We reported the performance of our model at 0.8M training iteration which turns out that the model was under-trained. After the submission, we kept training the model for another 0.8M iterations and the performance increased dramatically. We achieved FID 10.97 at 40 NFE compared with SoTA model [1] (to the best of our knowledge) which achieves 11.82 at 132 NFE. 

We kindly invite the reviewer to suggest any other baselines we are missing, and we are more than happy to incorporate them into our paper.

[1] Pooladian, Aram-Alexandre, et al. "Multisample flow matching: Straightening flows with minibatch couplings." arXiv preprint arXiv:2304.14772 (2023).

Reviewer comment (Presentation):
It will be helpful for the readers' understanding if you use the colored hyperlinks by reference (\ref) or citation (\cite, \citep, \citet) commands.

Author response:
We have added the hyperlink color back in the revision for better visualization as the reviewer suggested!