Neural Flow Diffusion Models: Learnable Forward Process for Improved Diffusion Modelling
Authors: Grigory Bartosh
Dmitry Vetrov
Christian A. Naesseth
Venue: NeurIPS 2024 poster (2024)
Keywords: diffusion
generative models
variational inference
Primary area: diffusion_based_models

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.

Neural Flow Diffusion Models: Learnable Forward Process for Improved Diffusion Modelling
Grigory Bartosh University of Amsterdam g. bartoshèuva.nl
Dmitry Vetrov
Christian A. Naesseth University of Amsterdam c.a ..nsseethuuva.n.
Constructor University, Bremen dvetrov@constructor university
Abstract
Conventional diffusion models often rely fixed forward process, which im- plicitly 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), novel framework that enhances diffusion models by supporting broader range of forward processes beyond the standard linear Gaussian. We also propose a novel parameterization technique for learning the forward process. Our framework provides an end-to-en end, simulation-Fi free optimization objective, effectively minimizing variational upper bound on the negative log-likelihood. Experimental results demonstrate NFDM't s strong performance, evidenced by state-of-the-ar likelihoods across a range image generation tasks. Furthermore, we investigate NFDM's capacity for learning generative dynamics with specific characteristics, such as deterministic straight lines trajectories, and demonstrate how the framework can be adopted for learning bridges between two distributions. The results underscores NFDM's versatility and its potential for wide range of applications.
Introduction
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 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 Nevertheless, most existing diffusion models fix the forward process to be predefined, usually linear, Gaussian which makes unable adapt the task at hand simplify the target for the reverse process. At the same time there is growing body of work that demonstrates how modifications of the forward process improve performance terms of generation quality (39 66 9
Diffusion models 16 20
21
48 63 69
65).
likelihood estimation 28 40 4
or sampling speed (32 44 644
In this paper, present Neural Flow Diffusion Models (NFDM), 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 conditional Gaussian forward process, NFDM can accommodate any continuous (and learnable) distribution that can be expressed as an invertible mapping applied noise. We also derive, and leverage, free optimization procedure for NFDM, that minimizes a variational upper bound on the negative log-likelihood (NLL).
end-to-end simulation-
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
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

experimental results on CIFAR- 10, ImageNet 32 and 64, attaining state-of-the art NLL results, which crucial for many applications such as data compression [22 anomaly detection 7 [3 and
out-of-distribution detection 51 70
Then, leveraging the flexibility of NEDM, 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 case study, we discuss curvature and obstacle avoidance penalties on the deterministic generative trajectories. Our empirical results indicate improved computational efficiency compared baselines on CIFAR-10, downsampled ImageNet, and synthetic data.
We summarize contributions follows:
1. We introduce Neural Flow Diffusion Models (NFDM), improving diffusion modelling through learnable forward process.
2. We develop end-to-ent optimization procedure that minimizes upper bound on the negative log-likelhhood in a simulation-Fi manner.
3. We demonstrate state-of of-tthe-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 signifieantly faster sampling speeds and enhanced generation quality with fewer sampling steps.
2 Background
Diffusion models are generative latent variable models consisting of two processes: the forward and the reverse (or generative) process. The forward process is dynamic process that takes data point x ~ g(x). x E TR,,, and perturbs it over time by injecting noise. This generates trajectory of latent variables (=(())=()011) conditional on the data x, where [0, 1] a fixed time horizon and z, zt) E RD. The (conditional) distribution can be described by an initial distribution g(zo|x) and Stochastic Differential Equation (SDE) with a linear drift term f (zz, t) : RR x [0, 1] + R. scalar variance g(t): [0, 1] + Rt, and a standard Wiener process w:
dz. = je (Zz, 1)dt + g(t)dw.
Due to the linearity of fr, we can reconstruct the conditional marginal distribution q(z, [x) (Zz ayx, 1. Typically, the conditional distributions evolve from some low variance distribution q(zo xx) ö(zo- (i)oo>))a x) to unit Gaussian (z (x) s N(zz; 0, 1). This forward process is then reversed by starting from the prior 21 ~ NN(11 0, I), and following the reverse SDE 33:
dz, = j t)dt + g())dw, where jB(2,, = j (2,
(2,,t)dt
t) f" 9² (1) 2, log (zg))
Here, w denotes standard Wiener process where time flows backwards. Diffusion models ap- proximate this reverse process by learning Vz, log ((x)), known as the score function, through Ae-weighted denoising score matching loss:
E
De (ee]ss(Z),t)) Vz, log (2gg(2;(2))
x(1)g(x.:)
where u(t) represents a uniform distribution over the interval [0. 1]. and 8p : RR. x (0, 10 R. is a learnable approximation. With a learned score function so so(z,, 1). one can generate sample from the reverse process by first sampling from the prior 21 ~ "N(z;; 0, I). and then simulating the reverse SDE, resulting in sample Zo
po (Zo) Re (Zo)) 2 q(x):
dz FF (z., 92 (z) t) g(t)dw.
Diffusion models possess several important properties. For example, for specific X, the objec- tive (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 simulation-free procedure. This means that simulating either the forward or reverse processes through its SDE is not necessary for sampling 2z, nor is it necessary for estimating the gradient of the loss function. Instead, we can

---

directly sample z₃ ~ g(z])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(2+) as the SDE (eq.
dz: = f(ze, t)dt,
where
(5)
g Vn, log (pz)))
f(z, t) = 8 (z., 1)
(6)
2
This implies that we can sample from diffusion models deterministically, allowing the use of off- the-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 17
3 Neural Flow Diffusion Models
In diffusion models the forward process defines stochastic conditional trajectories (z(()) and the reverse process tries to match the marginal distribution of trajectories. This construction can be viewed as specific type of hierarchical Variational Autoencoders (VAEs) 30 45). However, in conventional diffusson models the latent variables are inferred through a pre-specified linear combination of the data point and Gaussian noise. This formulation limits diffusson models in terms of the flcxibility of their latent space, and makes learning of the reverse process more challenging. To address this limitation, we propose generalized form of the forward process that enables the definition and learning of a broad range of distributions the latent space. From a practical perspective, 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 hierarchical VAE, which gives tighter bound model's NLL (see an extended discussion in Appendix B.T..
In this section, introduce Neural Flow Diffusion Models (NFDM) framework that generalizes conventional diffusion models. The key idea in NFDM is to define the forward process' conditional SDE implicitly via learnable transformation (= f, 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-brai basedd and simulation free training. Previous diffusion models emerge special cases when the data transformation is linear, fime-independent, and/or additive Gaussian.
3.1 Forward Process
We approach the forward process constructively. The ultimate goal is learnable distribution over trajectories, (Ze0)) given x, realized by conditional SDE constructed in three steps.
First, (implicitly) define the conditional marginal distributions qe x) for / E [o, 1] using (e, 1, x). Then, we introduce the corresponding conditional ODE that together with an initial distribution q,e (z022 x) matches the conditional marginal distribution O.p (z, /x). Finally, define conditional SDE that defines a distribution over trajectories (z(t))x<(o.1) (o.1))) with marginal distributions 4, (z(2//)) xx for each t.
Forward Marginal Distribution. We characterize the marginal distribution d (z, xx) of the forward process trough a function that transforms noise sample e into 2,, conditional on the time step and data point x:
2: Fo,,t,,
(7)
where F.. R [0. RD + RD and q(E) NNee; 0). This defines the conditional distribution of latent variables 9:p (z, x). Additionally, eq. facilitates direct and efficient sampling from q (z, x) through F. in eq.
Conditional ODE. We assume that Fg differentiable with respect to e and I, invertible with respect to s. Further, we assume that fixing specific values of and and varying from 0 to 1 results in
3

---

smooth trajectory from Zo to Z1. Differentiating these trajectories over time yields velocity field corresponding to the conditional distribution 9: (z()(x)) thereby defining a conditional ODE:
where (z,, = 0F (E,
dz, fip x,d,
(8)
Ot
(z(,t,))
Therefore, if we sample 4ep (z)lx) and solve eq. 8 until time I, we have (p(2)
The time derivative of F.p may be calculated efficiently using automatic differentiation tools like PyTorch [42] or JAX [5] (see details in Appendix(B..
Conditional SDE. The function F& and the distribution of the noise q(s) together defines (z x). To completely define the distribution of trajectories {z(t))eejo,1) we introduce a conditional SDE that starts from sample Zy and runs forward in time.
With access both the ODE and score function 2z, log (o x), an SDE 61 with marginal distributions Uu (Z, (x), is:
dz: = iF (Ze, 1, x)dt + 90 (t)dw, where
g
1e (z,, 1, (,,
Vz, log (2+
2
Here, 9e. [0, + R_ is scalar function, and represents standard Wiener process. Note that 9e only influences the distribution of trajectories (E(o)) The marginal distributions are the same, (2, xx), for any choice of 9p-
The SDE in eq. (9) requires access to the conditional score function Xx, log Qu [ze ()), which in the general case be computationally costly. However, for Fe that allow efficient evaluation of the log-determinant of the Jacobian matrix, the score function can be calculated efficiently. Examples of such Fúp are functions linear with respect e or RealNVP style architectures [[4] 29 The calculation of this score function is further discussed in Appendix B.3
3.2 Reverse (Generative) Process
To define the reverse (generative) process. we specify reverse SDE that starts from z p(z1) and rums backwards in time. To do so we first introduce a conditional reverse SDE that reverses the conditional forward SDE (eq. g Following 3, we define:
dz, 7 (2,
(2,, 1, x)dl + 9b (1)dw, where
(10)
F (zt, 1, x) = f&p (2,, g2
Vz. log (z.lx).
2
Secondly, leveraging this reverse SDE (see eq.
(1Ob) we incorporate a prediction of x:
dz, = Se (2,, t)dt 9 (1)dw, where le.0 = f (Ze, So(2,, 1));
(11)
and Ry RR 10, Rº function that predicts the data point x. This SDE defines the dynamics of the generative trajectories (2(())=40011)
To fully specify the reverse process, it is also necessary to define a prior distribution p(z)) and reconstruction distribution p(x[zu). In all of our experiments, set the prior p(zz) to be unit Gaussian distribution N ((z1: 0, I) and let p(x|zg) be a Gaussian distribution with a small variance (xx Zo, I2I,, where 8² = 10
The above parameterization of the reverse process not the only possibility. However, 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.

---

Algorithm Optimization of NEDM
Algorithm Stochastic Sampling from NEDM
Require: q(x), F.p. Sp. &o
Require: Fip. 9 Xp, T - number of steps
for learning iterations do
At = 7. 21 ~ p(z1)
~ x(x), & &~ u(t)
for t = -..... 2 do
Z: ~ qp Ze xx)
w N/(0,1)
I)
7: T
[ 2 (4)
| (Ze, t2 - (z. (2),
Z:-Ar - 2t -
Gradient step on and w.r.t. L
end for
end for
x ~ p(x(zo)
3.3 Optimization and Sampling
With F& in eq. 7 parameterized such that 9: (Zolx) 8(x zo) and (z xx) Rs p(z1). propose to optimize the forward and reverse processes of NFDM jointly, minimizing the following objective:
20 (1) | (Ze, - (a 011
= (u(()g()))
(12)
Since the forward process is parameterized by so, we need to optimize the objective (eq. 12 with respect both yp and 9 jointly. We discuss parameterization of the forward process in Appendix B.5
As we demonstrate in Appendix I the objective & shares similarities with the standard diffusion model objective (0), and it provides variational bound on the model's 's log-likelihood log po. (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 also diseuss the role of 9y 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 necessarily lead to accurate predictions of x.
A key characteristic of the NFDM objective is its compatibility with the simulation saree paradigm, which is critical for efficient optimization. We summarize the training procedure in Algorithm
To sample from the trained reverse process we can simulate the SDE, as defined in Section 3 This procedure is summarized in Algorithm 2] Additionally, during the sampling process, we can adjust the level of stochasticity by modifying g--(t (t) (eq. (9)). It is important to note that changes to (t)
also influence 75 (eq. (9)) and fe.p (eq. 11)). In the extreme case where (t) O, the reverse process becomes deterministic, allowing us to utilize off-the-shelf numerical ODE solvers and to
estimate densities 66] 17.. We provide extended sampling discussion in Appendix B.4
Neural Flow Bridge Models
In this section discuss simple modification the NFDM framework that enables us learn bridges between two data distributions, g(xo) and g(x1), modification we refer to as Neural Flow Bridge Models (NFBM).
In the context of bridges, we consider a joint data distribution g(xo: x1) (which may be factorized as (xx))9(x1) if paired data is unavailable). Our goal is to learn generative process that starts from X1 and generates Xa such that ((xo) = (xo)) = [ ((x1) } (x [x )dx To turn NFDM into NFBM, modify both the forward and reverse processes. In NFDM. the forward process is defined by two functions: (e, f, x) (see eq. (7)) and g (see eq. ( For NFBM, we let Fs. depend on both Xy and X1, thereby conditioning the entire forward process these data points:
Zi (,,
(13)
Similar to the discussion in Section 3.2, for the reverse process of NFBM, we predict the data points Xo and X; from z, and f, and substitute them into the conditional reverse time SDE (eq. (10):
dz: = fe. t)dt + g (t)dw, where (0, t) = f (ze, t, x0 (zz, t)).
(14)
5

---

Here, the function X, (z,, f) returns predictions of both Xo and X1. Alternatively, the data point X1 can be reused as conditioning for intermediate steps of the reverse process instead of predicting both points, strategy further detailed in Appendix C.2
When Fo in eq- (13) is parameterized such that 0.p (zn) Xn, x1) 8(x zo) and qo [1 Xxo x1) R p(zi (x1) ) (see Appendix C.I for parameterization details of the NFBM forward process), we propose training NFBM by minimizing the following objective:
1 (2+, Xo, x1) - (2,, 1)
xx.t)))(xx.x)aataa Xo.x.)
(15)
(g)
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 po. (Xo). and is compatible with the iimulation-free paradigm. Furthermore, it allows for sampling with various levels of stochasticity (by adjusting (4)) including deterministic sampling (when Go (1) = 0).
Thus, the NFDM framework not only enables the construction of generative models, but also facilitates learning bridge models between two data distributions.
5 Restricted NFDM
We introduce NFDM as a powerful framework that enables learning of the forward process. However, there is in general 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 highly beneficial, it enables generation with far fewer steps in the ODE solver. One approach is introduce penalties on the curvature of the reverse process's ' trajectories and Ict the learnable forward process to adapt to align with the reverse process. As 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.I
We propose learning model with curvature penalty as suggested by 25
2
d (ze
(Ze, dl
LOT= = C + XLere,
where
Lerv = (xft)g
(16)
We refer this variant NFDM-O1 Lerv is additional curvature loss that penalizes the second time derivative of the generative ODE trajectories. L is estimated as in Section [.3 whereas when calculating Lerv we set (t) = 0. Lerv ensures that the generative trajectories are straight lines. In experiments,
set À = 10 Empirical evidence suggests that higher values of X 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 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 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.
6 Experiments
We first showcase results demonstrating that NFDM consistently achieves better likelihood compared to baselines, obtaining state-of-the-an diffusion modeling results on the CIFAR-10 (31] and down- sampled ImageNet
671 datasets. Then, we explore the NFDM-OT modification, which penalizes
'OT stands for optimal transport. This designation is used for convenience. While straight trajectories are necessary condition for dynamic optimal transport, they not sufficient.

---

Table I: Comparison of NFDM results with baselines density estimation tasks. We present results terms of BPD, lower is better. NFDM achieves state-of-the-ee results across all three benchmark tasks.
Model
CIFAR10
ImageNet 32
ImageNet 64
DDPM
20
3.69
Score SDE [6]
2.99
Improved DDPM
2.94
3.54
VDM [28
2.65
3.72
3.40
Score Flow
59
2.83
3.76
Flow Matching
36
2.99
3.53
3.31
Stochastic Interp.
2.99
3.48
i-DODE [76
2.56
3.43
NDM 4
2.70
3.55
3.35
MuLAN
49
2.55
3.67
NFDM (Gaussian q:p (Z)
2.49
3.36
3.20
NFDM Lon-Gaussian 9p(2)) (Z(()))) x))
2.48
3.34
3.20
(a) Score SDE 61
(b) NFDM-OT (this paper)
Figure 1: Comparison of trajectories between the data distribution (on the left) and the prior distribu- tion (on the right), as learned by conventional diffusion and NFDM-OT.
the curvature of the deterministic generative trajectories. The NFDM-OT reduces trajectory curvature, significantly lowering the number generative steps required for sampling. Finally, we demonstrate the NFBM modification, learning bridges the AFHQ [ dataset and several synthetic examples.
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.o (2x is a Gaussian with learnable mean and covariance (see details in Appendices B.5 and C.I). For detailed description of parameterizations and other experimental details, please refer Ap- pendix F
The primary aim of the experiments with NFDM-OT and NFBM is not to introduce 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 available at wttttttpp :/gg h///ggttuut thub. com/ rrrigoryBarrosshneuraaa diiilldffffooo
6.1 Likelihood Estimation
For the first experiments, in addition to Gaussian we also provide results with non Con-Gaussian parame- terization of the forward process (see Appendix B.5).
the NLL results on
CIFAR-10 10 and two downsampled ImageNet datasets. Notably, NFDM outperforms diffusion- baselines on all three datasets with both parameterizations, achieving state-of-the-arr performance. The improved performance due to NFDM is natural progression for the following reasons.
7

---

Table 2: Summary of FID results for few-step generation. The table divided into three sections, based on different types of methods: those that do not minimize curvature, solvers for pretrained models, and models that specifically aim to minimize curvature. For the DDPM, we include results corresponding to two distinct objectives: the full ELBO-based objective and a simplified objective (Esimple) NFDM-OT outperforms baselines with comparable NFE values.
CIFAR-10
ImageNet 32 FID 4
ImageNet 64 NFE
Model
NFE 4
FID 4
NFE
FID
DDPM (Lsimple)
20
1000
3.17
DDPM (ELBO)
20
1000
13.51
Flow Matching
33
142
6.35
122
5.02
138
14.14
DDIM
57
13.36
12
5.28
DPM Solver
38
24
2.75
Trajectory Curvature Minimization
32
5
18.74
4
17.28
4
38.45
Multisample Flow Matching [44
12
7.18
12
17.6
12.44
9.83
27.70
NFDM-OT (this paper)
4
7.76
6.13
17.28
12
5.20
12
4.11
12
11.58
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. (2)) used for training the NFDM is also variational bound on the likelihood.
Second, diffusion models can be seen as hicrarchical VAEs. From this perspective, most baselines in Table I resemble VAEs with either fixed 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 mean and covariance that is non-linear in and t. This flexibility allows the NFDM to better conform to the reverse process, consequently enhancing likelihood estimation.
6.2 Straight Trajectories
We next evaluate NFDM-OT, designed to penalize the deterministic generative trajectory curvature. First, compare NFDM-OT with conventional continuous-time diffusion model. Figure la illustrates deterministic trajectories between a two-dimensional data distribution and unit Gaussian distribution learnt by conventional diffusion model [61] Figure 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 desired. In Appendix G. we provide some additional reult, demonstrating the importance of learnable forward process for NFDM-OT.
Then, we present results NEDM-OT on image datasets. Table 2 reports the FID scores for 2, and 12 Number of Function Evaluations (NFE) with respect to the function fe.. (eq. 11 ). In this experiment, 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 minimize the curvature of generative trajectories
Importantly, NFDM-OT is trained with 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 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 NEDM-OT in Appendix G.2

---

Figure 2: Trajectories of NFBM between two-dimensional data distributions, trained with an addi- tional penalty to avoid obstacles.
6.3 Bridges
We consider the task of translating images of dogs into images of cats from the downsampled AFHQ dataset. Table 3 reports the FID scores images sampled with 1000 steps. Notably, NFBM outperforms the baselines, demonstrating the effectiveness of our framework for learning bridges.
Table 3: FID values dog > cat on AFHQ 64.
Model
FID +
Finally, we demonstrate that the framework flexibility also allows for learning dynamics with specific properties in the case of NFBM.
DSBM [54]
14. 16.16
GSBM (34
12.39
To illustrate this, we adopted the strategy discussed in Section [ and learn NFBM with additional penalty to avoid obstacles (see
NFBM
12.06
visualizes the learned stochastic trajectories. It is evident that NFBM has efficiently learned to translate distributions while avoiding obstacles.
details in Appendix [22). Figure 2
7 Related Work
Diffusion models, originally proposed by (56). have evolved significantly through subsequent de- velopments (60 20 These advancements have resulted remarkable generative quality for high- dimensional 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 (5), combining blurring with Gaussian noise injection 46 g 233. diffusion in the wavelet spectrum 43
and forward processes based on the exponential family [1 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 learn- able 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 transforma-
tions [8 55 40
These methods 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 662 35 38 [3). Additionally, distillation techniques have been applicd [50] 58 75)
enhance sampling speed, albeit 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
8 Conclusion and Limitations
In this paper we introduced NFDM, a novel simulation-free framework for improved diffusion modeling through a Iearnable forward processes. NFDM outperforms baseline diffusion models on

---

standard benchmarks, showcasing its effectiveness. However, NFDM is not just another model, but rather 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-off's. 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 parameteriistions with functions that match specific conditions.
Despite these challenges, 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.
optimistic about the potential of NFDM. NFDM opens up significant
References
[1] M. S. Albergo, N. M. Boffi, and E. Vanden-Eijndee Stochastic interpolants: A unifying framework for flows and diffusions, 2023.
[2] M. S. Albergo and E. Vanden-Eijnden. Building normalizing flows with stochastic interpolants.
In The Eleventh International Conference on Learning Representations, 2023.
[3] B. D. Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-335 326, 1982.
[4] G. Bartosh, D. Vetrov, and C. A. Naesseth. Neural diffusion models. arXiv:2310.08337, 2023.
arXiv preprint
[5] J. Bradbury, R. Frostig, P. Hawkins, M. Johnson, C. Leary, D. Maclaurin G. Necula, A. Paszke,
J. VanderPlas, Wanderman-Milng and Q. Zhang. JAX: composable transformations of Python+NumPy programs. 2018.
[6] R. T. Chen, Y. Rubanova. J. Bettencourt, and D. K. Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
[7] Z. Chen, C. K. Yeo, S. Lee, and T. Lau. Autoencoder- network anomaly detection. In 2018 Wireless telecommunications symposium (WTS), pages 1-5, IEEE, 2018.
[8] Y. Choi, Y. Uh, J. Yoo, and J.-W. Ha. Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 8188-8197, 2020.
[9] G. Daras, M. Delbracio, H. Talebi, A. G. Dimakis, and P. Milanfar. Soft diffusion: Score matching for general corruptions. arXiv preprint arXiv:2209.0544 2022.
[10] V. De Bortoli, J. Thortoon, J. Heng. and A. Doucet. Diffusion schrödinger bridge with applica- tions score-based generative modeling. Advances in Neural Information Processing Systems, 34:1 74955-9997779, 2021.
[11] J. Deng, W. Dong, R. Socher, Li, K. Li, and Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255. leee, 2009.
[12] P. Dhariwal and A. Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780 8794, 2021.
[13] M. Dias, C. L. C. Mattos, L.. da Silva, J. de Macedo, and W. C. Silva. Anomaly detection in trajectory data with normalizing flows. In 2020 international joint conference on neural networks (IJCNN), pages 1-8. IEEE, 2020.
[14) I. Dinh, J. Sohl-Dickstein, and S. Bengio. Density estimation using real NVP. International Conference Learning Representations, 2017.
10

---

[15] J. R. Dormand and P. J. Prince. A family of embedded runge-kutta formulae. Journal of computational and applied mathematics, 6(1):19-26, 1980.
[16) W. Du, H. Zhang. T. Yang, and Y. Du. A flexible diffusion model. In International Conference on Machine Learning, pages 8678-8696. PMLR, 2023.
[17] W. Grathwohl, R. T. Q. Chen, J. Bettencourt, and D. Duvenaud. Scalable reversible genera- tive models with free-form continuous dynamics. International Conference on Learning Representations, 2019.
[18] J. Gu, S. Zhai, Y. Zhang, M. Á. Bautista, and J. M. Susskind. f-DM: A multi-stage diffusion model via progressive signal transformation. In The Eleventh International Conference on Learning Representations, 2023.
[19] M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[20] J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840- 33:6840-6851, 2020.
[21] J. Ho, C. Saharia, W. Chan, Fleet, M. Norouzi, and T. Salimans. Cascaded diffusion models for high fidelity image generation. The Journal of Machine Learning Research, 23(1):2249- 2281, 2022.
[22] Y.H. Ho, C.-C. Chan, W.-H. Peng, H.-M. Hang, and M. Domañski. Anfic: Image compression using augmented normalizing flows. IEEE Open Journal of Circuits and Systems, 2:613-626, 2021.
[23] E. Hoogeboom and Salimans. Blurring diffusion models. In The Eleventh International Conference on Learning Representations, 2023.
[24] Karras, M. Aittala, T. Aila, and S. Laine. Elucidating the design space of diffusion- generative models. Advances in Neural Information Processing Systems, 35:26565-26577, 2022.
[25] J. Kelly, J. Bettencourt, M. Johnson, and D. K. Duvenaud. Learning differential equations that are easy to solve. Advances in Neural Information Processing Systems, 33:4370-4380, 2020.
[26] D. Kim, B. Na, S. J. Kwon, D. Lee, W. Kang, and L-c. Moon. Maximum likelihood training of implicit nonlinear diffusion model. Advances Neural Information Processing Systems, 15:22770-32284 2022.
[27] D. Kingma and R. Gao. Understanding diffusion objectives as the elbo with simple data augmentation. Advances in Neural Information Processing Systems, 36, 2024.
[28] D. Kingma, T. Salimans, B. Poole, and J. Ho. Variational diffusson models. Advances in neural information processing systems, 34:21696-21777, 2021.
[29] D. P. Kingma and P. Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. Advances neural information processing systems, 31, 2018.
[30] D. P. Kingma and M. Welling. Auto-encoding variational Bayes International Conference on Learning Representations, 2014.
[31] A. Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[32] Lee, B. Kim, and C. Ye. Minimizing trajectory curvature of ode-based generative models. In International Conference on Machine Learning, pages 18957-18973. PMLR, 2023.
[33] Y. Lipman, R. T. Q. Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow matching for generative modeling. In The Eleventh International Conference Learning Representations, 2023.
11

---

[34] G.-H. Y. Lipman, M. Nickel, B. Karrer, E. Theodorou, and R. T. Q. Chen. Generalized schrödinger bridge matching, In The Twelfth International Conference on Learning Representa- lions, 2024.
[35] L. Liu, Y. Ren, Z. Lin, and Z. Zhao. Pseudo numerical methods for diffusion models on manifolds. In International Conference on Learning Representations, 2022.
[36] Q. Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint WrXiv:2209.145777 2022.
[37] X. Liu, C. Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference Learning Representations. 2023.
[38] C. Lu, Y. Zhou, F. Bao, J. Chen, C. Li, and J. Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5788 2022.
[39] A. Q. Nichol and P. Dhariwal. Improved denoising diffusion probabilistic models. In Interna- tional conference machine learning, pages 8162-8171. PMLR, 2021.
[40] B. M. G. Nielsen, A. Christensen, A. Dittadi, and O. Winther. Diffenc: Variational diffusion with learned encoder. The Twelfth International Conference on Learning Representations, 2024.
[41] A. Okhotin, D. Molchanov, A. Vladimir, G. Bartosh, V. Ohanesian, A. Alanov, and D. P. Vetrov. Star-shaped denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 36, 2024.
[42] A. Paszke, S. Gross, S. Chintala, G. Chanan, E. Yang. Z. DeVito, Z. Lin, A. Desmaison, L. Antiga, and A. Lerer. Automatic differentiation in pytorch. 2017.
[43] H. Phung, Q. Dao, and A. Tran. Wavelet diffusion models fast and scalable image generators. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10199-10208 2023.
[44] A.-A. Pooladian, H. Ben-Hamu, C. Comingo-Enriche B. Amos, Y. Lipman, and R. Chen. Multisample flow matching: Straightening flows with minibatch couplings. In International Conference Machine Learning. PMLR, 2023.
[45] D. J. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning, pages PMLR, 2014.
[46] S. Rissanen, M. Heinonen, and A. Solin. Generative modelling with inverse heat dissipation. In The Eleventh International Conference on Learning Representations, 2023.
[47] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pages 10684-10695, 2022.
[48] C. Saharia, W. Chan, S. Saxena, L. Li, Whang, L. Denton, K. Ghasemipour, Gon- tijo Lopes. B. Karagol Ayan, T. Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 35:36479-36444 2022.
[49] S. S. Sahoo, A. Gokaslan, C. D. Sa, and V. Kuleshov. Diffussonn models with learned adaptive noise processes, 2024.
[50] T. Salimans and J. Ho. Progressive distillation for fast sampling of diffusion models. In International Conference on Learning Representations, 2022.
[51] J. Serrà, D. Álvarez, V. Gómez, O. Slizovskaia, J. F. Núñez, and J. Luque. Input complexity and out-of-distribution detection with likelihood-based generative models. In International Conference Learning Representations, 2020.
12

---

[52] N. Shaul, R. T. Chen, M. Nickel, M. Le, and Y. Lipman. On kinetic optimal probability paths for generative models. In International Conference on Machine Learning, pages 30883-30907. PMLR, 2023.
[53] N. Shaul, J. Perez, R. T. Q. Chen, A. Thabet, A. Pumarola, and Y. Lipman. Bespoke solvers for generative flow models. In The Twelfth International Conference on Learning Representations 2024.
[54] Y. Shi, V. De Bortoli, A. Campbell, and A. Doucet. Diffusion schrödinger bridge matching. Advances in Neural Information Processing Systems, 36, 2024.
[55] R. Singhal, M. Goldstein, and R. Ranganath, Where to diffuse, how to diffuse, and how to get back: Automated learning for multivariate diffusions. In The Eleventh International Conference on Learning Representations, 2023.
[56] J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning. pages
PMLR, 2015.
[57] Song, C. Meng, and S. Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021.
[58] Y. Song, P. Dhariwal, M. Chen, and I. Sutskever. Consistency models. In The Eleventh International Conference on Learning Representations, 2023.
[59] Y. Song, C. Durkan, I. Murray, and S. Ermon. Maximum likelihood training of score-based diffusion models. Advances in Neural Information Processing Systems, 34:1415-1422 2021.
[60] Y. Song and S. Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[61] Y. Song, J. Sohk-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score- generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
[62] H. Tachibana, M. Go, M. Inahara, Y. Katayama, and Y. Watanabe. Ita"fo)ttaylor sampling scheme for denoising diffusion probabilistic models using ideal derivatives. arXiv preprint arXiv:2112.13339, 2021.
[63] X. Tan, Chen, H. Liu, Cong, C. Zhang, Y. Liu, X. Wang, Y. Leng, Y. Yi, L. He, al. Natu- ralspeech: End-to-end text-to-speech synthesis with human-level quality. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024.
[64] A. Tong, N. Malkin, G. Huguet, Y. Zhang. J. Rector- Brooks, K. Fatras, G. Wolf, and Y. Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. In ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems, 2023.
[65] B. L. Trippe, J. Yim, D. Tischer, D. Baker, T. Broderick, R. Barzilay, and T. S. Jaakkola. Diffusion probabilistic modeling of protein backbones in 3D for the scaffolding problem. In The Eleventh International Conference Learning Representations, 2023.
[66] A. Vahdat, K. Kreis, and J. Kautz. Score-based generative modeling in latent space. Advances in Neural Information Processing Systems, 34:1 1287-1 2021.
[67] A. Van Den Oord. N. Kalchbrenner, and K. Kavukcuoglu. Pixel recurrent neural networks. In International conference on machine learning, pages 1747-1756. PMLR, 2016.
[68] P. Vincent. A connection between score matching and denoising autoencoders. Neural compu- tation, 23(7):1661 2011.
[69] J. L. Watson, D. Juergens, N. R. Bennett, B. L. Trippe, J. Yim, H. E. Eisenach, W. Ahern, A. J. Borst, R. J. Ragotte, I.. F. Milles, et al. De novo design of protein structure and function with rfdiffusion. Nature, - 2023.
13

---

[70] Z. Xiao, Q. Yan, and Y. Amit. Likelihood regret: An -itto--dittrbtttion detection score for variational auto-encoder. Advances in neural information processing systems, 33:20685-20666. 2020.
[71] R. Yang and S. Mandt. Lossy image compression with conditional diffusion models, Advances in Neural Information Processing Systems, 36, 2024.
[72] T. Yin, M. Gharbi, R. Zhang, E. Shechtman, F. Durand, W. T. Freeman, and T. Park. One-step diffusion with distribution matching distillation. arXiv preprint arXiv:2311.18828, 2023.
[73] Q. Zhang and Y. Chen. Diffusion normalizing flow. Advances in Neural Information Processing Systems, 34:16280-16222 2021.
[74] Z. Zhang, Z. Zhao, J. Yu, and Q. Tian. Shiftddpms: exploring conditional diffusion models by shifting diffusion trajectories. Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 3552-3560, 2023.
[75] H. Zheng, W. Nie, A. Vahdat, K. Azizzadenesheli, and A. Anandkumar. Fast sampling of diffusion models via operator learning. International Conference on Machine Learning, pages 42390-42400 PMLR, 2023.
[76] K. Zheng, C. Lu, J. Chen, and J. Zhu. Improved techniques for maximum likelihood estimation for diffusion odes. In International Conference on Machine Learning, pages 42363-42389. PMLR, 2023.
14

---

A Derivations
A.1 Derivation of NFDM's Objective
Our goal is to derive a variational bound on the negative log-likelihood NFDM. For the conve- nience of the following derivations will consider the forward process conditional reverse SDE (eq. (10) that starts from 2 ~ (z) () and flows backwards in time:
B dz: f (Ze, t, x)dt + (t)(w...
(17)
Similarly to a conditional SDE (eq. (9) flowing forwards in time, this SDE also corresponds to conditional marginal distribution qu (Z x), which is implicitly defined by the invertible transformation
Z; (e x) (eq. (7). where
q(E).
For the reverse process we use the definition from Section 3.2 as reverse SDE (eq. (TT) that starts from the prior distribution p(z1), flows backwards in time, and finally samples p(x
dz: ff (Ze, t)dt 9vp (t)
(18)
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) [e[[o.] to discrete-time trajectories Zo, Z 21, where T
represents the number of discrete steps of size +.
The discrete conditional reverse process starts from distribution (Z (2) and follows conditional reverse Markov chain:
go ap(Zi-aelz: x) -V (Z St, (t))
(19)
Similarly, the discrete generative process shares distributions p(21) = p(zz) and ((x p(x|z.)
with the continuous one. However, instead of following the reverse SDE (eq. (8)) it follows the reverse Markov chain:
= No (Zz (Zz-sii
Atg
(20)
It's important to note that once we discretize the trajectories, q. (2, x) A (2,lx) and po.p + Pe.p(Z)) (Z:) except for t 1. However, this discretization corresponds to Euler- Baller-Maruyama method of integrating SDEs. Thus, we know that as At + 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 [66 20
By(x)) |- so ) S Pyve Lait + Epriors
where
(21)
Lror Exg, log (ogg((((22))
(22)
Laift = Σ (x.()) DKI. (9 (Z, A.,,, Pl (Z)-
(23)
Eprior Ey(x) [DKL (To (Z) (x)))
(24)
Equation 21 establishes variational upper bound on the negative log-likelihood of discretized
reverse process. Taking the limits of this upper bound, the reconstruction term Crec Crec and the
15

---

prior term Cprior = Eprior retain their forms. However, we can rewrite the diffusion term:
Caift = lim Cait
(25)
At-0
lim At-0
Σ t
(x.r,)
[DII arlz,, 2/
(26)
lim At-0
Bgg (x.2c)
3kd A1S0 7/6 AIS (e. 1,x)
(27)
2
Atf
( ((,,t.)))) - (E)s
lim At-0
(x.2,)
(28)
2 (t)
1=
PB (2 (z,, f2 - Sor 1) 1/
= =Bu(t)g(xx2.)
(29)
(t)
where u(t) is a uniform distribution over unit interval.
Therefore, the discretized trajectory tends to the continuous one when A O, we can the obtained limits as a variational bound on the negative log-likelihood of the continuous model:
exx) log Pa (x)] S L Lree + Laitf + Ppiior-
(30)
In general case, when function F is actually parameterized by the objective L implies that we must optimise all three terms with respect to parameters yp and However, when F. is parameterized (zı (z (x) = p(z1) for any yp the reconstruction Lree and prior
such that 9, (Zulx) = 8(x zo) and q, Lprior terms are constants with respect to the parameters.
A.2 Derivation of NFBM's Objective
Similarly to the derivation of the NFDM objective in Appendix A. here we also apply discretization to derive the objective of NFBM.
The conditional reverse process (eq. (17) and it's discretization (eq. 199 stays the same, except now
the conditional reverse process are conditioned on pair of points = (Xo. x1), sampled from
joined data distribution (Xo, x1) ~ g(xo, x1). 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) X1 = p(z (x1) and at the end samples from p(xp 20) = p(xol20).
We derive the variational bound on NLL of NFBM:
Bxx(x))
(31)
/
E Eg(xa)
log
((xollo)Po.ooo (Zo11 (1 ( dx
(32)
((x0120)
a Zp:1
At Zn (o
xx))(x))
Eggxa)
log Egixi 2on (xo)
(33)
Thp [Z0:1
2 Xo Xi 700 (2) No, xi (1(())(())))))
Z0. p
or
7 (z)
EEg(xo
log p(xo Zo)
(34)
(20:1-1-1 Xo,
x
=Erve + Laitt + Eprior + const,
where
(35)
Eree = Rggienzz) [- log (xoo(x))
(36)
Laitt
Σ (ao(e [KL (1p (Z.- Xo, xo. (2 #
(37)
Eprior
Ex(xx0x))) DQL (x
(38)
const = Ey(xo) [DKI. (x(x X [x)))(
(39)
Here, the reconstruction term Lree = Lree and the prior term Lprior = Eprior do not depend on discretization, and the constant term does not depend on model's 's parameters. For the diffusion term
16

---

Lair 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 Laitt transforms to:
Lai (x0xx))
f (Z,, xo: x1) fo. (z.
(40)
(t)
where u(t) is a uniform distribution over unit interval.
This gives us a variational bound on the negative log-likelihood of the NFBM, with very similar structure as the NFDM objective. Additionally, for NFBM we also may omit optimization of Cee and prior Cprior when 9p (Zo] Xo, x1) 8(x Zo) and q, (z) Xo, x1) ss p(21/x)))
B Details of NEDM
B.1 Motivation of Generalized Forward Process
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 Z0:1 (for simplicity, let's consider discrete-time model with some T time steps). However, the variational objective is equal to:
know from (30
DKL (4pp x, Z0:1 Olpo (x, 20:1) DKL. (1(x))lpo(x)) DRL (1p(Zaa (x) (x))
(41)
objective
entropy likelihood
posterior error
Thus, by minimizing the variational objective in diffusson models, 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 qup (Zo:1 xx) to match the model's posterior distribution (o(20.1/2 xx).
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 match the reverse process. This approach simplifies the task of matching the data distribution by the reverse process.
B.2 Calculation of the Time Derivative
As discussed in Section 3.1 the definition of forward process requires access to the time derivative of the function Fg. This derivative crucial for constructing a conditional ODE (see eq. (8)
fo x,) OF (F,(e,t,x
(42)
ØL
E= (z.
First, let clarify this notation: to find the ODE drift term fe at point 2, and time i conditionally x, we first apply the inverse function Fol to identify e, which corresponds to 2g, t,, and Subsequently, we compute the partial derivative of F,p with respect to 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.T have time derivatives available in closed form. However, in the general case, the time derivative is not readily available.
Fortunately, we can employ automatic differentiation tools like PyTorch 42 JAX (5). The naive approach to determining the time derivative involves running backpropagation procedure for each output of Fe and reconstructing each coordinate of the time derivative. Unfortunately, this method is not scalable for high- -igennimensional Fg-
Instead of backpropagation, which corresponds to Vector Jacobian Product (VJP), forward differentiation mode, or Jacobian Vector Product (JVP). By fixing the inputs e and x of the function F.p and focusing solely the input I, the Jacobian of F. appears as vector. Consequently, computing a JVP with a unit one-dimensional vector precisely produces the time derivative of F.g.
use
This approach allows for scalable computation of time derivatives.
17

---

B.3
Calculation of the Conditional Score Function
In Section 3.1
we define the forward process as conditional SDE (eq. ( dz, f. (zz, 1, x) + ()) where
(43)
f (z,, t,x) f (Z,, gZ((t)
2. log (z,
2
To construct this SDE, access the conditional score function Vz, log qup (z) (x) required. This section discusses methods for calculating this score function.
First, consider the case where Fip is linear in (further details this parameterization discussed in Appendix B:5):
Fo(E,t,x) (x (x,
(44)
where Me R. 0. R2 and de Rd [0. + RD, with the product (x,( being dlmenttwwise Consequently, the marginal probability densities of the forward process are Gaussian:
(z(z) x) = ((z (x, f), og (x,
(45)
In the case of Gaussian distributions, the score function can be easily found as:
(x, -x
E
V log (z((x)
(46)
(x,
ox
t) e= (2.,1.x)
Now, consider the general case where the distribution 9,p /x) is defined implicitly through an invertible transformation F,, of the random variable & ~ q(E) into 2t. To find the score function, we apply the change of variable formula:
OF (2t, x)
].
Zz, log q (2,x)) Vz,
log ((E)
+ log
(47)
F5 (x..1x)
Oz,
The first part of the equation is the log-density of the noise distribution, which is straightforward to calculate. The second part the log- dog-determinant of the Jacobian matrix of the inverse transformation. log |J If we have access log | formula and subsequently compute the gradient using automatic differentiation tools such as Py PyToch
calculate log gyp(z, x) using the change of variable
42 or JAX [5)
The log-determinant log |J can easily calculated when Fe linear in E, a case already considered above. For nonlinear, low-dimensional F.,, the Jacobian Jp' can be computed using O(D) backpropagations, and the log-determinant subsequently calculated. However, this approach is not scalable for high-ding digenssinni cases. It is challenging to compute the log-determinant log 1 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 degemmnermsssns..
Nevertheless, the forward process can parameterized with functions F& that inherently provide access to the log-determinant log [Jp') by design. For example, normalizing flow architectures 14 22 facilitate the construction of forward processes with non-Gaussian learnable distributions 4. (Ze x).
B.4
Stochastic and Deterministic Sampling
In Section 3.2 we define the reverse process as the following SDE (see eq. 1II)
dz; = fe.o(Ze, t)dt + (t)dw, where ferp (Ze, = f. (ze, t, (e t)).
(48)
Let us now rewrite this equation, substituting the definition of jß from 100:
(0((),))) (t) Vn, log gr (2.180(2e)))) di + (l)d..
dz. 1 (,
(49)
18

---

This expression represents the full form of the generative SDE, for which we derive the objective in Appendix/A.1 A.I 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 very simple dependence (t),
which determines the level of stochasticity. Therefore, in practice, we can adjust the stochasticity
level by modifying the function go
In the extreme case where 94 (t) = O, the generative SDE completely loses its stochasticity and becomes an ODE:
dz. for (Ze, L, 8o I))
(50)
This ODE allows for an interpretation of the model as Continuous Normalizing Flow (CNF) consequently enabling deterministic sampling from a pre-trained model and likelihood estimation.
6 17
B.5 Parameterization of the Forward Process
The parameterization of the variance function gp(() (eq. (9) and the data point predictor & (zi, 1) (eq. (11)) straightforward. In our experiments, we parameterize them directly using neural networks. However, parameterizing the transformation P. (eq. 77) is more complex. NFDM requires Fg to be differentiable and invertible with respect to E, as well as provide access to the log-determinant of the Jacobian matrix (see Appendix B.3). the forward process of NFDM.
We propose two parameterizations for
First, we propose parameterization of F that is linear in E:
E,(ett,x) pup (x, t) (x, where (x t) = (1 - t)x + t(1 tx) t), (x(x 5² (0 (xx (1-()
(51)
(52)
(53)
Here, A4g : RR >2 [0 RD and 3.p : -RP>(00 1] + R.P. with ö as constant. This parameterization
models q(2e)1 (x) as a conditional Gaussian distribution:
4p (2, x) =
Me Mx ox,
(54)
In our experiments, we set 8² = 10 10-4 If Pa is differentiable and o is differentiable and non- negative, this parameterization meets the requirements on Fig.
The functions tx, and o xx(xx) t) are specifically designed such that M'p 0) x and (x, t) = 8, and at = 1, they become (x,(,,1 1) = O and (x, = T. At intermediate times < f < 1, these functions can take arbitrary values, thereby ensuring that q(Za x) (N(2;; x, 5² and q(z xx) (2z,; 0, 7), with arbitrary Gaussian distributions at inter- mediate steps. These restrictions are not necessary for Fg, however, eliminates the need to optimize the reconstruction loss Lree and the prior loss Eprior, they become independent of the parameters &p and (see Appendix A.I
they equal
Secondly, we propose a more general parameterization:
xx(e, = (1 - 1)x (8+ (1-5)) F). f)x),
(55)
where Fg : RD x [0,1]x R is an invertible and differentiable function that allows access to the logedeterminant the Jacobian matrix. If F.p is sufficiently flexible (e.g., a neural network with a normalizing flow architecture), it can parameterize non-Gaussian distributions qu (z.
The additive term (1 - 1)x provides an inductive bias that shifts the mean the distribution go (z, /x) towards x when f. is near O. The linear function 8 + (1 8)t scales the distribution from a factor of 8 at f = 0 to 1 at = 1, also imparting an inductive bias for increasing variance over time. The scaling (1 f)x in the third argument of F. ensures that g,p (z||x) |x) at t = 1 does not depend allowing the reverse process to start from p(z1) = q (z, 00) and thereby guaranteeing matching of prior distributions between the forward and reverse processes. Consequently, this eliminates the need for optimizing the prior loss Cprior. However, unlike the previous parameterization that guarantees q(2o x) N (Zo; x, 52 T), in this case, there are no guarantees on the form of q(2g) xx), necessitating the optimization of the reconstruction loss Lree with respect to parameters &p (see Appendix A.I))
19

---

Importantly, even the Gaussian parameterization of F. in eq. (11 offers significantly more flexible forward process compared to prior works. Although linear in E, it may exhibit complex non-linear dependencies on and I.
Furthermore, we emphasize that the parameterizations of P.p, (t), and R (ze, t) discussed here 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.
C Details of NFBM
C.1 Parameterization of the Forward Process
To parameterize the forward process of NFBM, we adopt the Gaussian parameterization discussed for NFDM in Appendix B.5
(E, Xn, x1) = (xp:
x1 x .) + (X.,
where
(56)
Mep (Xo, X1, t) ()- tx1 + t(1 - (X0,
(57)
6 (xo, x1, 1) = 8 (ā x1, x1, ())(
(58)
Here, fhp RD x 00. 1] + I and o RD R [0, 1] + RD, with 82 10 in our experiments. Similar to the NFDM, this parameterization models q(z,|xx (x) as a conditional Gaussian distribution:
2: (2 x1) = -Nize Mup (xo, x1, x1, 02, t)).
(59)
In contrast to the parameterization of NFDM, here (xp 1) and o 1) are designed to ensure that ((z Xo: x = N(zo: Xo, 821) and ( |x0, x1) = (Zz;; (1;0)), At intermediate times, < < 1, ( Xu, X1) may still be an arbitrary Gaussian distribution. Therefore, with this parameterization of NFBM, we do not need to optimize the reconstruction loss Lree and the prior loss Cprior. as they are independent of the parameters y and 9 (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 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.
C.2 Parameterization of the Reverse Process
In Section we define the reverse process of NFBM by the following SDE:
dz: Pe t)dt + (t)dw, where (2,, () = 1 (Ze, t, &0 0.1
(ze, ()).
(60)
Here, the function & (zz, t) predicts both Xo and X1. Assuming oo (2 (xo, x
Re 8(x
2o) and
go (2 [2,]|x0,x0 x1) px this formulation leads to the objective (see Appendix
C.I.:
Eu (1) xx.)
[ (zz, xo, x1) - fa (Ze, 1)
(61)
(1)
As previously mentioned, alternatively, we could define the reverse process as being fully conditional on the initial data point x1. This leads to the following SDE:
dz. fo dt + 9 (t)dw, where fe.p 1) f (ze, L, &g(Ze,, L, x1),
(62)
This approach is akin to classifier- free guidance, where each step of the generative process depends on conditioning X1. Under the same assumptions- 9 (zo Xo, x1) xe 5(x-2z) and ( | [Xo. x1) R p(z, (x)) this definition leads to 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):
this
C Eu(t)g(No (x x1)
|1 (Zr, t, Xo. x1) fo.. (zi, 2,2,,t,0 X
(63)
(t
20

---

Intuitively, the difference between these two definitions implies that the first reverse process, at point 2, time t, follows the average trajectory of forward processes passing through 3₁. Conversely. the second reverse process tracks the average trajectory of forward processes that pass through and terminate 21-
The second approach may be beneficial when dealing with joint data distributions (Xx. x1). However, in this paper, we work with unpaired distributions q(xo) (x1). we have opted for the first approach in Section 4
D
NFDM with restrictions
D.I Additional Discussion of NFDM with Restrictions
In Section | 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 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 aim to learn a reverse process with some specific properties. To achieve this, 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 [. 11.. will lead to biased solutions, as is impossible to match two processes when
fl 112 = O.
when restriction is feasible, if the model parameterization is too simplistic, may not be able to match processes and satisfy the restriction, leading to bias. However, in our experiments, 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.
Additionally,
D.2 Details of NFDM with Curvature Penalty
In Section 5 we suggest penalizing the curvature of deterministic generative trajectories specifically with the Lerv (eq. (16)) penalty.
To calculate the curvature penalty Lerv (eq. (6)),, we proceed as follows. First, the chain rule to rewrite the time derivative:
(2.,t t)
Dlo.p(z. Oz
(64)
dt
Oz. Øt
Ot
a 0 1) Oz.
0So.(2n.1)
fo (Ze, 1)
(65)
Ol
The second term in eq. 665
the time derivative of function. As discussed in appendix B.2 the time
derivatives can be determined as a Jacobian Vector Product (JVP) with one-dimensional unit vector.
21

---

Notably, the first term in eq. (65) is also JVP. Therefore, we can combine these two operations. For this purpose, we define RPFT dimensional adjoint vector 2:
[0,0(20))) 11)
(66)
Consequently, Lerv can be efficiently computed as the JVP of (z,p(z,,t) 1) with the vector v.
Additionally, would like to note couple of aspects of learning the NFDM with curvature penalty.
First, unlike the regular NFDM, the curvature penalty necessitates that Fg be twice differentiable
with respect to time. Second, although we parameterize fors (eq. (11) through f (eq. (10)) fery influences the forward process but does not penalize the curvature of the conditional forward trajectories (see illustrations in Appendix G.I..
E Connections of NEDM with Other Approaches
In this section we continue discussing connections between NFDM and related works.
E.1 Special Cases of NFDM
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. [ consider the case of a fixed forward processes qzz x) = (zz; C2x, o,1), which can be reparameterized F(e,t, x) C,x + ø,s. Such distributions include, for example, combinations of blurring and the injection of Gaussian noise.
Star-Shaped Diffusion. [4]] extended conventional diffusion models to include distributions from the exponential family. Although Star-Shaped Diffusion a discrete-time approach and does not directly correspond to NFDM, the latter can work with exponential family distributions through reparameterivation functions. For instance, for some continuous one-dimensional distribution g(2 [x). NFDM could use tt,,t, x) (f,()),,,)) where is the inverse Cumulative Distribution Function (CDF) of q(2 x), and is the CDF of a unit Gaussian.
[28] proposed forward conditional distributions (Z, as N(zz; (t)x, of (t) with learnall parameters In the context of NFDM, this can be real- ized by Fg (s, x, = (x())x (t)E.
Variational Diffusion Models.
LSGM. [66] suggest alternative approach for parameterizing the forward process, proposing diffusion in a latent space of a VAE. Therefore, the forward process is characterized by distribution (,(2)/ x) = N(z,; at Exp I)., where E.p is the encoder of the VAE. To parameterize the same forward process with NFDM, one could use (e x) 0u E (x) + o To align the reverse process, the reconstruction distribution should p(x(zo) = (x; D Zo) 821), where D is VAE's decoder.
ShiftDDPM. proposed learning a function that shifts the mean and covariance of (z(2)))) x) by conditional information. In our work, we do not focus on conditional generation. However, we believe that the NFDM framework can readily adapted for conditional generation by incorporating the conditioning variable into both the transformation (,, c) and the predictor X(Zz, E, c).
NDM and DiffEnc.
proposed more general forward process, (z [x] ((z t), o? which, unlike LSGM, transforms in time-dependent manner. This forward process can also be described in terms of NFDM as the special case F,(e,t, = (x, + o.E.
40
Curvature minimization. 332 suggested learning the distribution (2, of the forward process at time step f = 1 conditional Gaussian. They parameterize the forward process via linear interpolation between x and 2ı. Within the NFDM framework, we can express this as F3 4, x, = (1 f)x + 1(1 (x) + a (x)e). Subsequently, the authors train their model by optimizing the Flow Matching objective in conjunction with minimizing the Kullback-Leible (KL) divergence between
22

---

go(Z1(x) (x) and p(z). This objective corresponds to minimization of diffusion term Lair and prior term ppior NFDM objective (see Appendix A.I for the forward process proposed by [32]
E.2
Compatible and Orthogonal Approaches
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 orthogonal line of works alternative sampling methods have been studied, such as deterministic sampling in discrete [57] and continuous time [6]] and novel numerical
solvers [62 35 38 53). Additionally, distillation techniques have been applied to both discrete [50] and continuous time models [58] [75 722 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-Frs 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.
E.3 Other Related Works
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.
Schrödinger Bridge models. Recent studies 73 10 54
Nevertheless, it is important 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.
44 and 664 proposed to construct the forward process with optimal data-noise couplings to learn straighter generative trajectories. While this method is theoretically justified, 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 the generative trajectories.
Miníbatch optimal transport.
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 acknowl- edge 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 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 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 with NFDM.
In their practical applications, Stochastic Interpolants are demonstrated with a simple parameteriza- tions 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
23

---

which they do not provide experimental results. Moreover, their work does not clearly articulate how Stochastic Interpolants might applied learning other generative dynamics.
In contrast, NEDM introduces a more generic method for learning generative dynamics. Moreover, when NEDM is learned with restrictions penalties (see Section [p. it remains within min-min optimization paradigm, which can be addressed more efficiently.
E.4 Connections of NFDM Objective
In this section, delve into the details of the NFDM's objective function (eq. will consider the diffusion term of the objective Laift (see Appendix[A.I]
12 Specifically
First, let's unpack the Lait by substituting the definitions of f and fo.p:
1
1 (2,,,,x) - -fo (t,,,t)) $11
Laift
E
(67)
uif (x)
(t)
(fo(2),t)x) (2 Xo (0 t))
E
(68)
(t)
Flow Matching term
(t(() 2
Vz, log Uu (2 Sep Ze. ()) Vz, log (2(2211x)) (x))
Score Matching term
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 (t). When gp((t) is small, the forward and reverse processes exhibit smoother trajectories, and the objective is dominated by the first term. Conversely, when 9 (t) is large, the processes exhibit more stochastic trajectories, and the objective is dominated by the second term.
Crucially, in the extreme scenarios where (() approaches either oo, the diffusion loss Laitt corresponds to either a reweighted Flow Matching loss [33 or reweighted Score Matching loss
68 respectively:
lim Laitt lim (x( (=(
(fp. f (z,, 30 (Z,, 1)]
(69)
(x)
29 (t)
(t)
lim Lai (t)=
lim
E (x(t)
|/ log 9p (2)1x0 Ve, log (z
8e 000
8
(70)
This connection highlights the importance of optimizing the full objective in eq. both Flow Matching and Score Matching components. Simplifying the objective to include only one component renders intractable from a variational perspective. Specifically, when ga (L) approaches either or oc, the objective, a variational bound on likelihood, becomes infinitely large.
67 which contains
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 (z,, x) and the conditional score function Vz, log 9 [ x) are merely linear combinations of 2, and x. As (27) demonstrated, for such models, Flow Matching and Score Matching be viewed reweighted ELBO. However, for more general non-linear forward process in NEDM, optimizing Flow Matching and Score Matching objectives is
equivalent.
Secondly, 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 either

---

Table 4: Training syper--arameters.
CIFAR-10
ImageNet 32
ImageNet 64
AFHQ 64
Channels
256
256
192
192
Depth
2
3
3
Channels multipliers
1,2,2,
1,2,2,
1,2,3,4
1,2,2,2
Heads
4
4
Heads Channels
64
64
Attention resolution
16
16,8
32,16,8
32,16,8
Dropout
0.0
0.0
0.0
0.0
Effective Batch size
256
1024
2048
1024
GPUs
16
16
Epochs
1000
200
250
100
Iterations
391k
250k
157k
102k
Learning Rate
4e-4
le-4
le-4
le-4
Learning Rate Scheduler
Polynomial 45k
Polynomial 20k
Constant
Constant
Warmup Steps
objective in the general case. In NFDM, we learn the forward process end-to-end with the reverse process. Thus, reducing the NFDM's objective could lead to the model's collapse, as not the correct objective to optimize.
We would also like emphasize that the NFDM framework, including its objective in eq. (67). represents a generalization of conventional diffusion models, rather than an alternative approach to training. Consequently, if take existing forward process from conventional diffusion models and incorporate it into the objective L, we will not create a new model or optimization procedure. Instead, this will align with approaches that focus on likelihood-baai training of conventional diffusion models, as explored in [59 76).
NFDM facilitates the convenient pre-specification of fixed forward processes. Hence, it might be utilized to define a fixed forward process and subsequently train with the Flow Matching objective. However, this is topic for future research and is beyond the scope of this work.
Implementation Details
Our evaluation of NFDM and NFBM includes tests on synthetic data, CIFAR-10 [31] two down- sampled ImageNet [T] 67 datasets, and downsampled AFHQ 8. To maintain consistency with baselines, we employ horizontal flipping data augmentation technique in training models on
CIFAR-IO and ImageNet 61 59] For density estimation of discrete data, uniform dequantization is
used (see Appendix
We parameterize Xo (eq. (1)) in the reverse process using 5-layer perceptrons with 512 neurons in each layer for synthetic data and the U-Net architecture from [12] for images. In all experiments, a 3 3-layer perceptrons with 64 neurons in each layer is employed to parameterize I.- (eq. (9)
For the parameterization of (eq.qq. (7 in NFDM as Gaussian, we employ a neural network identical to that used for Xg. The only difference is that for Fs, we double the output of the last layer to
parameterize both nh (eq. (52)) and a (eq. (53) using the same model (see Appendix [B.5).
For parameterizing Fg (eq. [13 ) in NFDM, we adopt the same approach but double the input size of the first layer condition the Torward process on both xo and x (see Appendix C.T.
To parameterize Fg (eq. (7) in NFDM as non-Gaussian, we apply the Glow (2 architecture tailored neural network identical to that X, to generate an embedding u(x, t). We then add this conditioning to the input of each non-linear transformation in the coupling layers.
to the corresponding dataset to transform into z₁. To incorporate conditioning,
To ensure the constraints dg 2 0 (eq. and ā 2 (eqs. (53) and (58) are met, we apply the
softplus function to the outputs of the neural networks.
25

---

The models were trained using the Adam optimizer with the following parameters: 81 = 0.9. 32 0.999, weight decay of 0.0, and 10 The training process was facilitated by polynomial decay learning rate schedule, which includes warm-up phase for predefined number of training steps. During this phase, the learning rate linearly increased from 10
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 Table Training was carried out on Tesla V100 GPUs.
F.1 Dequantization
For reporting the NLL, we employ standard uniform dequantization. The NLL is estimated using an importance-weighter average, given by
K
log K
po(x + uz), where ~ (1,00,)) 1).
(71)
1=1
where x E [O..... ,55]..
F.2
NFBM with Obstacles
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 points closer to obstacles. Specifically, applied the following objective:
Lobs ((())((z;)) (f
(72)
where f(z)) is penalty function. In our experiments, we use the probability density of mixture of Gaussian distributions, truncated some low values This truncation is necessary ensure that the model does not push trajectories infinitely away from the obstacle.
In contrast to Lerv proposed for NFDM-OT, which penalizes the reverse process influencing the forward process as a result, Lon directly penalizes the forward process, thereby influencing the properties of the reverse process.
G Additional Results
G.1
Visualisations of Forward Trajectories of NEDM-OT
In this section, provide visualizations of the forward trajectories of NFDM-OT, as illustrated in the two-dimensional data distribution experiment shown Figure [b] Figure depicts the deterministic trajectories of NFDM-ON OT. It is easy 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 NEDM-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 signifi- cantly enhances the flexibility diffusion models compared to conventional diffusion approaches, enabling the learning of complex forward processes.
G.2 Additional Results of NEDM and NFDM-OT Image Datasets
In this section we provide additional results, that directly compare NFDM and NFDM-OT, This results extend Tables and 2 As observed, NFDM without any additional restrictions provides
26

---

3
2
2. first coordinate
0.0
0.2
0.4
0.6
1.0
Figure 3: First coordinates of forward deterministic trajectories of NEDM-OT started from points (-1. and x (1, 1).
Table 5:
Comparison of NEDM with NFDM-OT results on density estimation tasks. We present results in terms of BPD, lower is better.
Model
CIFARI0
ImageNet 32
NFDM
2.49
3.36
NFDM-OT
2.62
3.45
Table 6: Comparison of NFDM with NFDM-OP -T results on few-step generation. We present results in terms of FID, lower is better.
CIFAR-10
ImageNet 32
Model
NFE
FID
NFE
50.12
57.60
NEDM
21.88
24.74
7.76
6.13
NFDM-OT
12
5.20
12
4.11
better ligeliiel estimation compared NFDM-OT. However, expected, NFDM exhibits much worse sample quality compared to NFDM-OT in few steps generation. We attribute this property the higher curvature NFDM's generative trajectories.
G.3 Generated Samples
In this section, we present generated samples from the NEDM, NEDM-OT, and NFBM models, trained on various datasets.
NFDM: Figure shows samples from the NFDM model with Gaussian qu (z.|x) (see Appendix B.5 We provide samples from NFDM trained on the CIFAR-10, ImageNet 32, and ImageNet 64 datasets.
NFDM-OT. Figure digulaasss
displays samples from NFDM trained on ImageNet 64 dataset. We generate these samples by simulating the ODE across 12, 4, and 2 steps.
NFBM. Figure 6 F demonstrates the generative trajectories of the NFBM model trained the AFHQ 64 dataset. This model starts from an image of a dog at time step /. = and simulate the SDE to time step f 0.
27

---

(a) CIFAR-10
(b) ImageNet
(c) ImageNet 64
Figure 4: Generated samples from NFDM trained on various datasets.
Figure 5: Generated samples from NFDM-OT trained ImageNet 64. Top: NFE 2; Middle:
NFE = 4; Bottom: NFE- 12.
t=0.0
t= 0.17
0.33
t=0.5
t=0.67
t=0.83
t=1.0
Figure 6: Generative trajectories from NFBM trained on AFHQ 64.
28

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our paper follows the structure of the abstract and introduction directly.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations in general in Section 8 and in Appendix B.3 in more
detail.
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 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]

---

Justification:
We provide proofs in Appendix A
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We provide detailed description of the model in Section and appendix
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
30

---

Answer:
[No]
Justification:
No new data has been generated, and the code will be made available late.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/pips..ccppullc//guui / Sow.llesss/oSSSeemiiseeee
wwt.Subbmiiiioooiiioooooooio for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The data configuration follows from previous work, while most of the experi- mental details are provided in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Computation of eтror bars is vary computationally demanding. We follow standard practice in the literature and report the average result.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence 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 capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall 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).
31

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
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 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 figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide compute resources details in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU 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).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
We read and checked the code of ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. 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 no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact 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.
32

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper does not pose such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 best faith effort.
12. Lieenses for existing assets
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?
Question:
Answer: [Yes]
Justification: All data and software used is properly credited and licenses respected.
Guidelines:
existing assets. The authors should cite the original paper that produced the code package or dataset.
The answer NA means that the paper does
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.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, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
33

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No assets introduced in the paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper did not require crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
No study participants were involved.
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
34


Reviews:
[{"review_id": "47ev2nT4HN", "reviewer": "Reviewer_xgai", "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.", "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.\n2. How do the authors ensure that the function $F$ is smooth between $z_0$​ and $z_1$​ (Line 120)? \n3. 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$ . \n4. 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. \n5. 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?\n6. 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?\n7. 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? \n8. 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? \n9. 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. \n\n[1]  Albergo, Michael S., and Eric Vanden-Eijnden. \"Building normalizing flows with stochastic interpolants.\" arXiv preprint arXiv:2209.15571 (2022). \n\n\n[2] Kim, Dongjun, et al. \"Maximum likelihood training of implicit nonlinear diffusion model.\" _Advances in neural information processing systems_ 35 (2022): 32270-32284.", "limitations": "The limitations are addressed in the paper.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is well-written and easy to follow. \n- The problem is conceptually well-motivated and interesting.\n- The introduced framework is clearly explained with all the necessary details.\n- Development of a new optimization framework preserving simulation-free.\n- Demonstration of state-of-the-art results on CIFAR-10, ImageNet 32, and 64.\n- Application of NFDM and NFBM with specific properties and improved sampling speed.", "weaknesses": "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. \n1. It is unclear how the introduced method would scale to larger datasets and image resolutions and what would be the benefit.\n1. The non-Gaussian parameterization does not show substantial improvement over the Gaussian parameterization, raising questions about its necessity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's positive feedback regarding the clarity and presentation of our paper. Below, we address the questions and comments raised in the review.\n\n**Weaknesses:**\n\n1. Please refer to the answer to Question 1 below.\n2. Unfortunately, due to limited computational resources, we were not able to conduct experiments on a significantly larger scale than those on ImageNet 64. Consequently, we cannot definitively ascertain the outcomes in a large-scale setting.\n    \n    Nevertheless, we would like to highlight that NFDM is a versatile framework that includes conventional diffusion as a special case which has been shown to scale. Furthermore, CIFAR-10 and ImageNet are standard benchmarks in the diffusion literature. Although we lack large-scale experiments to support our position fully, we see no reasons that would impede the scaling of NFDM.\n    \n3. Please refer to the answer to Question 6 below.\n\n**Questions:**\n\n1. The main contribution of our work is the neural flow diffusion model, introduced as a general framework.  It enables the construction of diffusion models with a broad spectrum of forward processes, including *non-linear*, *time-dependent*, and *learnable processes*. We show how to efficiently parameterize and train diffusion models with such forward processes, while retaining key properties of diffusion models, such as access to stochastic and deterministic generative dynamics as well as simulation-free training.\n    \n    In Appendix E, we discuss in more detail the connections between NFDM and related works that introduce learnable parameters into the forward process. Most existing prior work can be considered special cases of NFDM. For instance, to parameterize the forward process in [1] using NFDM, one could set $F_\\phi(\\varepsilon, t, x) = \\alpha_t E_\\phi(x) + \\sigma_t \\varepsilon$, where $E_\\phi(x)$ is an invertible transformation.\n    \n    We emphasize that NFDM is not merely an alternative approach, but a generalization of prior diffusion models. With a simple parameterization such as $F_\\phi(\\varepsilon, t, x) = \\alpha_t x + \\sigma_t \\varepsilon$, NFDM directly corresponds to DDPM and introduces no additional computational overhead. The computational cost of NFDM depends on the complexity of the forward process' parameterization, which lets the user interpolate between standard linear transformations to highly complex nonlinear ones.\n    \n    To the best of our knowledge, NFDM provides the most general diffusion modelling framework that allows learning the forward process while keeping the training procedure simulation-free. Besides achieving state-of-the-art log-likelihood modelling results, the flexibility of NFDM enables us to learn generative processes with specific properties like NFDM-OT, for which it is crucial that the forward process has a complex non-linear time-dependent forward processes (see Appendix G.1).\n    \n2. To ensure that $F_\\phi$ is smooth in time, we chose a parameterization that is continuous and differentiable. These requirements are standard for the neural networks and are not difficult to meet. We discuss the parameterization of $F_\\phi$ in more detail in Appendix B.5 and implementation details in Appendix F.\n3. The linear parameterization in $\\varepsilon$ of $F_\\phi$ is as follows: $F_\\phi(\\varepsilon, t, x) = \\mu_\\phi(x,t) + \\sigma_\\phi(x,t) \\varepsilon$, where $\\mu_\\phi$ and $\\sigma_\\phi$ are learnable functions. In terms of NFDM, we may describe the forward process from [2] as $F_\\phi(\\varepsilon, t, x) = (1 - t) x + t \\varepsilon$. In [3], the authors discuss other variations of linear combinations of $x$ and $\\varepsilon$, which can also be easily expressed within the NFDM framework. A key distinction between NFDM and [2, 3] is that the Stochastic Interpolants framework does not propose a method to learn the forward process. Their derivations rely on a fixed forward process, while NFDM allows for learnable parameters in the forward process.\n    \n    Indeed, when we parameterize the forward process of NFDM with learnable parameters, we need to learn both the function $F_\\phi$ and the predictor $x_\\theta$. In Appendix E, we referred to the fact that in [3], with a fixed forward process, two different functions need to be learned for inference, one for the velocity and one for the score. In contrast in NFDM, if we fix the forward process (i.e., fix $F$) like in [3], we only need to learn the predictor $\\hat{x}_\\theta$.\n\n[1] Kim, Dongjun, et al. \"Maximum likelihood training of implicit nonlinear diffusion model.”\n\n[2] Albergo, Michael S., and Eric Vanden-Eijnden. \"Building normalizing flows with stochastic interpolants.”\n\n[3] Albergo, Michael S., Nicholas M. Boffi, and Eric Vanden-Eijnden. \"Stochastic interpolants: A unifying framework for flows and diffusions.”", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "emzMxxrUss", "reviewer": "Reviewer_1ZLG", "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.", "questions": "See weaknesses", "limitations": "Adequately addressed", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- An elegant proposed framework for generalizing diffusion \n- I think that learning modifications to the forward process is a potentially very impactful research direction and the ability to express constraints on it opens up new avenues for exploration\n- State-of-the-art density estimation results", "weaknesses": "- 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\nlearns to predict, the direction stepped in during the reverse ODE, etc.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are pleased that the reviewer appreciates the elegance of our framework and recognizes the potential impact of our research direction. In response to the comments and questions raised, we provide the following clarifications and commitments:\n\n**Weaknesses:**\n\n1. We acknowledge the lack of visual explanations in our paper. Due to space constraints in the main manuscript, we unfortunately did not include them. In the camera-ready version, if accepted, we plan to include a series of figures in the main paper or the appendix. Specifically, we intend to add a diagram that visually explains how the $F_\\varphi$ function defines the conditional distributions $q_\\varphi(\\mathbf{z}_t|\\mathbf{x})$, and illustrates the connection between ODEs and SDEs corresponding to the forward process. \n    \n    Additionally, in our global response, we have provided visualizations of generated samples from the NFDM and NFBM models trained on various datasets.\n    \n2. We have provided a comparison between NFDM and NFDM-OT in Tables 5 and 6 in the appendix. Compared to NFDM, NFDM-OT trades off worse density estimation performance (lower log-likelihood) for better sample quality (lower FID) in the few-step generation mode.\n    \n    In the same 2D experiment illustrated in Figure 1, NFDM exhibits slightly less curved trajectories compared to conventional diffusion. We will include additional illustrations in the camera-ready version, if accepted.\n    \n\nWe believe that the additional visualizations and comparisons will strengthen your support for acceptance. We are also happy to address any further concerns or questions the reviewers may have.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cwIgFqZgeW", "reviewer": "Reviewer_AXnA", "summary": "This work proposes a generic framework to make diffusion model's encoder learnable.", "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?\n\n2. How $g_\\psi$ be determined? How it depends on $\\psi$?\n\n3. 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?\n\n \n\n[1] S. Lee, B. Kim, and J. C. Ye. Minimizing trajectory curvature of ode-based generative models. arXiv preprint arXiv:2301.12003, 2023.", "limitations": "The authors have discussed the limitation that using a learnable forward path may double the training runtime.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The authors propose a neat framework to make the encoding part of the diffusion model trainable. They also explain the differences and relationships between their proposal and related works.", "weaknesses": "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. \n\n2. 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.\n\n[Minor]\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are glad the reviewer found NFDM framework as neat. Below, we address the questions and comments raised in the review.\n\n**Weaknesses:**\n\n1. We define the distribution $q_\\phi(z_t|x)$ implicitly through the function $F_\\phi(\\varepsilon, t, x)$. Under certain assumptions, we can derive an ordinary differential equation (ODE) as shown in Equation 8 by simply differentiating $F_\\phi$ with respect to time. It is straightforward to see that integrating this ODE over time produces the same samples as those generated by $F_\\phi$. Consequently, the ODE in Equation 8 corresponds to the defined probability density $q_\\phi(z_t|x)$.\n    \n    The connection between deterministic and stochastic dynamics through the score function, as leveraged in Equations 9 and 10, is a classical result from [1]. Further proofs demonstrating that the dynamics discussed in Equations 8-10 share the same marginals  can be found in [2] (see Appendices A and D.1).\n    \n    We will include a more extensive discussion of these derivations in Appendix A in the revised version of the paper, if accepted.\n    \n2. Consider the distribution of latent variables $q(z_t)$ defined by the forward process. Given the data distribution $q(x)$ and the conditional distribution $q(z_t|x)$, we can derive the distribution of latent variables $z_t$ as follows: $q(z_t) = \\int q(x_t|x) q(x) dx$. Therefore, if the forward process that defines $q(z_t|x)$ is fixed, $q(z_t)$ is also implicitly predefined and fixed. At the same time, the generative process must adhere to all intermediate distributions $q(z_t)$ starting from $t=1$. This is why a fixed forward process limits the flexibility of the latent space.\n    \n    From a theoretical perspective, conventional diffusion with a fixed forward process that only injects Gaussian noise is a model that in the ideal scenario works for arbitrary distributions. However, conventional diffusion can be suboptimal, potentially leading to distributions of latent variables that are difficult to learn or highly curved trajectories (Fig. 1) that are expensive to integrate. A learnable forward process could simplify the target for the generative process, thereby improving the alignment of forward and generative processes and better capturing the data distribution as a result.\n    \n    We provide empirical evidence that learning the forward process enhances the performance of the generative process. A theoretical motivation for learning the forward process is also detailed in Appendix B.1.\n    \n    Additionally, the flexibility of the NFDM framework allows us to learn not just any generative process, but one with desired properties such as straight-line deterministic trajectories (see Section 5), which is unachievable with conventional diffusion that predetermines the distributions of the latent variables.\n    \n3. We are pleased that the reviewer appreciated our discussion of connections with related works. We will elaborate on the differences with these approaches in more detail in the revision of the paper, if accepted. We note, however, that in most cases, the parameterizations of NFDM discussed in Appendix E.1 lead to the same objectives as proposed in related works.\n\n**Questions:**\n\n1. In our experiments, we intentionally parameterize the forward process to ensure that $q_\\phi(z_0|x) \\approx \\delta(x - z_0)$ and $q_\\phi(z_1|x) \\approx p(z_0)$ (see Appendix B.5). In this way, the model can not collapse by design. For instance, in Eq. 12, $\\tilde{f}_\\phi(z_t, t, x)$ can not always equal $0$. Furthermore, the NFDM objective is also a variational bound on the log-likelihood (Appendix A.1).\n    \n    In practice we did not observe any instabilities in the training procedure and did not apply any techniques to stabilize training.\n    \n2. We describe the parameterization of $g_\\phi$ in Appendix F. In all our experiments, $g_\\phi$ is parameterized using a 3-layer perceptron with 64 neurons in each layer.\n3. When sampling with an adaptive ODE solver from models trained on the CIFAR10 dataset, we achieve an FID score of 4.22 for NFDM (with an average of 82 NFE) and 4.85 for NFDM-OT (with 45 NFE on average). For reference, a DDPM trained with the ELBO objective has an FID score of 13.51 (with 1000 NFE). We plan to calculate the FID metrics for the ImageNet dataset by the end of the discussion period.\n    \n    Some modifications of diffusion models trained with better parameterizations and different objectives can achieve lower FID scores, but they tend to show worse log-likelihood estimates. This phenomenon is a well-known attribute of diffusion models. We believe that with improved parameterization and hyperparameter- and objective tuning (which was done for conventional diffusion), the flexibility of NFDM could lead to even better performance in terms of FID score, but we leave this for future research.\n    \n    Although we have provided experimental results on image datasets as a standard benchmark, we introduce NFDM as a versatile generative modeling framework. In this context, likelihood emerges as the most commonly employed metric in probabilistic modeling across various fields, and NFDM achieves state-of-the-art results. While the FID score is a popular metric for images, it does not always correlate with image quality [3]. In contrast, likelihood is based on robust theoretical principles. We will include a table with these FID scores in the appendix.\n    \n\nWe trust that the clarifications further explanations of motivation, technical details and other will enhance your support for our paper's acceptance!\n\n[1] B. D. Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313–326, 1982.\n\n[2] Song, Yang, et al. \"Score-based generative modeling through stochastic differential equations.”\n\n[3] Stein, George, et al. \"Exposing flaws of generative model evaluation metrics and their unfair treatment of diffusion models.\"", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1N8pgtSqWG", "reviewer": "Reviewer_DMsJ", "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.", "questions": "Please see weaknesses.", "limitations": "The authors have adequately described the limitations and potential negative societal impact of their work.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The paper introduces the concept of a learnable forward process in diffusion models, which is a novel approach. Traditional diffusion models typically rely on fixed forward processes, and the shift to a learnable process represents a significant departure from established methodologies. This innovation opens up new avenues for research and development in generative modeling.\n2. The paper provides a solid theoretical underpinning for the proposed NFDM. The authors thoroughly explain the mathematical principles behind the learnable forward process and its optimization. The experimental validation is robust and extensive. By demonstrating state-of-the-art results on widely recognized benchmarks such as CIFAR-10 and ImageNet, the paper establishes the practical effectiveness of NFDM. The experiments are well-designed and the results are presented with appropriate statistical rigor.\n3. The paper is well-written and clearly structured, making it accessible to both experts and those new to the field.\n4. By introducing a learnable forward process, the paper makes a significant contribution to the field of generative modeling. This advancement has the potential to influence future research and applications in areas such as image generation.", "weaknesses": "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.\n2. 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.\n\n[1] Shiftddpms: Exploring conditional diffusion models by shifting diffusion trajectories", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are delighted to see that the reviewer finds our approach novel, the theoretical explanations solid, and the experiments robust, and even recognizes the potential to influence future research. Below, we address the questions and comments raised in the review.\n\n**Weaknesses:**\n\n1. We thank the reviewer for bringing the Shiftddpms [1] work to our attention, which we had overlooked. While Shiftddpms focuses on conditional generation, it indeed shares similarities with our NFDM, proposing learnable transformations of conditioning. Although we do not discuss conditional generation in this paper, we believe that the NFDM framework could be adapted for conditional generation and this is an interesting avenue for future research.\n    \n    We will definitely include a discussion of this work in the revision of the paper if it is accepted.\n    \n2. In this work, we aimed to focus on the introduction of the NFDM itself and demonstrate its capabilities compared to prior work on standard benchmarks used in those works. Indeed, in the current work, we do not discuss topics such as the application of NFDM to domains other than imaging, scaling, conditional generation, or training with non-ELBO objectives. However, we view these as very promising directions for future research.\n\n[1] Shiftddpms: Exploring conditional diffusion models by shifting diffusion trajectories", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their valuable feedback, which helps us to improve the paper. Below, we have responded individually to the specific comments and questions from each reviewer.

Here we would like to provide some visualizations of samples generated with the NFDM and NFBM models trained on different datasets. Please see the attached PDF file.

We hope that the additional visualizations, along with the clarifications and explanations provided below, will strengthen your support for the acceptance of our paper!

Meta-review: 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.
Decision: Accept (poster)