Generative Modeling with Phase Stochastic Bridge

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. 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.

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

We extend our sincere gratitude to the reviewer for the valuable comments. Please kindly find below our itemized responses, presented in an effort to address each of the reviewer’s concerns.

1. The exposition is rather dense and, as a consequence, the paper is somewhat difficult to read.

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.

2. pseudocode in Algorithm 1 and 2 to be rather uninformative.

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

3. For conditional generation, it seems to me that it will likely lead to a substantial drop in diversity.

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$.

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. Larger scale dataset

We understand the reviewer's concern and we agree that AGM may face unexpected difficulties when scaled to larger resolutions which is also the known issue [1,2] posed for general Dynamical Generative Modeling including Diffusion Model and Flow Matching. This issue partially (informally) explains the success of latent-space-based Diffusion Models [3,4] in high-resolution scenarios.

Due to the computation resource and time limitation, we were unable to conduct experiments at larger scale about which we apologize.

This is something of high priority in our todo list for further work.

[1] Jiatao Gu et al. 'Matryoshka Diffusion Models'

[2] Zahra Kadkhodaie et al. 'Learning multi-scale local conditional probability models of images.'

[3] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models. 2022 IEEE." CVF Conference on Computer Vision and Pattern Recognition (CVPR). 2021.

[4] Vahdat, Arash, Karsten Kreis, and Jan Kautz. "Score-based generative modeling in latent space." Advances in Neural Information Processing Systems 34 (2021): 11287-11302.

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"

We would like to express our sincere gratitude for your valuable feedback and comments. We truly appreciate the time and effort you invested in assessing our submission.

Please kindly see our itemized replies below in order to address the reviewer’s concerns.

1. The clarity of the paper will be better if the conditions of the Lemmas and Propositions written in this paper is stated more concretely

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).

2. Both this method and CLD make predictions of the data from both the current state and the velocity, providing the data prediction result from CLD.

Thank you for mentioning this aspect which is indeed valuable.

It has come to our attention that the training objective of CLD solely focuses on the score function with respect to velocity (scaled $\epsilon_1$). Consequently, there is no apparent method for CLD to effectively reconstruct the data point using the intermediate state and velocity. Specifically, in order to reconstruct $x_1$,the information of stochasticity in the position channel, denoted as $\epsilon_{0}$, is required. However, it is important to note that this information is not available after the training phase of CLD. Therefore, the reconstruction of $x_1$ is not feasible.

In our case, the special structure of acceleration, which includes the stochasticity information of state and velocity channel (eq.9 in revision), allows us to reconstruct the $x_1$ (Proposition.5).

To address any confusion resulting from our previous insufficient explanation, we have included additional clarification in Section 3.2 in revision.

3. Notation abuse in eq.7.

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

4. provide an elementary introduction of the stochastic optimal control?.

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

5. the ImageNet64 performance is not yet optimized, the generative performance of the SoTA models are expected to be much better than the paper has proposed.

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).

6. HyperLink color

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

7. Trajectory of acceleration

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.

8. The problem formulation ignores the regularization of acceleration

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), 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.