Neural Diffusion Models

We thank the reviewer for their review and suggestions which will help improve the paper. We will include the suggested reference in the revision. Below we address the questions and comments raised in the review.

Weaknesses and Questions:

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.

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.

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.

[1]:Dongjun Kim et al. 'Maximum Likelihood Training of Implicit Nonlinear Diffusion Models.'

[2]: Tianrong Chan et al. 'Likelihood Training of Schrödinger Bridge using Forward-Backward SDEs Theory.'

[3]: Qinsheng Zhang et al. 'Fast Sampling of Diffusion Models with Exponential Integrator'

[4]: Fan Bao et al. 'Analytic-DPM'

1. I totally agree with authors' argument that the full ELBO will lead to lower FID and better NLL, however I still insist that the numerical value FID lags too much. A kind suggestion is that, the authors can demonstrate some experimental setup in which the NLL is more important than other quantitative metrics. The image generation tasks rely quite much on FID because the generated quality counts more. The generated quality is not satisfying even though the higher NLL you achieved (Fig.2).

2. At the end of the day, the time dimension will be considered as part of batch size for both NDM and SB. The only difference is that NDM chops the time dimension (just randomly sample one timestep which results in [batch, 1,dimension]), and SB chops the batch dimension into small pieces (which results in [small-batch, T, dimension]). I still want to argue that the proposed method is not really 'simulation-free'.

EDIT****** My apology for the incorrect term. What I would like to argue is that the complexity will still be high. The proposed method is indeed simulation-free but with comparable complexity like simulation-based approaches (SB).

We thank the reviewer for the suggestions and review which will help to improve the paper. We are glad that the reviewer finds our idea to make sense and empirical results to be illustrative. Below we address the questions and comments raised in the review.

Weaknesses:

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

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

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

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

We are glad that the reviewer founds the description of NDM to be clear and the visualizations to be helpful. Below we address the questions and comments raised in the review.

Weaknesses:

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

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

• We don’t utilize any balancing coefficients. We always train models with the pure ELBO (eq. 9). However, thanks to parameterization of the transformation that guarantees $F_\varphi(\mathbf{x}, t) = \mathbf{x}$ for $t = 0$ we don’t need to optimize the reconstruction term (see Appendix C.2, (eq. 37)).

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

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

We are delighted that the reviewer found our idea of generalization and the qualitative results of learnable transformations intriguing. Below we address the questions and comments raised in the review.

Weaknesses:

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

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

Questions:

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

To incorporate the feedback and suggestions from reviewers, we have revised the paper. The key modifications are following:

• In Section 3, we introduced algorithm boxes (Algorithms 1 and 2) outlining the training and sampling procedures within the NDM framework.
• In Table 4, alongside FID scores for stochastic DDPM-style sampling, we have now included FID scores for deterministic DDIM-style sampling.
• In Appendix B, we have expanded discussions and comparisons of NDM with other existing models, including Stochastic Interpolants [1] and Schrödinger Bridge-based models.

We thank the reviewers for their valuable assistance in improving the quality of our work.

[1] Building normalizing flows with stochastic interpolants. https://arxiv.org/abs/2209.15571