Normalizing flow neural networks by JKO scheme

Summary

The paper introduces a novel approach to train continuous normalizing flows/score-based diffusion models that exploits the Wasserstein gradient flow theory and the JKO iterative scheme. The main idea is to approximate the density evolution of the variance-preserving forward dynamics of a diffusion model model (i.e. an Ornstein -Uhlenbeck process) with a series of deterministic transport maps. Under some conditions, the resulting deterministic dynamics is invertible and can then be used to map the stationary distribution back to the (implicitly learned) data distribution. In practice, this is very similar to what is done in ordinary score-based diffusion models, but the networks are trained in a radically different way since with the JKO approach the score is trained indirectly through a deterministic algorithm.

Strengths

I find this work exciting. The approach is very original as introduces a brand new way to train score-based models, which is backed by serious theoretical arguments. This both increases my confidence in the algorithm and opens the door for possibly very fruitful connection between several important research areas. The paper is well-explained, although it relies on several complex concepts that could alienate some readers. However, I think that in this case the mathematical complexity is unavoidable and gives added value. I appreciate that the authors not not solely focus on the mathematical framework but also discuss several important issues in the algorithmic implementation. The experimental section is convincing and rather comprehensive. The results on the low-dimensional datasets are good when compared with other state-of-the-art methods. The results on images are less impressive but overall convincing. While, in general, the application to generative computer vision needs to be fully flashed out, I agree with the authors that this goes beyond the scope of this work.

Weaknesses

I think that, generally speaking, the paper does not have major weaknesses. However, the part of the experiment section dedicated to image generation would have benefited from having a larger scope, for example by including experiments on CelebA HD and ImageNet. The selection of evaluation metrics should be complemented with FID scores, since this is the standard metric in the image generation literature and it is generally more reliable than the BPD score. Note that I do not think that having a lower FID score than the baseline is an argument for rejection, given the theoretical and algorithmic novelty of the work. Ther exposition would benefit from a more extended section connecting the JKO approach to the standard score-matching diffusion theory. This would greatly help readers coming from the diffusion literature, who in my opinion are one of the most important audiences for this work.

Questions

- Can you elaborate on the connections between the JKI iterative training and the score-matching training of diffusion models? It would be informative to discuss differences and similarities and give an overview of the strengths and weaknesses of each approach. - I would like to see the FID scores on MNIST and CIFAR10.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

Nothing to discuss.

Summary

This paper presents an innovative method for enhancing the stability of trajectories in CNF (Continuous Normalizing Flow) models. The authors propose incorporating the JKO scheme, which introduces regularization between the current density and the base distribution. The core concept revolves around leveraging the Wasserstein distance to learn an optimal transport map that guides the data distribution towards a normal equilibrium. To assess the effectiveness of the model, experiments are conducted on various scenarios, including toy examples, tabular data, and datasets characterized by higher dimensions like CIFAR or MNIST.

Strengths

- The idea behind the paper is novel and non-trivial. - Optimizing the velocity field instead of the transport map using JKO-Scheme is intriguing. - The procedure of training the model given the Algorithm is very efficient compared to reference baselines. - The method that can be seen as a particular case of regularisation of CNF models achieves outstanding results on tabular benchmark datasets. - The authors consider higher-dimension cases by conducting experiments on CIFAR. It is not common for such papers.

Weaknesses

- The proposed method still does not provide good-quality image samples. This group of models is rather dedicated to obtaining high values of NLL and modeling the distributions of low-dimensional data. - The variety of experiments conducted in this work may be richer. In particular, it would be interesting to see how the proposed approach works in some scenarios where CNFs perform well, like modeling multidimensional probabilistic regression, using the model in latent space of VAEs, or as a plugin model to large models (like StyleFlow).

Questions

I do not have any particular questions for the authors.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The main limitation of the proposed approach is the lack of capabilities to generate good-quality of image data, but this group of methods is rather not designed for such problems. Moreover, the evaluation may be a bit broader and include some other applications mentioned in the "weaknesses" section. Besides that, I admire the contribution of the work. and I think it should be accepted to the conference.

Summary

This paper proposes learning an invertible normalising flow using a neural ODE as a unique solution to the Fokker-Planck Equation (FPE) of the transport problem from the data distribution to the equilibrium solution. The solution of the FPE is obtained by using the JKO scheme by formulating it as a variational problem. The critical contribution of the paper is that each residual block of invertible neural ODE corresponds to a JKO step and the training objective can be computed from pushed data samples through the previous blocks. This leads to a block-wise procedure to train the JKO-iFlow model which is computationally more efficient than an end-to-end approach in the literature. The paper also proposes a way to determine the number of blocks adaptively by adaptively reparameterize the computed trajectory in the probability space with refinement to improve the model accuracy and the overall computational efficiency. Lastly, they empirical show that their method is competitive or better generative performance compared to existing flow and diffusion models on synthetic and real data,

Strengths

This paper is novel, and the idea of using iterative block-wise training is a good one. This strategy can be applied to other works as well and I think is a good idea. The formulation of the loss function from the JKO-step is also very well-written and clear. One strength of the paper is that their proposed method is very flexible as they used neural ODE with sufficiently small step sizes in the residual blocks to invertibility making their method very general and doesn't have any restrictive architecture like is typical of neural of other invertible flow models.

Weaknesses

The experiments are on relatively small datasets. It would be instructive to test how the model performs on larger datasets like miniImagNet or even CelebA dataset. The MMD loss could try to use different kernels or even mixed kernels to evaluate the quality better of the generative models. When comparing JKO-iflow against DDPM, it might be valuable to discuss what noise scheduling for DDPM is being used for each dataset, as DDPM is sensitive to the noise schedule used. Thus, it might lead to an unfair comparison if DDPM are not properly tuned. The authors can also include inverse problems in their experiments as many one of the advantages of DDPM is the fact they can easily be adapted to tackle inverse problems as well.

Questions

Could there be probability trajectories that are particularly "stiff" much like how there are stiff ODE that require very small step sizes for the numerical ODE method to be stable and have the resulting solution be close to the actual trajectory? The FPE does not only appear in image generation, in mean-field game theory the fokker-planck equation describes the dynamics of the aggregate distribution of agents can JKO-iflow work well on these other problems ?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors adequately address the limitations of their work, and there are no obvious potential negative social impacts of their work.

Summary

The authors introduce a novel normalizing flow training algorithm that integrates continuous normalizing flows and the JKO scheme. The objective of the training algorithm is to minimize KL divergence between the current density and the equilibrium density. The proposed method is inspired by the JKO scheme to enable adaptive blockwise training of residual networks. Furthermore, the authors demonstrate their method on synthetic and open-source datasets.

Strengths

1). The proposed method is computationally efficient compared to other normalizing flow methods or diffusion. 2). Novel idea of optimizing the velocity field instead of the transport map motivated JKO-Scheme. 3). The paper is well-motivated and easy to understand.

Weaknesses

1). The paper lacks metrics/empirical validation on larger images. FID metrics for Figure 4 should be necessary considering the paper is about generative modeling and its challenging to tell how adequate the CIFAR-10 samples are because they have such low resolution via eye test. 2). A significant highlight is that block-wise training leads to a large reduction in computation cost compared to other popular generative models (Figure 4), but the authors do not highlight that the proposed scheme uses a pre-trained auto-encoder, where some of the other methods do not.

Questions

1). Is it possible for the authors to provide FID scores for both MNIST and CIFAR-10? 2). How good is the proposed method without using a pre-trained auto-encoder?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The limitations are clearly discussed.

Thank you for the response

I want to thank the authors for their detailed responses. I am happy to keep my score the same and recommend this paper for acceptance.

Thank you

Dear authors, thank you for the rebuttal. I maintain my decision, and I believe the paper should be accepted.