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.
Paper
Similar papers
Peer review
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.
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. 2. 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. 3. The paper is well-written and clearly structured, making it accessible to both experts and those new to the field. 4. 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. 2. The paper does not extensively discuss the generalization of the results to other settings or the robustness of the results to violations of the assumptions made. This omission can be critical as it limits the confidence in applying the findings to different datasets or under different conditions than those explicitly tested. [1] Shiftddpms: Exploring conditional diffusion models by shifting diffusion trajectories
Questions
Please see weaknesses.
Rating
8
Confidence
4
Soundness
4
Presentation
4
Contribution
4
Limitations
The authors have adequately described the limitations and potential negative societal impact of their work.
Summary
This work proposes a generic framework to make diffusion model's encoder learnable.
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. 2. The authors claimed: "*conventional diffusion the latent variables are inferred through a pre-specified linear combination of the data point and Gaussian noise. This formulation limits diffusion models in terms of the flexibility of their latent space, and makes learning of the reverse process more challenging...*". However, is there any evidence showing that a diffusion model's linear combination of data and noise may lead to degradation or affect the flexibility of the latent space? Also, if the forward path is simply a pre-specified linear combination of the data point and Gaussian noise, wouldn't it result in a more easily learnable backward process? I suggest the authors provide either theoretical or empirical evidence to support their claim and motivation. [Minor] - I appreciate that the authors explained the differences and relationships between their proposal and related works, which is informative. I suggest extending Appendix E.1 to include more comparisons of each method's assumptions, such as the priors used, their final objective functions, and why their parameterizations fall short.
Questions
1. The proposed algorithm suggests learning the encoding and decoding parts together (though alternately). However, the loss function shown in Eq. (12) seems prone to collapsing to a trivial solution, despite the parametrizations potentially mitigating this issue somewhat. Indeed, some literature, such as [1], includes additional regularizers to stabilize the training. Did the authors observe any collapse or instability during training? Any trick to avoid the collapse if it is the case? 2. How $g_\psi$ be determined? How it depends on $\psi$? 3. What is the FID performance of NFDM (with/without curvature regularizer) if the total sampling timestep $T$ is set to the standard diffusion model value of $1000$, compared to DDIM or DDPM? The authors claimed diffusion model's pre-defined encoding may lead to inflexible and complicated latent for reverse process. The training objective is derived in a favor for likelihood computation. However, sample quality should be a more important metric to justify the claim. Also, what does the counterpart of Fig. 2 without curvature regularizer? [1] S. Lee, B. Kim, and J. C. Ye. Minimizing trajectory curvature of ode-based generative models. arXiv preprint arXiv:2301.12003, 2023.
Rating
5
Confidence
3
Soundness
2
Presentation
3
Contribution
3
Limitations
The authors have discussed the limitation that using a learnable forward path may double the training runtime.
Thanks for clarification and replies.
I appreciate the reviewers' clarification and have decided to increase my score.
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.
Strengths
- An elegant proposed framework for generalizing diffusion - 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 - 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 learns to predict, the direction stepped in during the reverse ODE, etc. - It would be interesting to see comparisons between NFDM and NFDM-OT included in the paper. E.g. if both of them could be included in each of the results tables. It would be very interesting to see how they compare FID and NLL, and also how the learned trajectories of NFDM compare to those of Fig 1a.
Questions
See weaknesses
Rating
7
Confidence
3
Soundness
3
Presentation
2
Contribution
3
Limitations
Adequately addressed
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.
Strengths
- The paper is well-written and easy to follow. - The problem is conceptually well-motivated and interesting. - The introduced framework is clearly explained with all the necessary details. - Development of a new optimization framework preserving simulation-free. - Demonstration of state-of-the-art results on CIFAR-10, ImageNet 32, and 64. - 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. 1. It is unclear how the introduced method would scale to larger datasets and image resolutions and what would be the benefit. 1. The non-Gaussian parameterization does not show substantial improvement over the Gaussian parameterization, raising questions about its necessity.
Questions
1. Can the authors provide a clear statement of their contributions and methods within the context of related works that also aim to extend the diffusion framework to include a learnable and potentially non-linear forward process (e.g. [2])? Specifically, it would be helpful to highlight the advantages and disadvantages of these methods compared to theirs, from both methodological and computational perspectives. 2. How do the authors ensure that the function $F$ is smooth between $z_0$ and $z_1$ (Line 120)? 3. Can the authors clarify how the linear in $\epsilon$ parameterization of $F_\varphi$ relates to the stochastic interpolant framework in [1]. Such comparison is already outlined in Appendix E where the authors point out the different optimization problems (min-min vs min-max) and state "*Stochastic Interpolants necessitate learning two separate functions for the reverse process: the velocity field and the score function*". Could the authors clarify the latter statement? It seems that also in the proposed NFDM framework one would need to learn both the $F_\varphi$ (which moreover needs to be derived wrt to time and invertible wrt to $\varepsilon$) and the prediction network $\hat{x}_\theta$ . 4. For clarity of exposition, it would be beneficial to better introduce the function $\hat{x}_\theta$, particularly explaining the difference with the score network in conventional diffusion. 5. Why is the framework still considered simulation-free despite requiring a prediction network $\hat{x}_\theta$? Can the authors explain why this step does not imply that the model, parameterized by $\theta$, implicitly learns to simulate the process? 6. In Table 1, there appears to be no significant improvement when NFDM uses a Gaussian parameterization compared to a non-Gaussian one. Could the authors comment on this result? 7. In App B.3, the authors also mention that a non-Gaussian parameterization is not scalable for high-dimensional cases. Have the authors observed any case where a non-Gaussian parameterization leads to an improvement? 8. Can the authors offer a more detailed comparison of their _Restricted NFDM_ compared to previous works that also constrain the generative process to have straight trajectories? Specifically, what are the computational advantages/disadvantages? 9. Can the authors better discuss the computational cost of NFDM compared to previous approaches? If I am not mistaken I only see a mention in the conclusion "*this leads to approximately 2.2 times longer optimization iteration of NFDM compared to conventional diffusion models*" but it is not clear to which NFDM parameterisation this estimate refers. Moreover, while Table 2 depicts a clear advantage of NFDM in terms of NFE, it would be useful to analyze these numbers in the context of the computational complexities of the respective methods. [1] Albergo, Michael S., and Eric Vanden-Eijnden. "Building normalizing flows with stochastic interpolants." arXiv preprint arXiv:2209.15571 (2022). [2] Kim, Dongjun, et al. "Maximum likelihood training of implicit nonlinear diffusion model." _Advances in neural information processing systems_ 35 (2022): 32270-32284.
Rating
6
Confidence
4
Soundness
3
Presentation
3
Contribution
2
Limitations
The limitations are addressed in the paper.
4. In conventional diffusion, where the forward process is parameterized as $F(\varepsilon, t, x) = \alpha_t x + \sigma_t \varepsilon$, there is a straightforward connection between the score predictor and the predictor of the data point: $s_\theta(z_t, t) = \frac{\alpha_t x_\theta(z_t, t) - z_t}{\sigma_t^2}$ (see [4]). Thus, during the training procedure, we can learn to predict either the data point predictor or the score function. Then, after training, any of these functions can be solved for from the other. From a theoretical standpoint, all these parameterizations are equivalent. When the forward process of NFDM is parameterized linearly as discussed above, NFDM exactly corresponds to conventional diffusion, and the function $x_\theta$ maintains the same relationship with the score function. However, in the general case, when $F_\phi$ is an arbitrary function, these connections between $s_\theta$ and $x_\theta$ do not exist. In our work, the parameterization of the reverse process through the function $x_\theta$ is a design choice. As discussed, $x_\theta$ retains relationships with the score function in the linear case. However, in the general case, we do not expect it to possess any special properties. We will include a discussion on this in the camera-ready, if accepted. 5. In the literature on diffusion models, the term "simulation-free" is most often taken to mean that during training the method does not require numerical simulations of the SDE or the integration of any dynamics at the training stage. For instance, approaches like [5], or others based on Schrödinger Bridge theory, require numerical simulation of an SDE that describes the forward process to sample $z_t \sim q(z_t|x)$. Such requirements make the training procedure significantly more expensive and are difficult to scale. In contrast, the NFDM framework allows for direct sampling of $z_t$ by inferring the function $F_\phi$, which includes no numerical simulations. We also do not need to simulate the generative dynamics during training, making the NFDM framework simulation-free. 6. As previously mentioned, the main contribution of our work is a comprehensive framework that supports a wide range of forward processes. The primary objective of the experiments with non-Gaussian parameterization of the forward process is not to surpass the Gaussian parameterization but to demonstrate the flexibility of the NFDM framework and its capability to handle even such complex parameterizations. We believe there are two main reasons why the non-Gaussian parameterization of the forward process does not show significant improvement compared to Gaussian parameterization. First, we have only evaluated our approach on image datasets, where Gaussian noise appears to naturally fit well. Secondly, we chose a relatively simple parameterization of the non-Gaussian forward process, which may be suboptimal. It is possible that in different domains or with a more powerful parameterization, non-Gaussian parameterization may prove more beneficial, but we leave this for future research. In this work, we aimed to demonstrate the functionality of the NFDM framework. Additionally, we would like to emphasize that even the Gaussian parameterization of NFDM with data-dependent variance is a generalization of existing diffusion models and leads to state-of-the-art log-likelihood estimation results. 7. Indeed, in Appendix B.3, we discuss that *a naive parameterization* of $F_\phi$ with arbitrary non-Gaussian functions is not scalable (density evaluation can be O(dim^3)). However, if we parameterize $F_\phi$ with functions that provide easy access to the log-determinant of the Jacobian of transformation, we can apply this parameterization even in the high-dimensional case. This is why, in our experiments on images with non-Gaussian parameterization of the forward process, we use a Glow neural network architecture. This architecture is scalable. Our experiments in Table 1 demonstrate that this non-Gaussian parameterization can lead to better log-likelihood estimation results, achieving state-of-the-art *on both* CIFAR-10 and ImageNet. [4] Kingma, Diederik, and Ruiqi Gao. "Understanding diffusion objectives as the elbo with simple data augmentation.” [5] Zhang, Qinsheng, and Yongxin Chen. "Diffusion normalizing flow.”
8. In the paper, we consider two related works that also propose approaches to learn straighter generative trajectories: [6] and [7]. The approach in [6] may be seen as a special case of NFDM (see the last paragraph in Appendix E.1). As soon as the authors in [6] propose to learn an additional function to parameterize the forward process, the training complexity of their method is comparable with NFDM. However, in NFDM-OT, we also need to calculate the curvature penalty (eq. 16), which makes the training more expensive compared to [6]. In [7], the authors proposed constructing the forward process with optimal transport data-noise couplings. The computational disadvantage of this approach is the necessity to find an optimal coupling for each mini-batch, which may be costly for large batch sizes and some complicated metrics (e.g., distance between 3D structures). In both [6] and [7], the inference works like in conventional diffusion models. However, due to the chosen parameterization, the inference of NFDM is more computationally expensive. In our experiments, instead of predicting the velocity field directly, we predict the data point $x$ and then substitute this prediction into the forward process. Therefore, we need to infer two functions: $x_\theta$ and $F_\phi$. The primary goal of NFDM-OT is to demonstrate the NFDM framework flexibility and its capability to learn generative dynamics with desired properties. In contrast to [6, 7], which only allows minimisation of trajectories curvature, NFDM allows learning generative dynamics with a variety of different properties like we demonstrate in Figure 2. Nevertheless, NFDM-OT is able to efficiently learn straight-line generative trajectories and demonstrates better performance compared to prior works as illustrated in Table 2. 9. In our experiments, we observe a similar slowdown for both Gaussian and non-Gaussian parameterizations of NFDM. The training iteration of NFDM takes approximately 2 times longer compared to conventional diffusion models that don’t parameterize the forward process. This happens because we use neural networks of similar sizes to parameterize both the forward and the reverse processes. Similarly, due to the parameterization of the generative process which uses the forward process, one inference iteration also takes approximately 2 times longer compared to the inference of conventional diffusion. We would like to emphasize that such properties of the models are consequences of these design choices. For simplicity of the experimental setup, we use neural networks of the same size to parameterize the forward process, which might be redundant. Additionally, it’s possible to choose an alternative parameterization of the reverse process that wouldn’t rely on the forward process for inference. That’s why we think that the comparison under the same number of function evaluations is more important. Nevertheless, even with this current parameterization, NFDM-OT demonstrates better performance compared to prior work under the same computational budget (if we double the NFE in Table 2). We hope that the clarifications provided, highlighting the novelty and strong performance of our work, along with further explanations of technical details will enhance your support for our paper's acceptance! [6] Lee, Sangyun, Beomsu Kim, and Jong Chul Ye. "Minimizing trajectory curvature of ode-based generative models.” [7] Pooladian, Aram-Alexandre, et al. "Multisample flow matching: Straightening flows with minibatch couplings.”
I appreciate the time and effort the authors invested in addressing the concerns I raised. Together with your clarifications, I consider your work to be a relevant contribution to the field, and therefore, I have decided to raise my score.
I acknowledge having read the authors' rebuttal. My overall assessment of the paper remains unchanged, and I continue to support my current rating.
Discussion
Please participate to the discussion.
Thank you for the response. It addressed my concerns and I've raised my score to a 7.
We would like to thank reviewers for their time and valuable contributions and discussion, which enhance the quality of our work. We are grateful for the reviewers' recognition of the NFDM framework as a “significant contribution”, “elegant”, “neat”, “novel” and “clearly explained”. We also appreciate the acknowledgment of our experiments as "robust and extensive," and the recognition of the potential of this research direction to be "impactful" and "influence future research." We are especially thankful for the reviewers’ unanimous recommendations. Here we summarise the main contributions of NFDM and the discussion: - We present NFDM, a novel general framework for diffusion models, which allows non-linear, time-dependent and learnable forward processes. - We present novel theoretical results, where we derive a simulation-free objective for NFDM, which is crucial for efficient training. We also derive ordinary and stochastic differential equations for sampling. - NFDM subsumes a large portion of existing approaches as special cases, enabling the free specification of the forward process (Section 7 and Appendix E). Moreover, it extends these much further, offering additional flexibility by introducing learnable parameters. - NFDM, with a simple parameterization of learnable transformations and without hyperparameter tuning, consistently and significantly outperforms baselines in density estimation across a range experiments (Table 1). - With a simple modification, our framework can also facilitate building bridges between two different distributions, NFBM. - NFDM opens a prospect for building new generative dynamics with specific properties, which conventional diffusion models like DDPM and DDIM are incapable of. - NFDM with additional penalty on curvature of deterministic generative trajectories (NFDM-OT), enables the learning of straight-line generative dynamics. This advancement permits the generation of samples in significantly fewer steps, outperforming baselines in few-step generation (Table 2). - NFDMs does not require any additional techniques to ensure stable training. We trust that the insights gained from our discussion will bolster your support for acceptance!
Decision
Accept (poster)