Neural Flow Diffusion Models: Learnable Forward Process for Improved Diffusion Modelling

We are delighted to see that the reviewer finds our approach novel, the theoretical explanations solid, and the experiments robust, and even recognizes the potential to influence future research. Below, we address the questions and comments raised in the review.

Weaknesses:

1. We thank the reviewer for bringing the Shiftddpms [1] work to our attention, which we had overlooked. While Shiftddpms focuses on conditional generation, it indeed shares similarities with our NFDM, proposing learnable transformations of conditioning. Although we do not discuss conditional generation in this paper, we believe that the NFDM framework could be adapted for conditional generation and this is an interesting avenue for future research.

   We will definitely include a discussion of this work in the revision of the paper if it is accepted.

2. In this work, we aimed to focus on the introduction of the NFDM itself and demonstrate its capabilities compared to prior work on standard benchmarks used in those works. Indeed, in the current work, we do not discuss topics such as the application of NFDM to domains other than imaging, scaling, conditional generation, or training with non-ELBO objectives. However, we view these as very promising directions for future research.

[1] Shiftddpms: Exploring conditional diffusion models by shifting diffusion trajectories

We are glad the reviewer found NFDM framework as neat. Below, we address the questions and comments raised in the review.

Weaknesses:

1. We define the distribution $q_\phi(z_t|x)$ implicitly through the function $F_\phi(\varepsilon, t, x)$. Under certain assumptions, we can derive an ordinary differential equation (ODE) as shown in Equation 8 by simply differentiating $F_\phi$ with respect to time. It is straightforward to see that integrating this ODE over time produces the same samples as those generated by $F_\phi$. Consequently, the ODE in Equation 8 corresponds to the defined probability density $q_\phi(z_t|x)$.

   The connection between deterministic and stochastic dynamics through the score function, as leveraged in Equations 9 and 10, is a classical result from [1]. Further proofs demonstrating that the dynamics discussed in Equations 8-10 share the same marginals can be found in [2] (see Appendices A and D.1).

   We will include a more extensive discussion of these derivations in Appendix A in the revised version of the paper, if accepted.

2. Consider the distribution of latent variables $q(z_t)$ defined by the forward process. Given the data distribution $q(x)$ and the conditional distribution $q(z_t|x)$, we can derive the distribution of latent variables $z_t$ as follows: $q(z_t) = \int q(x_t|x) q(x) dx$. Therefore, if the forward process that defines $q(z_t|x)$ is fixed, $q(z_t)$ is also implicitly predefined and fixed. At the same time, the generative process must adhere to all intermediate distributions $q(z_t)$ starting from $t=1$. This is why a fixed forward process limits the flexibility of the latent space.

   From a theoretical perspective, conventional diffusion with a fixed forward process that only injects Gaussian noise is a model that in the ideal scenario works for arbitrary distributions. However, conventional diffusion can be suboptimal, potentially leading to distributions of latent variables that are difficult to learn or highly curved trajectories (Fig. 1) that are expensive to integrate. A learnable forward process could simplify the target for the generative process, thereby improving the alignment of forward and generative processes and better capturing the data distribution as a result.

   We provide empirical evidence that learning the forward process enhances the performance of the generative process. A theoretical motivation for learning the forward process is also detailed in Appendix B.1.

   Additionally, the flexibility of the NFDM framework allows us to learn not just any generative process, but one with desired properties such as straight-line deterministic trajectories (see Section 5), which is unachievable with conventional diffusion that predetermines the distributions of the latent variables.

3. We are pleased that the reviewer appreciated our discussion of connections with related works. We will elaborate on the differences with these approaches in more detail in the revision of the paper, if accepted. We note, however, that in most cases, the parameterizations of NFDM discussed in Appendix E.1 lead to the same objectives as proposed in related works.

Questions:

1. In our experiments, we intentionally parameterize the forward process to ensure that $q_\phi(z_0|x) \approx \delta(x - z_0)$ and $q_\phi(z_1|x) \approx p(z_0)$ (see Appendix B.5). In this way, the model can not collapse by design. For instance, in Eq. 12, $\tilde{f}_\phi(z_t, t, x)$ can not always equal $0$. Furthermore, the NFDM objective is also a variational bound on the log-likelihood (Appendix A.1).

   In practice we did not observe any instabilities in the training procedure and did not apply any techniques to stabilize training.

2. We describe the parameterization of $g_\phi$ in Appendix F. In all our experiments, $g_\phi$ is parameterized using a 3-layer perceptron with 64 neurons in each layer.

3. When sampling with an adaptive ODE solver from models trained on the CIFAR10 dataset, we achieve an FID score of 4.22 for NFDM (with an average of 82 NFE) and 4.85 for NFDM-OT (with 45 NFE on average). For reference, a DDPM trained with the ELBO objective has an FID score of 13.51 (with 1000 NFE). We plan to calculate the FID metrics for the ImageNet dataset by the end of the discussion period.

   Some modifications of diffusion models trained with better parameterizations and different objectives can achieve lower FID scores, but they tend to show worse log-likelihood estimates. This phenomenon is a well-known attribute of diffusion models. We believe that with improved parameterization and hyperparameter- and objective tuning (which was done for conventional diffusion), the flexibility of NFDM could lead to even better performance in terms of FID score, but we leave this for future research.

   Although we have provided experimental results on image datasets as a standard benchmark, we introduce NFDM as a versatile generative modeling framework. In this context, likelihood emerges as the most commonly employed metric in probabilistic modeling across various fields, and NFDM achieves state-of-the-art results. While the FID score is a popular metric for images, it does not always correlate with image quality [3]. In contrast, likelihood is based on robust theoretical principles. We will include a table with these FID scores in the appendix.

We trust that the clarifications further explanations of motivation, technical details and other will enhance your support for our paper's acceptance!

[1] B. D. Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313–326, 1982.

[2] Song, Yang, et al. "Score-based generative modeling through stochastic differential equations.”

[3] Stein, George, et al. "Exposing flaws of generative model evaluation metrics and their unfair treatment of diffusion models."

I appreciate the reviewers' clarification and have decided to increase my score.

We are pleased that the reviewer appreciates the elegance of our framework and recognizes the potential impact of our research direction. In response to the comments and questions raised, we provide the following clarifications and commitments:

Weaknesses:

1. We acknowledge the lack of visual explanations in our paper. Due to space constraints in the main manuscript, we unfortunately did not include them. In the camera-ready version, if accepted, we plan to include a series of figures in the main paper or the appendix. Specifically, we intend to add a diagram that visually explains how the $F_\varphi$ function defines the conditional distributions $q_\varphi(\mathbf{z}_t|\mathbf{x})$, and illustrates the connection between ODEs and SDEs corresponding to the forward process.

   Additionally, in our global response, we have provided visualizations of generated samples from the NFDM and NFBM models trained on various datasets.

2. We have provided a comparison between NFDM and NFDM-OT in Tables 5 and 6 in the appendix. Compared to NFDM, NFDM-OT trades off worse density estimation performance (lower log-likelihood) for better sample quality (lower FID) in the few-step generation mode.

   In the same 2D experiment illustrated in Figure 1, NFDM exhibits slightly less curved trajectories compared to conventional diffusion. We will include additional illustrations in the camera-ready version, if accepted.

We believe that the additional visualizations and comparisons will strengthen your support for acceptance. We are also happy to address any further concerns or questions the reviewers may have.

We appreciate the reviewer's positive feedback regarding the clarity and presentation of our paper. Below, we address the questions and comments raised in the review.

Weaknesses:

1. Please refer to the answer to Question 1 below.

2. Unfortunately, due to limited computational resources, we were not able to conduct experiments on a significantly larger scale than those on ImageNet 64. Consequently, we cannot definitively ascertain the outcomes in a large-scale setting.

   Nevertheless, we would like to highlight that NFDM is a versatile framework that includes conventional diffusion as a special case which has been shown to scale. Furthermore, CIFAR-10 and ImageNet are standard benchmarks in the diffusion literature. Although we lack large-scale experiments to support our position fully, we see no reasons that would impede the scaling of NFDM.

3. Please refer to the answer to Question 6 below.

Questions:

1. The main contribution of our work is the neural flow diffusion model, introduced as a general framework. It enables the construction of diffusion models with a broad spectrum of forward processes, including non-linear, time-dependent, and learnable processes. We show how to efficiently parameterize and train diffusion models with such forward processes, while retaining key properties of diffusion models, such as access to stochastic and deterministic generative dynamics as well as simulation-free training.

   In Appendix E, we discuss in more detail the connections between NFDM and related works that introduce learnable parameters into the forward process. Most existing prior work can be considered special cases of NFDM. For instance, to parameterize the forward process in [1] using NFDM, one could set $F_\phi(\varepsilon, t, x) = \alpha_t E_\phi(x) + \sigma_t \varepsilon$, where $E_\phi(x)$ is an invertible transformation.

   We emphasize that NFDM is not merely an alternative approach, but a generalization of prior diffusion models. With a simple parameterization such as $F_\phi(\varepsilon, t, x) = \alpha_t x + \sigma_t \varepsilon$, NFDM directly corresponds to DDPM and introduces no additional computational overhead. The computational cost of NFDM depends on the complexity of the forward process' parameterization, which lets the user interpolate between standard linear transformations to highly complex nonlinear ones.

   To the best of our knowledge, NFDM provides the most general diffusion modelling framework that allows learning the forward process while keeping the training procedure simulation-free. Besides achieving state-of-the-art log-likelihood modelling results, the flexibility of NFDM enables us to learn generative processes with specific properties like NFDM-OT, for which it is crucial that the forward process has a complex non-linear time-dependent forward processes (see Appendix G.1).

2. To ensure that $F_\phi$ is smooth in time, we chose a parameterization that is continuous and differentiable. These requirements are standard for the neural networks and are not difficult to meet. We discuss the parameterization of $F_\phi$ in more detail in Appendix B.5 and implementation details in Appendix F.

3. The linear parameterization in $\varepsilon$ of $F_\phi$ is as follows: $F_\phi(\varepsilon, t, x) = \mu_\phi(x,t) + \sigma_\phi(x,t) \varepsilon$, where $\mu_\phi$ and $\sigma_\phi$ are learnable functions. In terms of NFDM, we may describe the forward process from [2] as $F_\phi(\varepsilon, t, x) = (1 - t) x + t \varepsilon$. In [3], the authors discuss other variations of linear combinations of $x$ and $\varepsilon$, which can also be easily expressed within the NFDM framework. A key distinction between NFDM and [2, 3] is that the Stochastic Interpolants framework does not propose a method to learn the forward process. Their derivations rely on a fixed forward process, while NFDM allows for learnable parameters in the forward process.

   Indeed, when we parameterize the forward process of NFDM with learnable parameters, we need to learn both the function $F_\phi$ and the predictor $x_\theta$. In Appendix E, we referred to the fact that in [3], with a fixed forward process, two different functions need to be learned for inference, one for the velocity and one for the score. In contrast in NFDM, if we fix the forward process (i.e., fix $F$) like in [3], we only need to learn the predictor $\hat{x}_\theta$.

[1] Kim, Dongjun, et al. "Maximum likelihood training of implicit nonlinear diffusion model.”

[2] Albergo, Michael S., and Eric Vanden-Eijnden. "Building normalizing flows with stochastic interpolants.”

[3] Albergo, Michael S., Nicholas M. Boffi, and Eric Vanden-Eijnden. "Stochastic interpolants: A unifying framework for flows and diffusions.”