Phase-aware Training Schedule Simplifies Learning in Flow-Based Generative Models

Thank you for your thorough review. We address your comments as follows.

W1. The paper misses some very important references on phase transitions in generative diffusion. In the revised version of the paper, we now refer and discuss the relations with the work of Raya and Ambrogioni 2023, Ambrogioni 2023, and Li and Chen 2024. Please let us know if this addresses your concern.

Q1. Could you connect your result to the theory of spontaneous symmetry breaking in generative diffusion models? Our results can be stated as follows in terms of symmetry breaking. (1) The results from Biroli et al (2024) translated to flow-based generative models show that for generating samples from the two-mode GM in dimension $d$, there is a spontaneous symmetry breaking at times of order $1/\sqrt{d}.$ (2) We then show in Proposition 1 that the critical window where the symmetry breaking happens can be made of constant length in $d$ through a time dilation, which we provide in Eq. (10) in the original paper. (3) Finally, analyzing the learning problem using this time dilation shows that the generative model can resolve the spontaneous symmetry breaking, stated formally in Results 1 and 2, and Corollary 5.

W2/Q2. The time dilation formula in Eq.10 can only be calibrated on a single symmetry-breaking point. ... It would be more useful to have a formula that can recalibrate the sampling of data with multiple decision points. The time dilation formula in Eq. 10 can be extended to more than two modes as follows. Consider $\mu=\sum_{i=1}^{m}p_{i}\mathcal{N}\left(r_{i},\sigma^{2}\text{I}\right)$ where $r_{i}\in\mathbb{R}^{d}$ and $|r_{i}|$ goes to infinity with $d,$ but $m,p_{i},\sigma^{2}$ are constant with respect to $d.$ If $X_{t}$ is the generative model associated with the interpolant $I_{t}=(1-t)z+ta$ where $z\sim\mathcal{N}(0,\text{I})$ and $a\sim\mu,$ then $X_{t}$ estimates $p_{i}$ at times of the order $1/|r_{i}|.$ We show this in Proposition 2 in the Appendix, Section E, of the revised version of the paper by arguing that it is only at times of order $1/|r_{i}|$ that the denoiser associated to $r_{i}\cdot X_{t}/|r_{i}|$ is nontrivial. Hence, to estimate $p_{i}$ we require a time dilation $\tau_{t}$ such that there exists $a$ and $b$ with $b-a=\Theta_{d}(1)$ where

From here we can derive the more general dilation formula that you asked for. We need to specify a dilation that for every $i$ ensures that the condition on $\tau_t$ is fulfilled. Assume $|r_{1}|\leq|r_{2}|\leq\cdots\leq|r_{m}|$, let $n=m+1$ and let $\kappa>0.$ Then

Then we have that $p_{i}$ is learnt when $t\in[(m-i)/n,(m-i+1)/n]$ and the $\sigma^{2}$ will be learnt when $t\in[m/n,1],$ giving rise to $m+1$ different phases. In the special case of $|r_{i}|=|r_{i+1}|,$ both $p_{i}$ and $p_{i+1}$ will already be learnt in $[(m-i)/n,(m-i+1)/n]$ so that the phase on the interval $[(m-i+1)/n,(m-i+2)/n]$ is unnecessary. Taking this consideration into account when using the general formula we have derived applied for the two-mode GM gives the time dilation formula from Eq. (10) in our paper. The only difference is that the time dilation here maps $[0,1]$ to $[0,1]$ and in the paper we map $[0,1]$ to $[0,2].$ We included this in the Appendix, Section E, of the revised version of the paper. Note that the idea to detect transitions via comparing the coefficient of data and noise, expressed formally in Lemma 7, is quite general and is a way to extend the time dilation to even more general data distributions.

Thanks for your review. We address your comments as follows.

W1. The standalone paper does not show any experiments or proof of them. The original paper has experiments in Section 6.

W3. The paper should include clear picture with network architecture. The architecture we analyze is shown in Equation 8. Since it is a very simple two-layer Denoising Autoencoder, we believe that showing the formula mathematically is best to describe it. On the other hand, the architecture used for the experiments is described in the Appendix, Section D, and is also simple. We stress that these are standard network architectures widely used by previous work.

W4. It is almost only theoretical without clearly presenting the necessary experiments or methods that would be expected. ICLR is a multidisciplinary conference that regularly accepts papers that are purely theoretical. What experiments or methods are you referring that would be expected?

W5. The paper is hard to follow and has some weirdly written sentences. We fixed the typo in the sentence you mentioned and went through the paper to check its grammatical correctness. Is there anything else in the paper that makes it hard to follow?

W6. The paper lacks quantitative and qualitative comparisons. Two graphs included show more of an ablation study and not the comparison. We note that both of our experiments show the benefit of dilating time by comparing the dilated interpolant with the non-dilated interpolant. What comparison do you have in mind?

W7. It has a well-written theoretical side but without any visual or quantitative results, it is not sufficient to defend the theory. We provide quantitative results in Section 6.

Q1. Do you have a visual comparison of the MNIST dataset? If so, can you include it in the main paper? In the Appendix, Section C2, we show samples from MNIST using the dilated interpolant and the non-dilated interpolant. Does this address your concern about a visual comparison of the MNIST dataset?

Q3. Could you please check the grammatical correctness of the paper? Some sentences have unnecessary brackets and overall the paper is hard to follow and understand. We went through the paper and checked for its grammatical correctness. Please let us know if the revised version has anything that makes it hard to follow.

Thank you for your detailed review. We address your comments as follows

P1. For datasets with more than two modes, it is not obvious that there are always two phases along the diffusion path. Our theoretical argument claiming that a neural network can learn the different phases is only for the two-mode Gaussian Mixture (GM). However, we prove below the existence of more than two phases when the data comes from a GM with more than two modes. (Showing that a neural network can learn more than two phases remains an open problem.) Consider $\mu=\sum_{i=1}^{m}p_{i}\mathcal{N}\left(r_{i},\sigma^{2}\text{I}\right)$ where $r_{i}\in\mathbb{R}^{d}$ and $|r_{i}|$ goes to infinity with $d,$ but $m,p_{i},\sigma^{2}$ are constant with respect to $d.$ If $X_{t}$ is the generative model associated with the interpolant $I_{t}=(1-t)z+ta$ where $z\sim\mathcal{N}(0,\text{I})$ and $a\sim\mu,$ then $X_{t}$ estimates $p_{i}$ at times of the order $1/|r_{i}|.$ We show this in Proposition 2 in the Appendix, Section E, of the revised version of the paper by arguing that it is only at times of order $1/|r_{i}|$ that the denoiser associated to $r_{i}\cdot X_{t}/|r_{i}|$ is nontrivial. Hence, to estimate $p_{i}$ we require a time dilation $\tau_{t}$ such that there exists $a$ and $b$ with $b-a=\Theta_{d}(1)$ where

We specify next a time dilation that for every $i$ would ensure that this condition on $\tau_t$ is fulfilled. Assume $|r_{1}|\leq|r_{2}|\leq\cdots\leq|r_{m}|$, let $n=m+1$ and let $\kappa>0.$ Then

Then we have that $p_{i}$ is learnt when $t\in[(m-i)/n,(m-i+1)/n]$ and the $\sigma^{2}$ will be learnt when $t\in[m/n,1],$ giving rise to $m+1$ different phases. In the special case of $|r_{i}|=|r_{i+1}|,$ both $p_{i}$ and $p_{i+1}$ will already be learnt in $[(m-i)/n,(m-i+1)/n]$ so that the phase on the interval $[(m-i+1)/n,(m-i+2)/n]$ is unnecessary. Taking this consideration into account when using the general formula we just provided applied for the two-mode GM gives the time dilation formula from Eq. (10) in our paper. The only difference is that the time dilation here maps $[0,1]$ to $[0,1]$ and in the paper we map $[0,1]$ to $[0,2].$ We included this in the Appendix, Section E, of the revised version of the paper.

Although we believe this sharp asymptotic characterization of times where different parameters are learnt is novel, we stress that the main contribution of our paper is to provide the first theoretical argument showing that a neural network can learn the two phases in the two-mode GM, providing the next step to Cui et al 2024 who assume a symmetric two-mode GM. In fact, as far as we are aware it is the first theoretical result linking the choice of training schedule with speciation times. Since in practice the training schedule is just determined by drawing times uniformly, this result could open the door to new and more efficient training methods. It is characteristic of machine learning theory (and specially of new analysis methods) that they will unavoidably deal with simpler settings than those from practice. Indeed, the denoiser for the GM with more than two modes does not have a simple expression like the one for the two-mode GM (see Albergo et al 2023 for the general case.) This makes the analysis much harder, but it is a clear direction for future work.

P2. Could the authors comment on the expected behaviour for a Gaussian mixture with three clusters? You asked about the case where $r_{1}=(-1,\cdots,-1),r_{2}=(0,\cdots,0),r_{3}=(100,\cdots,100)$ where $r_{i}\in\mathbb{R}^{d}$ and $p_{1}=p_{2}=p_{3}=1/3,$ $\sigma^{2}=1.$ The analysis from above yields that $\tau_{t}$ needs to be of order $1/\sqrt{d}$ for an interval of constant length to recover $p_{1}$ and of order $1/(100\sqrt{d})$ in another interval of constant length to recover $p_{3}.$ We note that $|r_{2}|=0$ so that it does not fulfill our assumption that $|r_{2}|$ goes to infinity with $d.$ This assumption is not needed, and the analysis can be extended for general $|r_{i}|.$ However, to understand your question this is not needed: indeed, estimating $p_{1}$ and $p_{3}$ correctly ensures correct estimation of $p_{2}$ since these three numbers add to $1$. We hence have three phases in your proposed problem: for times of order $1/(100\sqrt{d})$ where the parameter $p_{3}$ is estimated, for times of order $1/\sqrt{d}$ the parameter $p_{1}$ (and hence $p_{2}$) is estimated, and for times of order $1$ the paramter $\sigma^{2}$ is estimated.

P3. There may be multiple occurrences of cluster splitting, and without prior knowledge of the underlying probability density, it is difficult to derive explicit formulas. This limits the broader applicability of the approach to more general probability distributions. It is true that for a general probability distribution, we cannot derive an explicit formula for the time where the phase transition occurs. However, explicit formulas are not needed to determine these times. Indeed, in Section 6.2 of the original paper we provided a practical approach for determining this time for any probability distribution in terms of the U-Turn and exemplified the method for the MNIST dataset. Does this address your concern?

P4. Additionally, it is unclear why the phase transitions along the diffusion path are not evident from the mollification process described by the Ornstein-Uhlenbeck SDE. It is worth clarifying that we are not considering the Ornstein-Uhlenbeck SDE but instead we are working under the framework of the probability flow ODE, so that given the initial condition there is no noise added in the generation process. On the other hand, it is True that the phase transition times can be determined by looking at the ratio of the data and noise coefficients of a noisy version of the data $I_{t}=(1-t)z+ta$ where $z$ is noise and $a$ is data (this is shown formally in Lemma A3 below). We used this idea implicitly for our results, but we will add this explicit formulation from Lemma A3 to the paper. Is this what you meant by the phase transitions being evident from mollification process?