Summary
The paper introduces a method to solve Fokker-Planck equations using continuous normalizing flows, with known drift and diffusion functions. While one can resolve the dynamics using an ordinary differential equation (ODE), it necessitates knowing the score of the distribution at a specific time, a challenging task. To address this, the authors employ a neural network to learn $\log p(x, t)$. The network is trained through a regression between i) $\log p$ obtained by solving the ODE and ii) the network's output. Presumably, an exact solution to this regression would mean precisely learning $\log p$. Additionally, the paper utilizes discrete normalizing flows, specifically RealNVP, to model an initial distribution that eventually diffuses into the steady-state solution of a Fokker-Planck equation. The proposed algorithms, both time-dependent and steady-state, have been validated on simple problems with known Gaussian distribution solutions, and the results showed good agreement.
Overall, the learning method is novel and interesting, but the paper is lacking in clarity, rigor and convincing, non-trivial experiments. I do not feel the paper meets the standard for ICLR in its current form, but could meet the standard at ICLR or another venue after these points have been addressed.
Strengths
- Framing the problem as regression is clever – reminiscent of score matching, the technique behind diffusion models. Thus the central contribution of the paper is novel and interesting
- Most of the exposition of the background and theory is good, see weaknesses below for examples where this is not true
- The experiments shown seem to have worked very well. However, they are too easy (see below)
Weaknesses
- Quite hard to understand some key points on a first reading. I only understood what was going on after looking at the algorithms. For example, it would have been very helpful to mention explicitly that by $\log p_{net}$ you mean $\phi$ (or just write $\phi$ in the equation instead of $\log p_{net}$)
- Doesn’t mention that the divergence of a network (trace of gradient) is only cheap (constant wrt data dimension) if you are happy with an approximation via the Hutchinson trace estimator. Otherwise the cost of the trace is the same as calculating the full Jacobian and scales linearly with data dimension. Seems to be the case in this work, but scaling behavior is not mentioned
- No justification in the text of why the specific parameterization of $\phi$ was chosen. It would be nice to get a rough picture without having to read the references. In addition, it is claimed that the design imposes hard constraints, but those constraints are not specified, nor is it clear why they are desirable
- There should be a proof that solving the regression problem is equivalent to $\phi = \log p$. It is clear that if $\phi = \log p$ then the MSE loss is zero. Is the reverse also true? It is implicitly assumed, and an explicit demonstration would greatly improve the theoretical soundness of the paper
- There is no mention of diffusion models or score matching, even though this is quite similar to this work. In both cases, you want to learn a simple function of $p(x, t)$ (here the log, in score matching the gradient of the log) which allows you to solve the dynamics of the system. Diffusion models are based on Fokker-Planck equations (see [1] for example)
- The “method of characteristics” is not explained even though it is mentioned multiple times. At least a reference that explains it is a must
- The PINN model is not explained or referenced, even though it is mentioned as a competitor
- The experiments are disappointing: all problems are essentially toy problems with linear drift functions and known Gaussian solutions. Clearly $\log p$ is just quadratic, so the form you have chosen for $\phi$ easily fits this. In fact, the neural net part of $\phi$ should probably just be 0. It’s nice that the method works on these problems and keeps working for higher dimensions, but the problems need to be harder to be convincing. In the discussion you mention applications of Fokker-Planck equations such as in optimal control. Why don’t you try one of these applications? If that’s outside the scope of this work, you should have some problems where the drift and/or diffusion are nonlinear. You can solve the ground-truth dynamics with an SDE and approximate the agreement between the ground-truth and learned distribution via MMD or something similar
- There should be some idea of how the method compares against competitors, e.g. a table or figure showing the difference
[1] Song, Yang, et al. "Score-Based Generative Modeling through Stochastic Differential Equations." International Conference on Learning Representations. 2020.
Questions
- Why is $z$ transposed in eq (6)?
- Can you change the $x’$ and $t’$ notation to $x_1$ and $t_1$ or similar? This would be more in line with other works that solve ODEs or SDEs
- Can you avoid using $\mathcal{N}$ to denote a network? This is usually used for normal distributions
- What is MAPE?
- Are you missing a minus in the exp in eq (31)?
- Am I right that the toy example of 5.1 involves no learning? Maybe make this clear
- What are $\phi^f$ and $\phi^b$ in algorithm 3? Please define them
- In line 3 of algorithm 3, should it be $\phi^b$ instead of $\phi$?
- Why do you need the discrete normalizing flow for the steady-state problem? Can’t you just start with a standard normal for $x_0$ and evolve it to the steady state? Do you know anything about what $x_0$ is learned by the flow (in theory and in practice)? What does this distribution look like in your experiments?
- What do you mean by “In future work, we may explore the application of PINF with diffusion for density estimation tasks to enhance model generalization and data generation quality.” in the conclusion? Is your aim to enhance diffusion generative models? How would your method help?
Rating
3: reject, not good enough
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.