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

Authors: grigory  bartosh
dmitry  vetrov
christian a naesseth
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 11:22:04

Abstract: Conventional diffusion models typically relies on a fixed forward process, which implicitly defines complex marginal distributions over latent variables. This can often complicate the reverse process’ task in learning generative trajectories, and results in costly inference for diffusion models. To address these limitations, we introduce Neural Flow Diffusion Models (NFDM), a novel framework that enhances diffusion models by supporting a broader range of forward processes beyond the standard Gaussian. We also propose a novel parameterization technique for learning the forward process. Our framework provides an end-to-end, simulation-free optimization objective, effectively minimizing a variational upper bound on the negative log-likelihood. Experimental results demonstrate NFDM’s strong performance, evidenced by state-of-the-art likelihood estimation. Furthermore, we investigate NFDM’s capacity for learning generative dynamics with specific characteristics, such as deterministic straight lines trajectories, and demonstrate how the framework may be adopted for learning bridges between two distributions. The results underscores NFDM’s versatility and its potential for a wide range of applications.

Main paper: 
 introduction: 
 Diffusion models [56, 20] are a class of generative models constructed by two key processes: the forward process and the reverse process. The forward process gradually corrupts the data distribution, transforming it from its original form to a noised state. The reverse process learns to invert corruptions of the forward process and restore the data distribution. This way, the model learns to generate data from pure noise. Diffusion models have demonstrated remarkable results in various domains [21, 48, 63, 69, 65] . Nevertheless, most existing diffusion models fix the forward process to be predefined, usually linear, Gaussian which makes it unable to adapt to the task at hand or simplify the target for the reverse process. At the same time there is a growing body of work that demonstrates how modifications of the forward process improve performance in terms of generation quality [39, 66, 9] , likelihood estimation [28, 40, 4] or sampling speed [32, 44, 64] .In this paper, we present Neural Flow Diffusion Models (NFDM), a framework that allows for the pre-specification and learning of complex latent variable distributions defined by the forward process. Unlike conventional diffusion models [20] , which rely on a conditional Gaussian forward process, NFDM can accommodate any continuous (and learnable) distribution that can be expressed as an invertible mapping applied to noise. We also derive, and leverage, a new end-to-end simulationfree optimization procedure for NFDM, that minimizes a variational upper bound on the negative log-likelihood (NLL). Furthermore, we propose an efficient neural network-based parameterization for the forward process, enabling it to adapt to the reverse process during training and simplify the learning of the data distribution. To demonstrate NFDM's capabilities with a learnable forward process we provide experimental results on CIFAR-10, ImageNet 32 and 64, attaining state-of-the-art NLL results, which is crucial for many applications such as data compression [22, 71] , anomaly detection [7, 13] and out-of-distribution detection [51, 70] .Then, leveraging the flexibility of NFDM, we demonstrate how this framework can be applied to learn bridges between two distributions using the AFHQ dataset. Finally, we explore training with constraints on the reverse process to learn generative dynamics with specific properties. As a case study, we discuss curvature and obstacle avoidance penalties on the deterministic generative trajectories. Our empirical results indicate improved computational efficiency compared to baselines on CIFAR-10, downsampled ImageNet, and synthetic data.We summarize our contributions as follows:1. We introduce Neural Flow Diffusion Models (NFDM), improving diffusion modelling through a learnable forward process. 2. We develop an end-to-end optimization procedure that minimizes an upper bound on the negative log-likelihood in a simulation-free manner. 3. We demonstrate state-of-the-art log-likelihood results on CIFAR-10, ImageNet 32 and 64. 4. We show how NFDM can be used in learning bridges and generative processes with specific properties, such as dynamics with straight line trajectories, leading to significantly faster sampling speeds and enhanced generation quality with fewer sampling steps. 
background: 
 Diffusion models are generative latent variable models consisting of two processes: the forward and the reverse (or generative) process. The forward process is a dynamic process that takes a data point x ∼ q(x), x ∈ R D , and perturbs it over time by injecting noise. This generates a trajectory of latent variables {z(t)} t∈[0,1] , conditional on the data x, where [0, 1] is a fixed time horizon and z t = z(t) ∈ R D . The (conditional) distribution can be described by an initial distribution q(z 0 |x) and a Stochastic Differential Equation (SDE) with a linear drift term f F (z t , t) : R D × [0, 1] → R D , scalar variance g(t) : [0, 1] → R + , and a standard Wiener process w:EQUATIONDue to the linearity of f F , we can reconstruct the conditional marginal distribution q(z t |x) = N (z t ; α t x, σ 2 t I). Typically, the conditional distributions evolve from some low variance distribution q(z 0 |x) ≈ δ(z 0 -x) to a unit Gaussian q(z 1 |x) ≈ N (z 1 ; 0, I). This forward process is then reversed by starting from the prior z 1 ∼ N (z 1 ; 0, I), and following the reverse SDE [3] : dz t = f B (z t , t)dt + g(t)d w, where f B (z t , t) = f F (z t , t) -g 2 (t)∇ zt log q(z t ).(Here, w denotes a standard Wiener process where time flows backwards. Diffusion models approximate this reverse process by learning ∇ zt log q(z t ), known as the score function, through a λ t -weighted denoising score matching loss:EQUATIONwhere u(t) represents a uniform distribution over the interval [0, 1], and s θ : R D × [0, 1] → R D is a learnable approximation. With a learned score function s θ (z t , t), one can generate a sample from the reverse process by first sampling from the prior z 1 ∼ N (z 1 ; 0, I), and then simulating the reverse SDE, resulting in a sample z 0 ∼ p θ (z 0 ) ≈ q(z 0 ) ≈ q(x):EQUATIONDiffusion models possess several important properties. For example, for a specific λ t , the objective (eq. ( 3)) can be reformulated [59] as an Evidence Lower Bound (ELBO) on the model's likelihood. Furthermore, the minimization of denoising score matching (eq. ( 3)) is a simulation-free procedure. This means that simulating either the forward or reverse processes through its SDE is not necessary for sampling z t , nor is it necessary for estimating the gradient of the loss function. Instead, we can directly sample z t ∼ q(z t |x). The simulation-free nature of this approach is a crucial aspect for efficient optimization.Another notable property is the existence of an Ordinary Differential Equation (ODE) corresponding to the same marginal densities q(z t ) as the SDE (eq. ( 1)):EQUATIONf (z t , t) = f F (z t , t) -g 2 (t) 2 ∇ zt log q φ (z t ).This implies that we can sample from diffusion models deterministically, allowing the use of offthe-shelf numerical ODE solvers for sampling, which may improve the sampling speed compared to stochastic sampling that requires simulating an SDE. Additionally, deterministic sampling enables us to compute densities by treating the model as a continuous normalizing flow, as detailed in [6, 17] . 
neural flow diffusion models: 
 In diffusion models the forward process defines stochastic conditional trajectories {z(t)} t∈[0,1] and the reverse process tries to match the marginal distribution of trajectories. This construction can be viewed as a specific type of hierarchical Variational Autoencoders (VAEs) [30, 45] . However, in conventional diffusion models the latent variables are inferred through a pre-specified linear combination of the data point and Gaussian noise. This formulation limits diffusion models in terms of the flexibility of their latent space, and makes learning of the reverse process more challenging.To address this limitation, we propose a generalized form of the forward process that enables the definition and learning of a broad range of distributions in the latent space. From a practical perspective, a more flexible forward process can simplify the task of learning the reverse process. From a theoretical perspective, learning of the forward process is analogous to learning the variational distribution in a hierarchical VAE, which gives a tighter bound on model's NLL (see an extended discussion in Appendix B.1).In this section, we introduce Neural Flow Diffusion Models (NFDM) -a framework that generalizes conventional diffusion models. The key idea in NFDM is to define the forward process' conditional SDE implicitly via a learnable transformation F φ (ε, t, x) that defines the marginal distributions. This lets the user define a broad range of continuous time-and data-dependent forward processes, that the reverse process will learn to invert. Importantly, NFDM retains crucial properties of conventional diffusion models, like likelihood-based and simulation-free training. Previous diffusion models emerge as special cases when the data transformation is linear, time-independent, and/or additive Gaussian.We approach the forward process constructively. The ultimate goal is a learnable distribution over trajectories, {z(t)} t∈[0,1] given x, realized by a conditional SDE constructed in three steps.First, we (implicitly) define the conditional marginal distributions q φ (z t |x) for t ∈ [0, 1] using F φ (ε, t, x). Then, we introduce the corresponding conditional ODE that together with an initial distribution q φ (z 0 |x) matches the conditional marginal distribution q φ (z t |x). Finally, we define a conditional SDE that defines a distribution over trajectories {z(t)} t∈[0,1] , with marginal distributions q φ (z t |x) for each t.Forward Marginal Distribution. We characterize the marginal distribution q φ (z t |x) of the forward process trough a function that transforms a noise sample ε into z t , conditional on the time step t and data point x:EQUATIONwhereF φ : R D × [0, 1] × R D → R D and ε ∼ q(ε) = N (ε; 0, I).This defines the conditional distribution of latent variables q φ (z t |x). Additionally, eq. ( 7) facilitates direct and efficient sampling from q φ (z t |x) through F φ in eq. ( 7).Conditional ODE. We assume that F φ is differentiable with respect to ε and t, invertible with respect to ε. Further, we assume that fixing specific values of x and ε and varying t from 0 to 1 results in a smooth trajectory from z 0 to z 1 . Differentiating these trajectories over time yields a velocity field corresponding to the conditional distribution q φ (z t |x), thereby defining a conditional ODE:EQUATIONTherefore, if we sample z 0 ∼ q φ (z 0 |x) and solve eq. ( 8) until time t, we have z t ∼ q φ (z t |x).The time derivative of F φ may be calculated efficiently using automatic differentiation tools like PyTorch [42] or JAX [5] (see details in Appendix B.2).Conditional SDE. The function F φ and the distribution of the noise q(ε) together defines q φ (z t |x).To completely define the distribution of trajectories {z(t)} t∈[0,1] , we introduce a conditional SDE that starts from sample z 0 and runs forward in time.With access to both the ODE and score function ∇ zt log q φ (z t |x), an SDE [61] with marginal distributions q φ (z t |x), is:EQUATIONf F φ (z t , t, x) = f φ (z t , t, x) + g 2 φ (t) 2 ∇ zt log q φ (z t |x).Here, g φ : [0, 1] → R + is a scalar function, and w represents a standard Wiener process. Note that g φ only influences the distribution of trajectories {z(t)} t∈[0,1] . The marginal distributions are the same, q φ (z t |x), for any choice of g φ .The SDE in eq. ( 9) requires access to the conditional score function ∇ zt log q φ (z t |x), which in the general case can be computationally costly. However, for F φ that allow efficient evaluation of the log-determinant of the Jacobian matrix, the score function can be calculated efficiently. Examples of such F φ are functions linear with respect to ε or RealNVP style architectures [14, 29] . The calculation of this score function is further discussed in Appendix B.3.To define the reverse (generative) process, we specify a reverse SDE that starts from z 1 ∼ p(z 1 ) and runs backwards in time. To do so we first introduce a conditional reverse SDE that reverses the conditional forward SDE (eq. ( 9)). Following [3] , we define:EQUATIONf B φ (z t , t, x) = f φ (z t , t, x) - g 2 φ (t) 2 ∇ zt log q φ (z t |x).Secondly, leveraging this reverse SDE (see eq. ( 10)) we incorporate a prediction of x:EQUATIONand xθ : R D × [0, 1] → R D is a function that predicts the data point x. This SDE defines the dynamics of the generative trajectories {z(t)} t∈[0,1] .To fully specify the reverse process, it is also necessary to define a prior distribution p(z 1 ) and a reconstruction distribution p(x|z 0 ). In all of our experiments, we set the prior p(z 1 ) to be a unit Gaussian distribution N (z 1 ; 0, I) and let p(x|z 0 ) be a Gaussian distribution with a small variance N (x; z 0 , δ 2 I), where δ 2 = 10 -4 .The above parameterization of the reverse process is not the only possibility. However, it is a convenient choice as it allows for the definition of the reverse process simply through the prediction of x, similar to conventional diffusion models [20] . We leave exploration of alternate parameterizations for future work.Require: q(x), F φ , g φ , xθ for learning iterations do x ∼ q(x), t ∼ u(t)z t ∼ q φ (z t |x) L = 1 2g 2 φ (t) f B φ (z t , t, x) -fθ,φ (z t , t) 2 2Gradient step on θ and φ w.r.t. L end for Algorithm 2 Stochastic Sampling from NFDMRequire: F φ , g φ , xθ , T -number of steps ∆t = 1 T , z 1 ∼ p(z 1 ) for t = 1, . . . , 2 T , 1 T do w ∼ N (0, I) z t-∆t = z t -fθ,φ (z t , t)∆t+g φ (t)w√ ∆t end for x ∼ p(x|z 0 )With F φ in eq. ( 7) parameterized such that q φ (z 0 |x) ≈ δ(x -z 0 ) and q φ (z 1 |x) ≈ p(z 1 ), we propose to optimize the forward and reverse processes of NFDM jointly, minimizing the following objective:EQUATION)Since the forward process is parameterized by φ, we need to optimize the objective (eq. ( 12)) with respect to both φ and θ jointly. We discuss parameterization of the forward process in Appendix B.5.As we demonstrate in Appendix A.1 the objective L shares similarities with the standard diffusion model objective [20] , and it provides a variational bound on the model's log-likelihood log p θ,φ (x).The objective L also exhibits strong connections with both the Flow Matching [33] and Score Matching [68] objectives. We explore these connections in Appendix E. 4 , where we also discuss the role of g φ in detail. It is important to note that despite the parameterization of the reverse process through the prediction of x, the objective L optimizes the generative dynamics and does not necessarily lead to accurate predictions of x.A key characteristic of the NFDM objective is its compatibility with the simulation-free paradigm, which is critical for efficient optimization. We summarize the training procedure in Algorithm 1.To sample from the trained reverse process we can simulate the SDE, as defined in Section 3.2. This procedure is summarized in Algorithm 2. Additionally, during the sampling process, we can adjust the level of stochasticity by modifying g φ (t) (eq. ( 9)). It is important to note that changes to g φ (t) also influence f F φ (eq. ( 9)) and fθ,φ (eq. ( 11)). In the extreme case where g φ (t) ≡ 0, the reverse process becomes deterministic, allowing us to utilize off-the-shelf numerical ODE solvers and to estimate densities [6, 17] . We provide an extended sampling discussion in Appendix B.4. 
neural flow bridge models: 
 In this section we discuss a simple modification to the NFDM framework that enables us to learn bridges between two data distributions, q(x 0 ) and q(x 1 ), a modification we refer to as Neural Flow Bridge Models (NFBM).In the context of bridges, we consider a joint data distribution q(x 0 , x 1 ) (which may be factorized as q(x 0 )q(x 1 ) if paired data is unavailable). Our goal is to learn a generative process that starts from x 1 and generates x 0 such that q(x 0 ) ≈ p θ (x 0 ) = q(x 1 )p θ (x 0 |x 1 )dx 1 . To turn NFDM into NFBM, we modify both the forward and reverse processes. In NFDM, the forward process is defined by two functions: F φ (ε, t, x) (see eq. ( 7)) and g t (see eq. ( 9)). For NFBM, we let F φ depend on both x 0 and x 1 , thereby conditioning the entire forward process on these data points:EQUATIONSimilar to the discussion in Section 3.2, for the reverse process of NFBM, we predict the data points x 0 and x 1 from z t and t, and substitute them into the conditional reverse time SDE (eq. ( 10)):EQUATIONHere, the function x0,1 θ (z t , t) returns predictions of both x 0 and x 1 . Alternatively, the data point x 1 can be reused as conditioning for intermediate steps of the reverse process instead of predicting both points, a strategy further detailed in Appendix C.2.When F φ in eq. ( 13) is parameterized such that q φ (z 0 |x 0 , x 1 ) ≈ δ(x -z 0 ) and q φ (z 1 |x 0 , x 1 ) ≈ p(z 1 |x 1 ) (see Appendix C.1 for parameterization details of the NFBM forward process), we propose training NFBM by minimizing the following objective:EQUATION)The derivation of this objective is provided in Appendix A.2. This objective shares key properties with the NFDM objective, it provides a variational bound on the model's log-likelihood log p θ,φ (x 0 ), and is compatible with the simulation-free paradigm. Furthermore, it allows for sampling with various levels of stochasticity (by adjusting g φ (t)) including deterministic sampling (when g φ (t) ≡ 0).Thus, the NFDM framework not only enables the construction of generative models, but also facilitates learning bridge models between two data distributions. 
restricted nfdm: 
 We introduce NFDM as a powerful framework that enables learning of the forward process. However, there is in general an infinite number of forward and reverse processes that correspond to each other.In this section we discuss how the flexibility of NFDM (and NFBM) allows learning generative dynamic with user-specified beneficial properties.Suppose our objective is to learn straight generative ODE trajectories, which can be highly beneficial, as it enables generation with far fewer steps in the ODE solver. One approach is to introduce penalties on the curvature of the reverse process's trajectories and let the learnable forward process to adapt to align with the reverse process. As a result, we would obtain a generative process that not only corresponds to the forward process, ensuring accurate data generation, but also features the desired property of straight trajectories. We discuss NFDM with restrictions in more details in Appendix D.1.We propose learning a model with curvature penalty as suggested by [25] :EQUATION)We refer to this variant as NFDM-OT1 . L crv is an additional curvature loss that penalizes the second time derivative of the generative ODE trajectories. L is estimated as in Section 3.3, whereas when calculating L crv we set g φ (t) ≡ 0. L crv ≡ 0 ensures that the generative trajectories are straight lines.In our experiments, we set λ = 10 -2 . Empirical evidence suggests that higher values of λ lead to slower convergence of the model. The specifics of the curvature loss are elaborated upon in more detail in Appendix D.2.We would like to emphasise that the purpose of this section is to provide an example of how the NFDM framework allows learning dynamics with specific properties. Not to propose the best way for learning generative straight-line generative dynamics. In addition, we note that conventional diffusion models are incapable of handling such penalization strategies. In diffusion models with a fixed forward process, the target for the reverse process is predetermined and the corresponding ODE trajectories are highly curved. Hence, imposing constraints on the reverse process, such as trajectory straightness, would lead to a mismatch with the forward process and, consequently, an inability to generate samples with high data fidelity. 
experiments: 
 We first showcase results demonstrating that NFDM consistently achieves better likelihood compared to baselines, obtaining state-of-the-art diffusion modeling results on the CIFAR-10 [31] and downsampled ImageNet [11, 67] datasets. Then, we explore the NFDM-OT modification, which penalizes We report the NLL in Bits Per Dimension (BPD) and sample quality measured by the Frechet Inception Distance (FID) score [19] . The NLL is calculated by integrating the ODEs using the RK45 solver [15] , with all NLL metrics computed on test data. For FID, we provide the average over 50k generated images.Unless otherwise stated, we parameterized the forward process of NFDM and NFBM such that q φ (z t |x) is a Gaussian with learnable mean and covariance (see details in Appendices B.5 and C.1). For a detailed description of parameterizations and other experimental details, please refer to Appendix F.The primary aim of the experiments with NFDM-OT and NFBM is not to introduce a novel model that surpasses others in few-step generation or learning bridges, but rather to showcase the ability of the NFDM framework to learn generative dynamics with specific properties. The straightness of the trajectories is just one example of such a property. We leave it for future research to explore different parameterizations and modified objectives for NFDM that may yield even better results.Please refer to appendix G.3 for generated samples. The code is available at https://github.com/ GrigoryBartosh/neural_diffusion.For the first experiments, in addition to Gaussian we also provide results with non-Gaussian parameterization of the forward process (see Appendix B.5). Table 1 summarizes the NLL results on the CIFAR-10 and two downsampled ImageNet datasets. Notably, NFDM outperforms diffusion-based baselines on all three datasets with both parameterizations, achieving state-of-the-art performance.The improved performance due to NFDM is a natural progression for the following reasons. First, it is well-established [59, 76] that diffusion models exhibit improved likelihood estimation when trained with the full ELBO objective. The objective L (eq. ( 12)) used for training the NFDM is also a variational bound on the likelihood.Second, diffusion models can be seen as hierarchical VAEs. From this perspective, most baselines in Table 1 resemble VAEs with either fixed or constrained variational distributions. In contrast, the NFDM extends beyond these baselines by providing a more flexible variational distribution. That is true even for the Gaussian parameterization with a mean and covariance that is non-linear in x and t. This flexibility allows the NFDM to better conform to the reverse process, consequently enhancing likelihood estimation.We next evaluate NFDM-OT, designed to penalize the deterministic generative trajectory curvature. First, we compare NFDM-OT with a conventional continuous-time diffusion model. Figure 1a illustrates deterministic trajectories between a two-dimensional data distribution and a unit Gaussian distribution learnt by a conventional diffusion model [61] . Figure 1b depicts trajectories learnt by NFDM-OT. Conventional diffusion, being constrained in its forward process, learns highly curved trajectories, whereas NFDM-OT successfully learns straight generative trajectories as desired. In Appendix G.1 we provide some additional reult, demonstrating the importance of learnable forward process for NFDM-OT.Then, we present results of NFDM-OT on image datasets. Table 2 reports the FID scores for 2, 4, and 12 Number of Function Evaluations (NFE) with respect to the function fθ,φ (eq. ( 11)). In this experiment, we employ Euler's method for sampling integration. For the specified NFEs, NFDM-OT demonstrates superior sample quality compared to other approaches with similar NFE values. Specifically, NFDM-OT outperforms approaches specifically designed to minimize the curvature of generative trajectories [44, 32] .Importantly, NFDM-OT is trained with an ELBO-based objective (eq. ( 12)), which is known to yield higher FID scores for diffusion models [59, 76] . In contrast, some of the approaches listed in Table 2 are trained with different objectives, leading to improved FID scores. Even so, for comparable NFE values NFDM-OT still achieves superior results. We provide additional results of NFDM-OT in Appendix G.2.Figure 2 : Trajectories of NFBM between two-dimensional data distributions, trained with an additional penalty to avoid obstacles. We consider the task of translating images of dogs into images of cats from the downsampled AFHQ dataset. Table 3 reports the FID scores on images sampled with 1000 steps. Notably, NFBM outperforms the baselines, demonstrating the effectiveness of our framework for learning bridges.Finally, we demonstrate that the framework flexibility also allows for learning dynamics with specific properties in the case of NFBM.To illustrate this, we adopted the strategy discussed in Section 5 and learn NFBM with an additional penalty to avoid obstacles (see details in Appendix F.2). Figure 2 visualizes the learned stochastic trajectories. It is evident that NFBM has efficiently learned to translate distributions while avoiding obstacles. 
related work: 
 Diffusion models, originally proposed by [56] , have evolved significantly through subsequent developments [60, 20] . These advancements have resulted in remarkable generative quality for highdimensional data distributions [12, 48] . Nevertheless, conventional diffusion models, typically relying on a linear Gaussian forward process, may not optimally fit some data distributions. To address this, alternative forward processes have been explored, such as alternative linear transformations [16, 55] , combining blurring with Gaussian noise injection [46, 9, 23] , diffusion in the wavelet spectrum [43] , and forward processes based on the exponential family [41] . These models are limited by their fixed forward processes and may be seen as specific instances of NFDM.Some studies have focused on making the forward process learnable. Approaches include a learnable noise injection schedule [28, 39, 49] and learning data transformations like time-independent transformations based on VAEs [66, 47] and normalizing flows [26] , or time-dependent transformations [18, 55, 40, 4] . These methods can also be considered special cases of NFDM with specific transformations F φ (eq. ( 7)).To improve the sampling efficiency, in orthogonal line of works alternative sampling methods have been studied [62, 35, 38, 53] . Additionally, distillation techniques have been applied [50, 58, 75, 72] to enhance sampling speed, albeit at the cost of training additional models. Most of these approches are compatible with NFDM and may be combined for further gains.Another line of works proposed technics for learning straighter generative trajectories. However, these approaches requires distillation [37, 36] , solving unstable min-max problem [1] , estimating an optimal coupling over minibatches [44, 64] of the entire dataset, which, for large datasets, may become uninformative, or can be viewed as specific instances [32, 52] of NFDM.The connections between NFDM and some related works is discussed further in Appendix E. 
conclusion and limitations: 
 In this paper we introduced NFDM, a novel simulation-free framework for improved diffusion modeling through a learnable forward processes. NFDM outperforms baseline diffusion models on standard benchmarks, showcasing its effectiveness. However, NFDM is not just another model, but rather it is a versatile framework that facilitates the predefined specification and learning of a forward process. Similarly, the penalty on the curvature of deterministic trajectories discussed in Section 5 represents just one specific example of the restrictions that NFDM can accommodate.Nonetheless, the advantages of NFDM come with certain trade-offs. Once the forward process is parameterized using a neural network, this leads to approximately 2.2 times longer optimization iteration of NFDM compared to conventional diffusion models. Additionally, as discussed in Section 3.1, NFDM is limited to forward process parameterisations with functions that match specific conditions.Despite these challenges, we are optimistic about the potential of NFDM. NFDM opens up significant prospects for exploring new generative dynamics. We believe that with alternative parameterizations, modifications of the objective, and the integration of orthogonal approaches, NFDM has the potential to achieve even better results across a range of data modalities.Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We provide compute resources details in Appendix F.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: We don't believe this work has any societal impact.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: The paper does not pose such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: All data and software used is properly credited and licenses are respected.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: Our goal is to derive a variational bound on the negative log-likelihood of NFDM. For the convenience of the following derivations we will consider the forward process as a conditional reverse SDE (eq. ( 10)) that starts from z 1 ∼ q φ (z 1 |x) and flows backwards in time:EQUATIONSimilarly to a conditional SDE (eq. ( 9)) flowing forwards in time, this SDE also corresponds to conditional marginal distribution q φ (z t |x), which is implicitly defined by the invertible transformation z t = F φ (ε, t, x) (eq. ( 7)), where ε ∼ q(ε).For the reverse process we use the definition from Section 3.2 as reverse SDE (eq. ( 11)) that starts from the prior distribution p(z 1 ), flows backwards in time, and finally samples x ∼ p(x|z 0 ):EQUATION)To derive the NFDM's variational bound on the NLL we discretize processes, derive an objective in discrete case and then consider the limit of discretization to switch back to continuous time.We discretize the conditional reverse SDE in eq. ( 17) and the reverse SDE in eq. ( 18), transitioning from continuous-time trajectories {z(t)} t∈[0,1] to discrete-time trajectories z0 , z 1 T , . . . , z1 , where T represents the number of discrete steps of size ∆t = 1 T . The discrete conditional reverse process starts from distribution qφ (z 1 |x) = q φ (z 1 |x) and follows conditional reverse Markov chain:EQUATIONSimilarly, the discrete generative process shares distributions p(z 1 ) = p(z 1 ) and p(x|z 0 ) = p(x|z 0 ) with the continuous one. However, instead of following the reverse SDE (eq. ( 18)) it follows the reverse Markov chain:EQUATIONIt's important to note that once we discretize the trajectories, qφ (z t |x) ̸ = q φ (z t |x) and pθ,φ (z t ) ̸ = p θ,φ (z t ) except for t = 1. However, this discretization corresponds to an Euler-Maruyama method of integrating SDEs. Thus, we know that as ∆t → 0, the discretized trajectories converge to the continuous ones. Now, considering these discrete processes and recognizing that they are Markovian conditional and unconditional processes moving backwards in time, we can leverage a key result from [56, 20] :EQUATIONEQUATIONEQUATIONEquation ( 21) establishes a variational upper bound on the negative log-likelihood of discretized reverse process. Taking the limits of this upper bound, the reconstruction term L rec = Lrec and the prior term L prior = Lprior retain their forms. However, we can rewrite the diffusion term:EQUATION(28)EQUATIONwhere u(t) is a uniform distribution over a unit interval.Therefore, as the discretized trajectory tends to the continuous one when ∆t → 0, we can use the obtained limits as a variational bound on the negative log-likelihood of the continuous model:EQUATIONIn general case, when function F φ is actually parameterized by φ, the objective L implies that we must optimise all three terms with respect to parameters φ and θ. However, when F φ is parameterized such that q φ (z 0 |x) ≈ δ(x -z 0 ) and q φ (z 1 |x) ≈ p(z 1 ) for any φ, the reconstruction L rec and prior L prior terms are constants with respect to the parameters.Similarly to the derivation of the NFDM objective in Appendix A.1, here we also apply discretization to derive the objective of NFBM.The conditional reverse process (eq. ( 17)) and it's discretization (eq. ( 19)) stays the same, except now the conditional reverse process are conditioned on a pair of points x = (x 0 , x 1 ), sampled from a joined data distribution (x 0 , x 1 ) ∼ q(x 0 , x 1 ). The reverse process (eq. ( 18)) and its discretization (eq. ( 20)) also have the same structure. However, for NFBM the reverse process starts from distribution p(z 1 |x 1 ) = p(z 1 |x 1 ) and at the end samples from p(x 0 |z 0 ) = p(x 0 |z 0 ).We derive the variational bound on NLL of NFBM:EQUATION=E q(x0) -log p(x 0 |z 0 )p θ,φ (z 0:1-∆t |z 1 )p(z 1 |x 1 )q(x 1 )dz 0:1 dx 1 (32)EQUATIONEQUATION= Lrec + Ldiff + Lprior + const, where (35)EQUATIONEQUATIONEQUATIONconst = E q(x0) D KL q(x 1 |x 0 )∥q(x 1 ) .Here, the reconstruction term L rec = Lrec and the prior term L prior = Lprior do not depend on discretization, and the constant term does not depend on model's parameters. For the diffusion termLdiff it is easy to see that in the case of NFBM it has the same structure as in NFDM. Therefore, taking the limit of the number of discretization steps T the diffusion term Ldiff transforms to:EQUATIONwhere u(t) is a uniform distribution over a unit interval.This gives us a variational bound on the negative log-likelihood of the NFBM, with a very similar structure as the NFDM objective. Additionally, for NFBM we also may omit optimization of L rec and prior L prior when q φ (z 0 |x 0 , x 1 ) ≈ δ(x -z 0 ) and q φ (z 1 |x 0 , x 1 ) ≈ p(z 1 |x 1 ).In this section, we provide an extended discussion on the motivation behind the learnable forward process.From a theoretical perspective, diffusion models can be viewed as a hierarchical VAE, where the objective function of diffusion models is the variational bound on model log-likelihood. Therefore, diffusion models aim to match the joint distributions of data x and latent variables z 0:1 (for simplicity, let's consider a discrete-time model with some T time steps). However, as we know from [30, 45] , the variational objective is equal to:D KL q φ (x, z 0:1 )∥p θ (x, z 0:1 ) objective = D KL q(x)∥p θ (x)entropy + likelihoodEQUATIONThus, by minimizing the variational objective in diffusion models, we address two aspects: the likelihood of the data and the posterior error. From this perspective, it is quite a natural choice to also optimize the variational distribution q φ (z 0:1 |x) to match the model's posterior distribution p θ (z 0:1 |x).In other words, by making the forward process learnable, we optimize not only the reverse process to match the forward process but also allow the forward process to match the reverse process. This approach simplifies the task of matching the data distribution by the reverse process.As discussed in Section 3.1, the definition of forward process requires access to the time derivative of the function F φ . This derivative is crucial for constructing a conditional ODE (see eq. ( 8)):EQUATIONFirst, let us clarify this notation: to find the ODE drift term f φ at point z t and time t conditionally on x, we first apply the inverse function F -1 φ to identify ε, which corresponds to z t , t, and x. Subsequently, we compute the partial derivative of F φ with respect to t at ε and x.For certain parameterizations of F φ , the time derivative can be determined analytically. For instance, F φ associated with conventional diffusion models (see Appendix E.1) have time derivatives available in a closed form. However, in the general case, the time derivative is not readily available.Fortunately, we can employ automatic differentiation tools like PyTorch [42] or JAX [5] . The naive approach to determining the time derivative involves running a backpropagation procedure for each output of F φ and reconstructing each coordinate of the time derivative. Unfortunately, this method is not scalable for high-dimensional F φ .Instead of backpropagation, which corresponds to a Vector Jacobian Product (VJP), we can use a forward differentiation mode, or Jacobian Vector Product (JVP). By fixing the inputs ε and x of the function F φ and focusing solely on the input t, the Jacobian of F φ appears as a vector. Consequently, computing a JVP with a unit one-dimensional vector precisely produces the time derivative of F φ . This approach allows for scalable computation of time derivatives.In Section 3.1, we define the forward process as a conditional SDE (eq. ( 9)):EQUATIONf F φ (z t , t, x) = f φ (z t , t, x) + g 2 φ (t) 2 ∇ zt log q φ (z t |x).To construct this SDE, access to the conditional score function ∇ zt log q φ (z t |x) is required. This section discusses methods for calculating this score function.First, consider the case where F φ is linear in ε (further details on this parameterization are discussed in Appendix B.5):EQUATIONwhereµ φ : R D × [0, 1] → R D and σ φ : R d × [0, 1] → R D, with the product εσ φ (x, t) being element-wise. Consequently, the marginal probability densities of the forward process are Gaussian:EQUATIONIn the case of Gaussian distributions, the score function can be easily found as:EQUATION)Now, consider the more general case where the distribution q φ (z t |x) is defined implicitly through an invertible transformation F φ of the random variable ε ∼ q(ε) into z t . To find the score function, we apply the change of variable formula:EQUATIONThe first part of the equation is the log-density of the noise distribution, which is straightforward to calculate. The second part is the log-determinant of the Jacobian matrix of the inverse transformation, log |J -1 F |. If we have access to log |J -1 F |, we can calculate log qφ(z t |x) using the change of variable formula and subsequently compute the gradient using automatic differentiation tools such as PyTorch [42] or JAX [5] .The log-determinant log |J -1 F | can be easily calculated when F φ is linear in ε, a case already considered above. For nonlinear, low-dimensional F φ , the Jacobian J -1 F can be computed using O(D) backpropagations, and the log-determinant subsequently calculated. However, this approach is not scalable for high-dimensional cases. It is challenging to compute the log-determinant log |J -1 F | in general high-dimensional cases, posing a limitation to the NFDM framework as it restricts the parameterization of the forward process to functions with accessible log-determinants.Nevertheless, the forward process can be parameterized with functions F φ that inherently provide access to the log-determinant log |J -1 F | by design. For example, normalizing flow architectures [14, 29] facilitate the construction of forward processes with non-Gaussian, learnable distributions q φ (z t |x).In Section 3.2, we define the reverse process as the following SDE (see eq. ( 11)):EQUATIONLet us now rewrite this equation, substituting the definition of f B φ from eq. ( 10):EQUATION)This expression represents the full form of the generative SDE, for which we derive the objective in Appendix A.1. After training, to sample data points from the models, we can simulate this SDE as outlined in eq. ( 49) and depicted in Algorithm 2.However, as evident from eq. ( 49), the generative process has a very simple dependence on g φ (t), which determines the level of stochasticity. Therefore, in practice, we can adjust the stochasticity level by modifying the function g φ (t).In the extreme case where g φ (t) ≡ 0, the generative SDE completely loses its stochasticity and becomes an ODE:EQUATIONThis ODE allows for an interpretation of the model as a Continuous Normalizing Flow (CNF) [6, 17] , consequently enabling deterministic sampling from a pre-trained model and likelihood estimation.The parameterization of the variance function g φ (t) (eq. ( 9)) and the data point predictor xθ (z t , t) (eq. ( 11)) is straightforward. In our experiments, we parameterize them directly using neural networks. However, parameterizing the transformation F φ (eq. ( 7)) is more complex. NFDM requires F φ to be differentiable and invertible with respect to ε, as well as provide access to the log-determinant of the Jacobian matrix (see Appendix B.3). We propose two parameterizations for the forward process of NFDM.First, we propose a parameterization of F φ that is linear in ε:EQUATIONEQUATIONEQUATIONHere, μφ : R D ×[0, 1] → R D and σφ : R D ×[0, 1] → R D + , with δ as a constant. This parameterization models q(z t |x) as a conditional Gaussian distribution:EQUATIONIn our experiments, we set δ 2 = 10 -4 . If μφ is differentiable and σφ is differentiable and nonnegative, this parameterization meets the requirements on F φ .The functions µ φ (x, t) and σ φ (x, t) are specifically designed such that at t = 0, they equal µ φ (x, 0) = x and σ φ (x, t) = δ, and at t = 1, they become µ φ (x, 1) = 0 and σ φ (x, t) = I. At intermediate times 0 < t < 1, these functions can take arbitrary values, thereby ensuring that q(z 0 |x) = N (z 0 ; x, δ 2 I) and q(z 1 |x) = N (z 1 ; 0, I), with arbitrary Gaussian distributions at intermediate steps. These restrictions are not necessary for F φ , however, it eliminates the need to optimize the reconstruction loss L rec and the prior loss L prior , as they become independent of the parameters φ and θ (see Appendix A.1).Secondly, we propose a more general parameterization:EQUATIONwhere Fφ : R D × [0, 1] × R D → R D is an invertible and differentiable function that allows access to the log-determinant of the Jacobian matrix. If Fφ is sufficiently flexible (e.g., a neural network with a normalizing flow architecture), it can parameterize non-Gaussian distributions q φ (z t |x).The additive term (1 -t)x provides an inductive bias that shifts the mean of the distribution q φ (z t |x) towards x when t is near 0. The linear function δ + (1 -δ)t scales the distribution from a factor of δ at t = 0 to 1 at t = 1, also imparting an inductive bias for increasing variance over time. The scaling (1 -t)x in the third argument of Fφ ensures that q φ (z 1 |x) at t = 1 does not depend on x, allowing the reverse process to start from p(z 1 ) = q φ (z 1 |0) and thereby guaranteeing matching of prior distributions between the forward and reverse processes. Consequently, this eliminates the need for optimizing the prior loss L prior . However, unlike the previous parameterization that guarantees q(z 0 |x) = N (z 0 ; x, δ 2 I), in this case, there are no guarantees on the form of q(z 0 |x), necessitating the optimization of the reconstruction loss L rec with respect to parameters φ (see Appendix A.1).Importantly, even the Gaussian parameterization of F φ in eq. ( 51) offers a significantly more flexible forward process compared to prior works. Although linear in ε, it may exhibit complex non-linear dependencies on x and t.Furthermore, we emphasize that the parameterizations of F φ , g φ (t), and xθ (z t , t) discussed here are not the only possible ones and may be sub-optimal. For the purposes of this paper, we maintain a simple setup and leave the exploration of more effective parameterizations for future research.To parameterize the forward process of NFBM, we adopt the Gaussian parameterization discussed for NFDM in Appendix B.5:EQUATIONEQUATIONEQUATIONHere, μφ :R D × R D × [0, 1] → R D and σφ : R D × R D × [0, 1] → R D +, with δ 2 = 10 -4 in our experiments. Similar to the NFDM, this parameterization models q(z t |x) as a conditional Gaussian distribution:EQUATION)In contrast to the parameterization of NFDM, here µ φ (x, t) and σ φ (x, t) are designed to ensure that q(z 0 |x 0 , x 1 ) = N (z 0 ; x 0 , δ 2 I) and q(z 1 |x 0 , x 1 ) = N (z 1 ; x 1 , δ 2 I). At intermediate times, 0 < t < 1, q(z t |x 0 , x 1 ) may still be an arbitrary Gaussian distribution. Therefore, with this parameterization of NFBM, we do not need to optimize the reconstruction loss L rec and the prior loss L prior , as they are independent of the parameters φ and θ (see Appendix A.2).As with NFDM, the parameterization of NFBM that we provide in this section represents a simple design choice. We believe there are more efficient ways to parameterize the forward process of NFBM. However, in this paper, we want to focus on presenting the framework itself, rather than dwelling on design choices.In Section 4, we define the reverse process of NFBM by the following SDE:EQUATIONHere, the function x0,1 θ (z t , t) predicts both x 0 and x 1 . Assuming q φ (z 0 |x 0 , x 1 ) ≈ δ(x -z 0 ) and q φ (z 1 |x 0 , x 1 ) ≈ p(z 1 |x 1 ), this formulation leads to the objective (see Appendix C.1):EQUATION)As previously mentioned, alternatively, we could define the reverse process as being fully conditional on the initial data point x 1 . This leads to the following SDE:EQUATION)This approach is akin to classifier-free guidance, where each step of the generative process depends on conditioning x 1 . Under the same assumptions-that q φ (z 0 |x 0 , x 1 ) ≈ δ(x-z 0 ) and q φ (z 1 |x 0 , x 1 ) ≈ p(z 1 |x 1 )-this definition leads to a slightly different objective (the derivation is not provided here but can be easily inferred by following the steps described in Appendix A.2 for the previous definition):EQUATION)Intuitively, the difference between these two definitions implies that the first reverse process, at a point z t at time t, follows the average trajectory of forward processes passing through z t . Conversely, the second reverse process tracks the average trajectory of forward processes that pass through z t and terminate at z 1 .The second approach may be beneficial when dealing with joint data distributions q(x 0 , x 1 ). However, in this paper, as we work with unpaired distributions q(x 0 )q(x 1 ), we have opted for the first approach in Section 4.In Section 5, we discussed training the NFDM with restrictions on the reverse process. This section aims to further explore the training under restrictions and to examine the limitations of this technique.Typically, there is an infinite number of forward and reverse processes that correspond to each other. So how to choose between them? The standard NFDM, as described in Section 3, learns just one such pair of processes. However, the flexibility of NFDM opens up the possibility of choosing specific pairs of processes depending on the task.Suppose we aim to learn a reverse process with some specific properties. To achieve this, we must impose certain restrictions on the reverse process. One approach is to encode these restrictions directly in the parameterization of the reverse process. Alternatively, we can introduce penalties on the reverse process that would penalize deviations from the desired properties. In both cases, we expect the forward process to adapt and align with the reverse process. Consequently, the reverse process will match the forward process and, as a result, the data distribution, while possessing the desired properties.In Section 5, we discussed the restriction of the reverse process to have straight deterministic generative trajectories. However, we may apply the same strategy to ensure arbitrary properties. We believe that the NFDM framework may help construct various new generative dynamics for a variety of tasks and domains.Nevertheless, training the NFDM with restrictions has limitations. The primary concern is the feasibility of these restrictions. For instance, imposing overly stringent constraints on the reverse process (such as fixing its drift term f ≡ 0) will render the forward process incapable of adapting, regardless of its flexibility. Similarly, unattainable penalties, such as those on the squared norm of the drift term ∥ fθ,φ ∥ 2 2 , will lead to biased solutions, as it is impossible to match two processes when ∥ fθ,φ ∥ 2 2 ≡ 0. Additionally, even when a restriction is feasible, if the model parameterization is too simplistic, it may not be able to match processes and satisfy the restriction, leading to bias. However, in our experiments, we find that even the Gaussian parameterization of the forward process (see Appendix B.5) is flexible enough in practice to successfully learn various restrictions, such as straight generative trajectories.In Section 5, we suggest penalizing the curvature of deterministic generative trajectories specifically with the L crv (eq. ( 16)) penalty.To calculate the curvature penalty L crv (eq. ( 16)), we proceed as follows. First, we use the chain rule to rewrite the time derivative:EQUATION)The second term in eq. ( 65) is the time derivative of a function. As discussed in appendix B.2, the time derivatives can be determined as a Jacobian Vector Product (JVP) with a one-dimensional unit vector.Notably, the first term in eq. ( 65) is also a JVP. Therefore, we can combine these two operations. For this purpose, we define a R D+1 dimensional adjoint vector v:EQUATIONConsequently, L crv can be efficiently computed as the JVP of fθ,φ (z t , t) with the vector v.Additionally, we would like to note a couple of aspects of learning the NFDM with a curvature penalty. First, unlike the regular NFDM, the curvature penalty necessitates that F φ be twice differentiable with respect to time. Second, although we parameterize fθ,φ (eq. ( 11)) through f B φ (eq. ( 10)), L crv influences the forward process but does not penalize the curvature of the conditional forward trajectories (see illustrations in Appendix G.1).In this section we continue discussing connections between NFDM and related works.Many existing works define the forward process in diffusion models as either fixed or simply parameterized processes. These processes can be considered special cases of NFDM. In this section, we review some of these approaches. Soft Diffusion. [9] consider the case of a fixed forward processes q(z t |x) = N (z t ; C t x, σ t I), which can be reparameterized as F (ε, t, x) = C t x + σ t ε. Such distributions include, for example, combinations of blurring and the injection of Gaussian noise.Star-Shaped Diffusion. [41] extended conventional diffusion models to include distributions from the exponential family. Although Star-Shaped Diffusion is a discrete-time approach and does not directly correspond to NFDM, the latter can work with exponential family distributions through reparameterization functions. For instance, for some continuous one-dimensional distribution q(z t |x), NFDM could use F (ε, t, x) = a(b(ε), t, x), where a is the inverse Cumulative Distribution Function (CDF) of q(z t |x), and b is the CDF of a unit Gaussian.Variational Diffusion Models. [28] proposed forward conditional distributions q φ (z t |x) as N (z t ; α φ (t)x, σ 2 φ (t)I) with learnable parameters φ. In the context of NFDM, this can be realized by F φ (ε, t, x) = α φ (t)x + σ φ (t)ε.LSGM. [66] suggest an alternative approach for parameterizing the forward process, proposing diffusion in a latent space of a VAE. Therefore, the forward process is characterized by a distributionq φ (z t |x) = N (z t ; α t E φ (x), σ 2 t I), where E φ is the encoder of the VAE. To parameterize the same forward process with NFDM, one could use F φ (ε, t, x) = α t E φ (x) + σ t ε. To align the reverse process, the reconstruction distribution should be p(x|z 0 ) = N (x; D φ (z 0 ), δ 2 I), where D φ is VAE's decoder.ShiftDDPM. [74] proposed learning a function that shifts the mean and covariance of q ϕ (z t |x) by conditional information. In our work, we do not focus on conditional generation. However, we believe that the NFDM framework can be readily adapted for conditional generation by incorporating the conditioning variable c into both the transformation F ϕ (ε, t, x, c) and the predictor x(z t , t, c). NDM and DiffEnc. [4, 40] proposed a more general forward process, q φ (z t |x) = N (z t ; α t f φ (x, t), σ 2 t I), which, unlike LSGM, transforms x in a time-dependent manner. This forward process can also be described in terms of NFDM as the special caseF φ (ε, t, x) = α t f φ (x, t) + σ t ε.Curvature minimization. [32] suggested learning the distribution q φ (z 1 |x) of the forward process at time step t = 1 as a conditional Gaussian. They parameterize the forward process via linear interpolation between x and z 1 . Within the NFDM framework, we can express this as F φ (ε, t, x) = (1 -t)x + t µ φ (x) + σ φ (x)ε . Subsequently, the authors train their model by optimizing the Flow Matching objective in conjunction with minimizing the Kullback-Leibler (KL) divergence between q φ (z 1 |x) and p(z). This objective corresponds to minimization of diffusion term L diff and prior term L prior of NFDM objective (see Appendix A.1) for the forward process proposed by [32] .Despite the simulation-free nature of the training procedure, diffusion models still necessitate full reverse process simulations for sample generation, leading to slow and computationally expensive inference. To address this, in an orthogonal line of works alternative sampling methods have been studied, such as deterministic sampling in discrete [57] and continuous time [61] and novel numerical solvers [62, 35, 38, 53] . Additionally, distillation techniques have been applied to both discrete [50] and continuous time models [58, 75, 72] to enhance sampling speed, albeit at the cost of training additional models. Most of these approaches are compatible with NFDM and may be adapted for further gains.The significance of deterministic trajectories in efficient sampling was highlighted by [24] . At the same time, building on diffusion model concepts, [33, 37, 2] introduced simulation-free methods for learning deterministic generative dynamics. Based on these ideas, [37] and [36] proposed distillation procedures to straighten deterministic generative trajectories. Since these methods are based on distillation, we consider them orthogonal to NFDM.Schrödinger Bridge models. Recent studies [73, 10, 54] have explored generative models based on Schrödinger Bridge theory and finite-time diffusion constructions. Unlike NFDM, these models deviate from the simulation-free paradigm by necessitating the full simulation of stochastic processes for training, which significantly increases their computational cost. Moreover, these studies typically aim to address the optimal stochastic control problem, which is outside the scope of our work.Nevertheless, it is important to highlight that within this research domain, some authors have proposed advanced techniques for parameterizing learnable conditional processes. For instance, in [34] , the authors introduced the use of splines to parameterize the conditional process, thereby constructing a more flexible and efficient model. This parameterization technique can be seen as a specific instantiation of the parameterization procedure proposed by the NFDM framework.Minibatch optimal transport. [44] and [64] proposed to construct the forward process with optimal data-noise couplings to learn straighter generative trajectories. While this method is theoretically justified, it relies on estimating an optimal coupling over minibatches of the entire dataset, which, for large datasets, may become uninformative as to the true optimal transport coupling. In contrast, NFDM-OT directly penalizes curvature of the generative trajectories.Stochastic Interpolants. In this section, we explore the connections between NFDM and Stochastic Interpolants, as proposed by [1] . Both NFDM and Stochastic Interpolants introduce a more general family of forward processes through the use of a reparameterization function. Although we acknowledge the relevance of Stochastic Interpolants to our work, there are notable differences between the two approaches. First, our methodology involves parameterizing the reverse process by substituting the prediction of x into the forward process, whereas Stochastic Interpolants necessitate learning two separate functions for the reverse process: the velocity field and the score function.Second, we present NFDM as a framework that enables both the pre-specification and learning of the forward process, in contrast to Stochastic Interpolants, which are derived under the assumption of a fixed forward process. Consequently, the objectives utilized by Stochastic Interpolants do not support the incorporation of learnable parameters for the forward process. Furthermore, these objectives are tailored towards learning the velocity field and the score function, rather than optimizing likelihood, as is the case with NFDM.In their practical applications, Stochastic Interpolants are demonstrated with a simple parameterizations of the forward process. They propose a theoretical framework for learning the forward process that would result in an reverse process characterized by dynamical optimal transport. However, this approach is contingent upon solving a high-dimensional min-max optimization problem, for which they do not provide experimental results. Moreover, their work does not clearly articulate how Stochastic Interpolants might be applied to learning other generative dynamics.In contrast, NFDM introduces a more generic method for learning generative dynamics. Moreover, when NFDM is learned with restrictions or penalties (see Section 5), it remains within a min-min optimization paradigm, which can be addressed more efficiently.In this section, we delve into the details of the NFDM's objective function (eq. ( 12)). Specifically we will consider the diffusion term of the objective L diff (see Appendix A.1).First, let's unpack the L diff by substituting the definitions of f B φ and fθ,φ :EQUATION(67)EQUATION)g 2 φ (t) 2 ∇ zt log q φ z t |x θ (z t , t) -∇ zt log q φ (z t |x) Score Matching term 2 2     .This formulation clearly delineates two components of the objective: the first calculates the difference between the ODE drift terms, and the second calculates the difference between the score functions. Moreover, this expression highlights the role of g φ (t). When g φ (t) is small, the forward and reverse processes exhibit smoother trajectories, and the objective is dominated by the first term. Conversely, when g φ (t) is large, the processes exhibit more stochastic trajectories, and the objective is dominated by the second term.Crucially, in the extreme scenarios where g φ (t) approaches either 0 or ∞, the diffusion loss L diff corresponds to either a reweighted Flow Matching loss [33] or a reweighted Score Matching loss [68] , respectively:EQUATIONEQUATIONThis connection highlights the importance of optimizing the full objective in eq. ( 67), which contains both Flow Matching and Score Matching components. Simplifying the objective to include only one component renders it intractable from a variational perspective. Specifically, when g ϕ (t) approaches either 0 or ∞, the objective, a variational bound on likelihood, becomes infinitely large.There are two primary reasons why many popular diffusion models are trained using just one of these components.Firstly, many diffusion models employ a simple linear parameterization of the forward process. For these models, both the conditional vector field f φ (z t , t, x) and the conditional score function ∇ zt log q φ (z t |x) are merely linear combinations of z t and x. As [27] demonstrated, for such models, Flow Matching and Score Matching can be viewed as reweighted ELBO. However, for more general non-linear forward process in NFDM, optimizing Flow Matching and Score Matching objectives is not equivalent.Secondly, in other models, the Flow Matching and Score Matching objectives are derived under the assumption that the forward process is fixed. Without this assumption, we cannot derive eitherThe models were trained using the Adam optimizer with the following parameters: β 1 = 0.9, β 2 = 0.999, a weight decay of 0.0, and ϵ = 10 -8 . The training process was facilitated by a polynomial decay learning rate schedule, which includes a warm-up phase for a predefined number of training steps. During this phase, the learning rate is linearly increased from 10 -8 to a peak value.After reaching the peak learning rate, it is then linearly decreased to 10 -8 by the final training step. The specific hyperparameters are detailed in Table 4 . Training was carried out on Tesla V100 GPUs.For reporting the NLL, we employ standard uniform dequantization. The NLL is estimated using an importance-weighted average, given byEQUATIONwhere x ∈ [0, . . . , 255].In Section 6.3, we demonstrate how the concept of learning NFDM with the restrictions discussed in Section 5 can be adapted to NFBM for learning generative dynamics that avoid obstacles. For this experiment, we propose training the NFBM models with an additional penalty that has higher values at points closer to obstacles. Specifically, we applied the following objective:EQUATIONwhere f (z t ) is a penalty function. In our experiments, we use the probability density of a mixture of Gaussian distributions, truncated at some low values to 0. This truncation is necessary to ensure that the model does not push trajectories infinitely away from the obstacle.In contrast to L crv proposed for NFDM-OT, which penalizes the reverse process influencing the forward process as a result, L obs directly penalizes the forward process, thereby influencing the properties of the reverse process.In this section, we provide visualizations of the forward trajectories of NFDM-OT, as illustrated in the two-dimensional data distribution experiment shown in Figure 1b . Figure 3 depicts the deterministic trajectories of NFDM-OT. It is easy to see that the forward process has learned some highly curved trajectories. This example highlights several key aspects of NFDM-OT and NFDM more broadly.Firstly, this example underscores the point discussed in Appendix D.2: while NFDM-OT penalizes the curvature of the reverse process, which is parameterized through the forward process (see Section 3.2), it does not penalize the curvature of the forward process itself.Secondly, this example demonstrates the importance of learning nonlinear forward processes. In attempting to adapt to the straight-line reverse process, the forward process ends up being highly curved, which in turn aids the reverse process. This emphasizes the importance of the NFDM framework, which facilitates the learning of such nonlinear forward processes that may help to learn better generative dynamics.Finally, we note that we parameterize the forward process with a Gaussian distributions featuring learnable mean and variance (see Appendix B.5). Even this relatively simple parameterization significantly enhances the flexibility of diffusion models compared to conventional diffusion approaches, enabling the learning of complex forward processes.In this section we provide additional results, that directly compare NFDM and NFDM-OT. This results extend Tables 1 and 2 . As observed, NFDM without any additional restrictions provides a NeurIPS Paper Checklist Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. • The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.Justification: We read and checked the code of ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).38th Conference on Neural Information Processing Systems (NeurIPS 2024).OT stands for optimal transport. This designation is used for convenience. While straight trajectories are a necessary condition for dynamic optimal transport, they are not sufficient. 


Decision: This paper proposes Neural Flow Diffusion Models (NFDM), a framework that allows one to introduce a learnable forward process into diffusion models, enabling more flexible forward processes beyond the standard Gaussian.

Reviewers recognize the paper's significant contribution, clear presentation, strong theoretical foundation, and impressive empirical performance on standard benchmarks. I concur with them that the paper presents an interesting methodology and I thus recommend the paper to be accepted. However, when preparing the final version of the paper, the authors should address carefully the remaining concerns of the reviewers. To strengthen the paper, the authors should provide a more thorough comparison with prior work (shiftddpms in particular), conduct more extensive experiments, and clarify the points of confusion regarding the framework's details. Addressing these concerns will significantly enhance the paper's impact.

Reviews: review: paper_summary: paper_summary: The paper introduces Neural Flow Diffusion Models (NFDM), a novel framework that enhances diffusion models by supporting a broader range of forward processes beyond the standard Gaussian. The proposed parameterization technique facilitates the learning of the forward process and minimizes a variational upper bound on the negative log-likelihood in an end-to-end, simulation-free manner. Experimental results on datasets like CIFAR-10, ImageNet 32, and 64 demonstrate state-of-the-art performance in terms of log-likelihood and generation quality.
strength_weakness: strength_weakness: 1. Previous works are mentioned but the introduced NFDM framework is not properly placed in the context of the related works. Therefore, it is difficult to fully assess the novelty and the contributions. 
1. It is unclear how the introduced method would scale to larger datasets and image resolutions and what would be the benefit.
1. The non-Gaussian parameterization does not show substantial improvement over the Gaussian parameterization, raising questions about its necessity.
questions: questions: 1. Can the authors provide a clear statement of their contributions and methods within the context of related works that also aim to extend the diffusion framework to include a learnable and potentially non-linear forward process (e.g. [2])? Specifically, it would be helpful to highlight the advantages and disadvantages of these methods compared to theirs, from both methodological and computational perspectives.
2. How do the authors ensure that the function $F$ is smooth between $z_0$​ and $z_1$​ (Line 120)? 
3. Can the authors clarify how the linear in $\epsilon$ parameterization of $F_\varphi$ relates to the stochastic interpolant framework in  [1].  Such comparison is already outlined in Appendix E where the authors point out the different optimization problems (min-min vs min-max) and state "*Stochastic Interpolants necessitate learning two separate functions for the reverse process: the velocity field and the score function*".  Could the authors clarify the latter statement? It seems that also in the proposed NFDM framework one would need to learn both the $F_\varphi$ (which moreover needs to be derived wrt to time and invertible wrt to $\varepsilon$) and the prediction network $\hat{x}_\theta$ . 
4. For clarity of exposition, it would be beneficial to better introduce the function $\hat{x}_\theta$​, particularly explaining the difference with the score network in conventional diffusion. 
5. Why is the framework still considered simulation-free despite requiring a prediction network $\hat{x}_\theta$​? Can the authors explain why this step does not imply that the model, parameterized by $\theta$, implicitly learns to simulate the process?
6. In Table 1, there appears to be no significant improvement when NFDM uses a Gaussian parameterization compared to a non-Gaussian one. Could the authors comment on this result?
7. In App B.3, the authors also mention that a non-Gaussian parameterization is not scalable for high-dimensional cases.  Have the authors observed any case where a non-Gaussian parameterization leads to an improvement? 
8. Can the authors offer a more detailed comparison of their _Restricted NFDM_ compared to previous works that also constrain the generative process to have straight trajectories? Specifically, what are the computational advantages/disadvantages? 
9. Can the authors better discuss the computational cost of NFDM compared to previous approaches? If I am not mistaken I only see a mention in the conclusion "*this leads to approximately 2.2 times longer optimization iteration of NFDM compared to conventional diffusion models*" but it is not clear to which NFDM parameterisation this estimate refers.  Moreover, while Table 2 depicts a clear advantage of NFDM in terms of NFE, it would be useful to analyze these numbers in the context of the computational complexities of the respective methods. 

[1]  Albergo, Michael S., and Eric Vanden-Eijnden. "Building normalizing flows with stochastic interpolants." arXiv preprint arXiv:2209.15571 (2022). 


[2] Kim, Dongjun, et al. "Maximum likelihood training of implicit nonlinear diffusion model." _Advances in neural information processing systems_ 35 (2022): 32270-32284.
limitations: limitations: The limitations are addressed in the paper.
paper_summary: paper_summary: The authors propose a variant of a diffusion model in which the forward process is parameterized by a neural network that maps from data, noise, and a timestep to the noised data. They show how, if this neural network has certain properties, the SDE and ODEs of the forward and reverse processes can be simulated. They demonstrate good density estimation results and the ability to penalize the learned forward process to encourage the straightness of ODE trajectories.
strength_weakness: strength_weakness: - Exposition could be clearer. Understanding all of the notation in Section 3 in particular took a while. I understand that some of this complexity might be irreducible, but it could be helpful to e.g. have a diagram showing relationships between e.g. the different variaions of , what
learns to predict, the direction stepped in during the reverse ODE, etc.
- It would be interesting to see comparisons between NFDM and NFDM-OT included in the paper. E.g. if both of them could be included in each of the results tables. It would be very interesting to see how they compare FID and NLL, and also how the learned trajectories of NFDM compare to those of Fig 1a.
questions: questions: See weaknesses
limitations: limitations: Adequately addressed
paper_summary: paper_summary: This work proposes a generic framework to make diffusion model's encoder learnable.
strength_weakness: strength_weakness: 1. The work starts from the prediction of forward marginal distribution, and derive the conditional ODE and SDE. However, whether these dynamics corresponding to the same marginal distributions need some rigorous reasoning (e.g., Fokker-Planck Equation or continuity equation). For instance, the claim in L132 and L132 should be explained. Please supplement it. 

2. The authors claimed: "*conventional diffusion the latent variables are inferred through a pre-specified linear combination of the data point and Gaussian noise. This formulation limits diffusion models in terms of the flexibility of their latent space, and makes learning of the reverse process more challenging...*". However, is there any evidence showing that a diffusion model's linear combination of data and noise may lead to degradation or affect the flexibility of the latent space? Also, if the forward path is simply a pre-specified linear combination of the data point and Gaussian noise, wouldn't it result in a more easily learnable backward process? I suggest the authors provide either theoretical or empirical evidence to support their claim and motivation.

[Minor]
- I appreciate that the authors explained the differences and relationships between their proposal and related works, which is informative. I suggest extending Appendix E.1 to include more comparisons of each method's assumptions, such as the priors used, their final objective functions, and why their parameterizations fall short.
questions: questions: 1. The proposed algorithm suggests learning the encoding and decoding parts together (though alternately). However, the loss function shown in Eq. (12) seems prone to collapsing to a trivial solution, despite the parametrizations potentially mitigating this issue somewhat. Indeed, some literature, such as [1], includes additional regularizers to stabilize the training. Did the authors observe any collapse or instability during training? Any trick to avoid the collapse if it is the case?

2. How $g_\psi$ be determined? How it depends on $\psi$?

3. What is the FID performance of NFDM (with/without curvature regularizer) if the total sampling timestep $T$ is set to the standard diffusion model value of $1000$, compared to DDIM or DDPM? The authors claimed diffusion model's pre-defined encoding may lead to inflexible and complicated latent for reverse process. The training objective is derived in a favor for likelihood computation. However, sample quality should be a more important metric to justify the claim. Also, what does the counterpart of Fig. 2 without curvature regularizer?

 

[1] S. Lee, B. Kim, and J. C. Ye. Minimizing trajectory curvature of ode-based generative models. arXiv preprint arXiv:2301.12003, 2023.
limitations: limitations: The authors have discussed the limitation that using a learnable forward path may double the training runtime.
paper_summary: paper_summary: The paper introduces Neural Flow Diffusion Models (NFDM), which incorporate a learnable forward process to enhance diffusion modeling. Traditional diffusion models rely on a fixed forward process, often complicating the reverse process and resulting in costly inference. NFDM aims to address these challenges by introducing flexibility and improved generative dynamics through a learnable forward process. The paper demonstrates NFDM’s capabilities through experiments on CIFAR-10 and ImageNet datasets, achieving state-of-the-art results in terms of negative log-likelihood. Additionally, NFDM is applied to learn bridges between distributions on the AFHQ dataset, further showcasing its flexibility and utility in different generative modeling tasks.
strength_weakness: strength_weakness: 1. A previous work Shiftddpms [1] share similar idea with NFDM. It incorporates a learnable forward process (and its corresponding reverse process) to enhance conditional diffusion modeling. The paper doesn't refer to and compare with Shiftddpms.
2. The paper does not extensively discuss the generalization of the results to other settings or the robustness of the results to violations of the assumptions made​. This omission can be critical as it limits the confidence in applying the findings to different datasets or under different conditions than those explicitly tested.

[1] Shiftddpms: Exploring conditional diffusion models by shifting diffusion trajectories
questions: questions: Please see weaknesses.
limitations: limitations: The authors have adequately described the limitations and potential negative societal impact of their work.

Comments: title: Discussion
Thanks for clarification and replies.
comment: We would like to thank reviewers for their time and valuable contributions and discussion, which enhance the quality of our work. We are grateful for the reviewers' recognition of the NFDM framework as a “significant contribution”, “elegant”, “neat”, “novel” and “clearly explained”. We also appreciate the acknowledgment of our experiments as "robust and extensive," and the recognition of the potential of this research direction to be "impactful" and "influence future research." We are especially thankful for the reviewers’ unanimous recommendations.

Here we summarise the main contributions of NFDM and the discussion:

- We present NFDM, a novel general framework for diffusion models, which allows non-linear, time-dependent and learnable forward processes.
- We present novel theoretical results, where we derive a simulation-free objective for NFDM, which is crucial for efficient training. We also derive ordinary and stochastic differential equations for sampling.
- NFDM subsumes a large portion of existing approaches as special cases, enabling the free specification of the forward process (Section 7 and Appendix E). Moreover, it extends these much further, offering additional flexibility by introducing learnable parameters.
- NFDM, with a simple parameterization of learnable transformations and without hyperparameter tuning, consistently and significantly outperforms baselines in density estimation across a range experiments (Table 1).
- With a simple modification, our framework can also facilitate building bridges between two different distributions, NFBM.
- NFDM opens a prospect for building new generative dynamics with specific properties, which conventional diffusion models like DDPM and DDIM are incapable of.
- NFDM with additional penalty on curvature of deterministic generative trajectories (NFDM-OT), enables the learning of straight-line generative dynamics. This advancement permits the generation of samples in significantly fewer steps, outperforming baselines in few-step generation (Table 2).
- NFDMs does not require any additional techniques to ensure stable training.

We trust that the insights gained from our discussion will bolster your support for acceptance!
Thank you for the response. It addressed my concerns and I've raised my score to a 7.
Please participate to the discussion.
I appreciate the time and effort the authors invested in addressing the concerns I raised. Together with your clarifications, I consider your work to be a relevant contribution to the field, and therefore, I have decided to raise my score.
I appreciate the reviewers' clarification and have decided to increase my score.
I acknowledge having read the authors' rebuttal. My overall assessment of the paper remains unchanged, and I continue to support my current rating.
8. In the paper, we consider two related works that also propose approaches to learn straighter generative trajectories: [6] and [7].
    
    The approach in [6] may be seen as a special case of NFDM (see the last paragraph in Appendix E.1). As soon as the authors in [6] propose to learn an additional function to parameterize the forward process, the training complexity of their method is comparable with NFDM. However, in NFDM-OT, we also need to calculate the curvature penalty (eq. 16), which makes the training more expensive compared to [6].
    
    In [7], the authors proposed constructing the forward process with optimal transport data-noise couplings. The computational disadvantage of this approach is the necessity to find an optimal coupling for each mini-batch, which may be costly for large batch sizes and some complicated metrics (e.g., distance between 3D structures).
    
    In both [6] and [7], the inference works like in conventional diffusion models. However, due to the chosen parameterization, the inference of NFDM is more computationally expensive. In our experiments, instead of predicting the velocity field directly, we predict the data point $x$ and then substitute this prediction into the forward process. Therefore, we need to infer two functions: $x_\theta$ and $F_\phi$.
    
    The primary goal of NFDM-OT is to demonstrate the NFDM framework flexibility and its capability to learn generative dynamics with desired properties. In contrast to [6, 7], which only allows minimisation of trajectories curvature, NFDM allows learning generative dynamics with a variety of different properties like we demonstrate in Figure 2. Nevertheless, NFDM-OT is able to efficiently learn straight-line generative trajectories and demonstrates better performance compared to prior works as illustrated in Table 2.
    
9. In our experiments, we observe a similar slowdown for both Gaussian and non-Gaussian parameterizations of NFDM. The training iteration of NFDM takes approximately 2 times longer compared to conventional diffusion models that don’t parameterize the forward process. This happens because we use neural networks of similar sizes to parameterize both the forward and the reverse processes.
    
    Similarly, due to the parameterization of the generative process which uses the forward process, one inference iteration also takes approximately 2 times longer compared to the inference of conventional diffusion.
    
    We would like to emphasize that such properties of the models are consequences of these design choices. For simplicity of the experimental setup, we use neural networks of the same size to parameterize the forward process, which might be redundant. Additionally, it’s possible to choose an alternative parameterization of the reverse process that wouldn’t rely on the forward process for inference. That’s why we think that the comparison under the same number of function evaluations is more important.
    
    Nevertheless, even with this current parameterization, NFDM-OT demonstrates better performance compared to prior work under the same computational budget (if we double the NFE in Table 2).

We hope that the clarifications provided, highlighting the novelty and strong performance of our work, along with further explanations of technical details will enhance your support for our paper's acceptance!

[6] Lee, Sangyun, Beomsu Kim, and Jong Chul Ye. "Minimizing trajectory curvature of ode-based generative models.”

[7] Pooladian, Aram-Alexandre, et al. "Multisample flow matching: Straightening flows with minibatch couplings.”
4. In conventional diffusion, where the forward process is parameterized as $F(\varepsilon, t, x) = \alpha_t x + \sigma_t \varepsilon$, there is a straightforward connection between the score predictor and the predictor of the data point: $s_\theta(z_t, t) = \frac{\alpha_t x_\theta(z_t, t) - z_t}{\sigma_t^2}$ (see [4]). Thus, during the training procedure, we can learn to predict either the data point predictor or the score function. Then, after training, any of these functions can be solved for from the other. From a theoretical standpoint, all these parameterizations are equivalent.
    
    When the forward process of NFDM is parameterized linearly as discussed above, NFDM exactly corresponds to conventional diffusion, and the function $x_\theta$ maintains the same relationship with the score function. However, in the general case, when $F_\phi$ is an arbitrary function, these connections between $s_\theta$ and $x_\theta$ do not exist. In our work, the parameterization of the reverse process through the function $x_\theta$ is a design choice. As discussed, $x_\theta$ retains relationships with the score function in the linear case. However, in the general case, we do not expect it to possess any special properties. We will include a discussion on this in the camera-ready, if accepted.
    
5. In the literature on diffusion models, the term "simulation-free" is most often taken to mean that during training the method does not require numerical simulations of the SDE or the integration of any dynamics at the training stage. For instance, approaches like [5], or others based on Schrödinger Bridge theory, require numerical simulation of an SDE that describes the forward process to sample $z_t \sim q(z_t|x)$. Such requirements make the training procedure significantly more expensive and are difficult to scale.
    
    In contrast, the NFDM framework allows for direct sampling of $z_t$ by inferring the function $F_\phi$, which includes no numerical simulations. We also do not need to simulate the generative dynamics during training, making the NFDM framework simulation-free.
    
6. As previously mentioned, the main contribution of our work is a comprehensive framework that supports a wide range of forward processes. The primary objective of the experiments with non-Gaussian parameterization of the forward process is not to surpass the Gaussian parameterization but to demonstrate the flexibility of the NFDM framework and its capability to handle even such complex parameterizations. 
    
    We believe there are two main reasons why the non-Gaussian parameterization of the forward process does not show significant improvement compared to Gaussian parameterization. First, we have only evaluated our approach on image datasets, where Gaussian noise appears to naturally fit well. Secondly, we chose a relatively simple parameterization of the non-Gaussian forward process, which may be suboptimal. It is possible that in different domains or with a more powerful parameterization, non-Gaussian parameterization may prove more beneficial, but we leave this for future research. In this work, we aimed to demonstrate the functionality of the NFDM framework.
    
    Additionally, we would like to emphasize that even the Gaussian parameterization of NFDM with data-dependent variance is a generalization of existing diffusion models and leads to state-of-the-art log-likelihood estimation results.
    
7. Indeed, in Appendix B.3, we discuss that *a naive parameterization* of $F_\phi$ with arbitrary non-Gaussian functions is not scalable (density evaluation can be O(dim^3)). However, if we parameterize $F_\phi$ with functions that provide easy access to the log-determinant of the Jacobian of transformation, we can apply this parameterization even in the high-dimensional case. This is why, in our experiments on images with non-Gaussian parameterization of the forward process, we use a Glow neural network architecture. This architecture is scalable. Our experiments in Table 1 demonstrate that this non-Gaussian parameterization can lead to better log-likelihood estimation results, achieving state-of-the-art *on both* CIFAR-10 and ImageNet.

[4] Kingma, Diederik, and Ruiqi Gao. "Understanding diffusion objectives as the elbo with simple data augmentation.”

[5] Zhang, Qinsheng, and Yongxin Chen. "Diffusion normalizing flow.”