Paper: Neural Diffusion Models (ICLR-2024)

Reviewer comment (Novelty):
There are a few other learnable forward process which generalizes the diffusion model [1,2]. The authors should consider citing or comparing with them.

Author response:
1. We present NDM not merely as a model learning the forward process but as a comprehensive framework that enables the predefined specification of forward dynamics capable of information destruction. Additionally, it facilitates the learning of these dynamics. The concept of a learnable forward process is not novel. However, unlike others, NDM allows time-dependent, non-linear, and learnable data transformations, all while remaining a simulation-free framework, which is an essential factor for the efficiency of the training procedure.
    
    For instance, in [1], the authors suggest training a conventional diffusion model in the latent space of a flow model. This model can be viewed as a special case of NDM, where the transformation $F^\varphi$ corresponds to a time-independent flow.
    
    The work [2], like many other Schrödinger Bridge models, does not fall into the special case of NDM. However, in contrast to NDM, the approach [2] is not simulation-free. In Schrödinger Bridge models, we typically lack access to the distribution $q(\mathbf{z}^t|\mathbf{x})$. Consequently, to sample the latent variable $\mathbf{z}^t$ in training time, we must simulate the full stochastic process, such as the SDE in [2]. This characteristic makes Schrödinger Bridge models expensive in training and not simulation-free. In contrast, NDM, by design, has access to $q(\mathbf{z}^t|\mathbf{x})$. Thus, with NDM, we can directly sample $\mathbf{z}^t$ without expending computational resources on simulation, enhancing the efficiency of the training procedure.

Reviewer comment (Experiments):
Actually I am not quite convincing by the effectiveness of proposed method. The numerical value is compared with the DDPM. The authors clearly present the differences with DDIM in Figure 1, but the experiment section is just focus on DDPM for fast sampling. The numerical value provided by authors is not impressive or competitive compared with the other fast sampling techniques on the market which is built on DDPM/DDIM [3,4].

Author response:
2. While NDM can be thought of as a generalisation of DDIMs, the DDIM-style deterministic sampling effectively implies sampling from a process different from the one with which the model was trained. Thus, for a comparison, we opted to sample from the same processes that the model was trained on, which coincides with the DDPM reverse process. However, since NDMs allow deterministic sampling we acknowledge the interest of this alternative comparison and plan to incorporate it into the camera-ready version of our paper to compare FID scores. The switch from DDPM to DDIM doesn’t affect evaluation of likelihood, since it relies on integration of deterministic process in both cases.
    
    Nevertheless, in addition to the comparison with DDPM presented in Table 4, we also provide a comparison of NDM with other newer models in Tables 2 and 3.
    
    Simultaneously, there exist studies proposing techniques to enhance the performance of diffusion models or expedite the sampling procedure [3, 4]. However, most of these approaches are orthogonal to NDMs and can be easily adapted for further gains.

Reviewer comment (Novelty):
What is the difference between the NDM compared with [1,2]? What is the benefits over them?

Author response:
1. We present NDM not merely as a model learning the forward process but as a comprehensive framework that enables the predefined specification of forward dynamics capable of information destruction. Additionally, it facilitates the learning of these dynamics. The concept of a learnable forward process is not novel. However, unlike others, NDM allows time-dependent, non-linear, and learnable data transformations, all while remaining a simulation-free framework, which is an essential factor for the efficiency of the training procedure.
    
    For instance, in [1], the authors suggest training a conventional diffusion model in the latent space of a flow model. This model can be viewed as a special case of NDM, where the transformation $F^\varphi$ corresponds to a time-independent flow.
    
    The work [2], like many other Schrödinger Bridge models, does not fall into the special case of NDM. However, in contrast to NDM, the approach [2] is not simulation-free. In Schrödinger Bridge models, we typically lack access to the distribution $q(\mathbf{z}^t|\mathbf{x})$. Consequently, to sample the latent variable $\mathbf{z}^t$ in training time, we must simulate the full stochastic process, such as the SDE in [2]. This characteristic makes Schrödinger Bridge models expensive in training and not simulation-free. In contrast, NDM, by design, has access to $q(\mathbf{z}^t|\mathbf{x})$. Thus, with NDM, we can directly sample $\mathbf{z}^t$ without expending computational resources on simulation, enhancing the efficiency of the training procedure.

Reviewer comment (Theory):
How do the authors obtain the $\hat{x}_{\theta}$ in eq.10? Does it include the inference of the network? If yes, then why the algorithm is simulation-free? From my understanding, it is implicitly simulating the dynamics which is also the part of reason for heavy training complexity as stated by the authors.

Author response:
3. We present the derivation of equation 10 in Appendix A.2. The computation of this equation requires the inference of the transformation $F^\varphi$. Importantly, this does not disrupt the simulation-free paradigm. As described earlier, a non-simulation-free term implies the necessity to simulate certain dynamics during the training procedure. In NDM, we can sample $\mathbf{z}^t$ with just one inference of $F^\varphi$, introducing increased complexity only as a constant without altering the order of complexity.
    
    In our experiments, the introduction of a learnable $F^\varphi$ extends the training time by approximately 2.3 times. However, in a non-simulation-free approach with $T$ timestamps, employing a learnable forward process would result in a $T$-fold slowdown.

Reviewer comment (Reproducibility):
The presentation of the paper should be polished. The learning and sampling algorithm is difficult to find. I suggest the authors shorten/defer the discussion section to Appendix and add algorithm boxes in the main text. Moreover, I have several questions on the training process: (1) Are $\phi$ and $\theta$ jointly trained or alternatively trained? (2) How are the hyper-parameters of $F_\phi$ set? Are they similar to the x-prediction network?

Author response:
We agree that adding algorithm boxes can help improve readability. We will include these boxes in the camera-ready version if accepted.
    
    (1) We optimize the parameters of the forward process $\varphi$ and parameters of the reverse process $\theta$ jointly by optimizing the full ELBO objective (eq. 9).
    
    (2) As we discuss in Section 4.1, in all our experiments involving learnable transformations, we employ the same neural network architecture and set of hyperparameters for both the transformations $F^\varphi$ and prediction of datapoints $\hat{x}^\theta$.

Reviewer comment (Theory):
Certain constraints of $F_\phi(x_t, t)$ should be satisfied. For example, when $t=0$, I think it must satisfy $F_\phi(x_0, 0) = x_0$. I wonder how the authors ensure that.

Author response:
Upon training NDM with the complete ELBO objective (equation 9), additional restrictions are not necessary. However, as outlined in Appendix C.2, in practice, we ensure that $F^\varphi(\mathbf{x}, t) = \mathbf{x}$ for $t = 0$. This is achieved through the following reparameterization: $F^\varphi(\mathbf{x}, t) = (1 - t) \mathbf{x} + t \bar{F}^\varphi(\mathbf{x}, t)$. This ensures that the reconstruction term $\mathcal{L}^\mathrm{rec}$ doesn’t depend on parameters $\varphi$.

Reviewer comment (Experiments):
The proposed method brings additional computation overhead in the inference time when simulating Eq. (12) and Eq. (13). Because $F_\phi$ is a neural network and the authors use the same U-Net for both $F_\phi$ and $\hat{x}$, it at least doubles the computation. Visually, Figure 2 shows that $F_\phi$ actually gives very close prediction to the real data $x$. I guess maybe simply double the size of the original $\hat{x}_\theta$ can get the similar results as NDM.

Author response:
To address the question of whether the improved performance of NDM is due to the proposed method or merely the result of increasing the number of model parameters we provide an ablation study in Appendix D.3, which demonstrates that the improved quality is not a consequence of an increased number of parameters.

Reviewer comment (Experiments):
Ablation studies are missing. At least the influence of (1) the various choices of $\alpha_t$ and $\sigma_t$, and (2) the various choices of the network structure and number of parameters of $F_\phi$ should be investigated to verify that NDM brings consistent improvement.

Author response:
We introduce NDMs as a framework that generalises existing approaches and enables further extensions through learnable transformations. The main purpose of our experiments is not to obtain state-of-the-art performance nor uniformly better performance compared to baselines. Instead, our aim is to establish the functionality of NDMs in a broader sense, highlighting their additional flexibility and studying learnt transformations.
    
    We believe that NDM with learnable transformations may exhibit enhanced performance with different neural network architectures, improved parameterization, and after hyperparameter tuning. However, we defer this aspect to future research. In this work, we wanted to keep experimental setup simple to showcase the inherent properties of NDM, prioritizing this over achieving superior results through engineering efforts.
    
    Nevertheless, as mentioned above, we conducted an ablation study that substantiates the enhanced performance observed in experiments, attributing it to the model itself.

Reviewer comment (Novelty):
Relationship to learnable interpolation in [1] should be discussed in more detail, although I notice there are several sentences in the related works. The two methods are actually very close to each other (I think the only difference is the objective).

Author response:
While NDM and stochastic interpolants [1] are similar in spirit, these methods differ significantly. Firstly, stochastic interpolants represent an approach to learning continuous-time deterministic generative dynamics, whereas NDM learns stochastic dynamics in either discrete or continuous time, which can subsequently may be transformed into a deterministic process.
    
    Secondly, in NDM, the model is trained by optimizing the variational bound on the log likelihood, while stochastic interpolants are trained by optimizing the generalization of the Flow Matching objective [2].
    
    Lastly, NDM joint learns both the forward and reverse processes by optimizing the likelihood, whereas stochastic interpolants learn the generative process with a fixed forward process. The authors of stochastic interpolants demonstrate the possibility of constructing an optimization procedure for the forward process through a max-min game to solve a dynamic optimal transport problem. However, the purpose of this optimization differs from that of NDM.
    
    Moreover, max-min optimization, as employed in stochastic interpolants, is notably less stable compared to min-min optimization in NDM. Additionally, stochastic interpolants do not present experimental results for the optimization of the forward process.
    
    Therefore, we find stochastic interpolants intriguing and promising, as well as related to our work. However, we believe that this approach differs significantly from NDM. We plan to include an extended discussion of stochastic interpolants along with a comparison to NDM in the camera-ready version if accepted.

Reviewer comment (Novelty):
It seems that the technical details are similar to the conventional DMs. So, what are the technical challenges and novelties here?

Author response:
The main novelty of our work is NDM as a simulation-free framework that generalises conventional diffusion models and allows non-linear, time-dependent and learnable transformations of the data distribution in both discrete and continuous time.
    
    While NDM is a generalisation of conventional diffusion models, it shares many technical similarities with them. However, implementing NDM involves addressing several technical challenges. 
    
    Firstly, in order to consistently define the forward process, we needed to shift to a non-Markovian definition of the model. This required deriving consistent distributions $q(\mathbf{z}_t|\mathbf{x})$ and $q(\mathbf{z}_s|, \mathbf{z}_t, \mathbf{x})$, as well as parameterizing the reverse process and deriving objective.
    
    Secondly, to extend NDM to continuous time, we derived stochastic and ordinary differential equations (SDE and ODE) along with the corresponding objective. Lastly, for an efficient practical implementation of the continuous-time model, we devised a method to calculate time derivatives $\frac{\partial F_\varphi}{\partial  t}$ using Jacobian Vector Product (JVP). Additionally, we applied importance sampling to ensure the stability of the training procedure.

Reviewer comment (Novelty):
In my opinion, the argument "a key limitation of most existing diffusion models is that they rely on a fixed and pre-specified forward process that is unable to adapt to the specific task or data at hand" is not convincing enough. The extensive empirical studies in the community reflect that conventional DMs have enough flexibility to accommodate diverse data. So, more clearly, I want the authors to clarify what practical consequences would the conventional "inflexible" modeling lead to. On the other hand, conventional modeling bears great simplicity, where the denoising process is training-free. Compared to that, NDMs may rely on more complicated training.

Author response:
Conventional diffusion models indeed exhibit remarkable generalization ability and performance across various domains. However, the forward process in these models implicitly predefines the distribution of latent variables and consequently generative trajectories. The fixed forward process thus determines the complexity of the target for the reverse process and the complexity of the sampling procedure. Consequently, the flexibility of the forward process has the potential to simplify the target for the reverse process, resulting in enhanced performance and simplified trajectories that boost sampling.
    
    Recent research in diffusion models indicates that techniques such as diffusion in latent or orthogonal space, a combination of diffusion and blurring, flow matching, and rectified flows contribute to improved performance. These techniques bring modifications to the conventional diffusion process, enhancing sample quality and sampling speed. The NDM framework presented in our work represents the next step in the quest for finding better methods to construct generative dynamics.

Reviewer comment (Evaluation):
More clarifications should be put on the training cost and stability of NDMs, compared to vanilla DMs.

Author response:
In all our experiments, we employ the same neural network parameterization for both transformations, $F_\varphi$ and prediction of datapoint $\hat{x}_\theta$. Therefore, with this parameterization, NDM has twice as many parameters compared to vanilla diffusion models and approximately 2.3 times longer training duration (refer to Section 6). However, in Appendix D.3, we present an ablation study demonstrating that performance improvements are not achieved by increasing the number of parameters.
    
    The only additional technique we employed to ensure training stability is importance sampling of time for continuous-time models, as described in Sections 3.2 and 6. No further techniques where necessary to ensure stable training of NDMs. We will clarify it in camera ready.

Reviewer comment (Experiments):
Is there an experiment on a larger dataset using a larger model? Are there technical challenges to achieving this?

Author response:
In our paper, we present experimental results on the ImageNet 64 and CelebA-HQ-256 datasets, which involve significantly more computational resources compared to experiments on the standard CIFAR-10 dataset. However, across datasets we consistently vary only the scale of the neural network architecture in all our experiments. Thus, the primary technical challenge we identify for experiments on larger datasets with larger models is the associated computational expense.

Reviewer comment (Evaluation):
One of the primary motivations for learnable transformations is that it simplifies the data distribution and therefore leads to predictions of x that are more aligned with data. Ideally, if transformations indeed helped with simplification of data distribution, one should have observed better quantitive metrics in fewer sampling steps. However, the actual gains in quantitative metrics like NLL and NELBO seem marginal. Further, there seems to be no consistent gains in terms of FID. In addition, as listed in limitations, the model uses 2.3$\times$ more training time than DDPM (which by itself is slow and needs hundreds of thousands of steps to get good FID). Therefore, I am not sure if the minor gains in NLL and NELBO can be justified when compared to 2.3$\times$ increase in training time as well as $\sim 2\times$ increase in model size.

Author response:
As soon as NDM with learnable transformation $F_\varphi$ has a more flexible forward process, an improved likelihood estimation is anticipated, as indeed observed. As illustrated in Section 4.2, the incorporation of learnable transformations consistently enhances both NLL and NELBO across a variety of datasets, including state-of-the-art results on ImageNet-64 dataset.
    
    While these learnable transformations simplify the generative process, making the approximation by neural networks easier, it does not necessarily ensure the simplicity of generative trajectories. In theory, it may potentially lead to increased curvature of trajectories. Therefore, the results presented for NDM with a reduced number of steps (Table 4, 1000 → 10 steps) are considered more of an empirical finding rather than an expected property.
    
    While extra flexibility of the forward process leads to consistently improved likelihood estimation, it doesn't always lead to improved FID score. This is a well-known attribute of diffusion models trained with the full ELBO objective (eq. 9).
    
    We introduce the NDM as a novel framework that generalizes many existing diffusion models and allows for a more complex parameterization of the forward process. Therefore, the primary focus of our experimental section is to demonstrate the feasibility of the framework, rather than presenting a stronger model that surpasses baselines. We believe that the NDM may exhibit enhanced performance with different parameterization and hyperparameter tuning. Therefore, we aimed to maintain a simple experimental setup. Nevertheless, our findings illustrate that the NDM, with learnable transformations, yields improved NLL estimation while remaining comparable in terms of FID score. In addition, in Appendix D.3, we present the ablation study, which demonstrates that the improved quality is not a consequence of an increased number of parameters.

Reviewer comment (Experiments):
Benefits of learnable transformations hasn’t been explored much in the paper and as a result its benefits remain unclear. There is a toy experiment on training optimal transport on 1D data in the appendix. However, larger scale experiments on real data would make this work much stronger. Overall, I feel that NDMs are well-motivated but it remains unclear from the experimental results that introducing learnable transformations as a standard practice for training diffusion models is beneficial.

Author response:
While we agree that advantages of NDMs are not yet fully exploited, we consider this a strength of the method rather than a weakness. NDM introduces a new flexible framework for diffusion modelling, derives novel time-discrete and time-continuous bounds on the log-marginal likelihood, studies the type of potential transformations that are learned, delivers promising NLL results on a variety of real and synthetic datasets, as well as discusses its potential limitations. Because of the flexibility of NDMs they enable interesting potential for adapting to the problem at hand and new avenues for future research in generative modelling, like learning dynamic optimal transport that we present in Appendix E.

Reviewer comment (Reproducibility):
After going through the implementation details (Section 4.1, Appendix C), the architectural details of neural network used to model the transformation $F_\psi$ remains unclear to me. Could the authors further elaborate on these details?

Author response:
In all our experiments with learnable transformations we use the same neural network parameterization for both transformations $F_\varphi$ and prediction of the datapoint $\hat{x}_\theta$.

Reviewer comment (Experiments):
Table 4 compares DDPM and NDM. However, NDM has non-Markovian forward and reverse process and for a fair comparison, results for DDIM should also be included, especially in the cases when the number of sampling steps is smaller than the number of training steps.

Author response:
While NDM can be thought of as a generalisation of DDIMs, the DDIM-style deterministic sampling effectively implies sampling from a process different from the one with which the model was trained. Thus, for a comparison, we opted to sample from the same processes that the model was trained on, which coincides with the DDPM reverse process. However, since NDM allows deterministic sampling we acknowledge the interest of this alternative comparison and plan to incorporate it into the camera-ready version of our paper to compare FID scores. The switch from DDPM to DDIM doesn’t affect evaluation of likelihood, since it relies on integration of deterministic process in both cases.