We extend the concept of self-consistency for the Fokker-Planck equation (FPE) to the more general McKean-Vlasov equation (MVE). While FPE describes the macroscopic behavior of particles under drift and diffusion, MVE accounts for the additional inter-particle interactions, which are often highly singular in physical systems. Two important examples considered in this paper are the MVE with Coulomb interactions and the vorticity formulation of the 2D Navier-Stokes equation. We show that a generalized self-consistency potential controls the KL-divergence between a hypothesis solution to the ground truth, through entropy dissipation. Built on this result, we propose to solve the MVEs by minimizing this potential function, while utilizing the neural networks for function approximation. We validate the empirical performance of our approach by comparing with state-of-the-art NN-based PDE solvers on several example problems.
Paper
Similar papers
Peer review
Summary
This work focuses on deriving a uniform-in-time bound to the KL divergence between the NN-derived solution $\rho_t^f$ and the true solution $\bar{\rho}_t$, where $\rho_t^f$ is obtained via a push forward of samples of $\bar{\rho}_0$ with velocity parameterized by the neural network $f$. The neural network $f$ is trained to match the operator $\mathcal{A}[\rho]$ via the the Entropy-dissipation Informed Neural Network (EINN) loss given as an expectation over $\bar{\rho}_0$ of the integral from time $0$ to $T$ of the difference between $f$ and $\mathcal{A}[\rho^f]$. The method is then demonstrated to perform better than SOTA on McKean-Vlasov equations with singular interaction kernels: Coulomb interaction and Biot-Savart interaction.
Strengths
1. Uniform-in-time bound to the KL divergence for finite duration $T$. 2. Method was shown to work on singular interaction kernels: Coulomb interaction,\ and Biot-Savart interaction. 3. Singularity for Biot-Savart kernel was removed in gradient computation during training.
Weaknesses
1. Addressed in the last paragraph of section 4, the similarity in methodology to Shen et al. [2022]. Though the implementation is similar, the theoretical contribution is different. Perhaps include a comparison of their result in Figure 1 with and without generalization to MVE. And perhaps discuss the generalization to MVE. 2. Also addressed in the second last paragraph of section 4, the similar in loss function to Random Deep Vortex Network (RDVN) of Zhang et al. [2022]. Though the losses are similar, the authors mention a lack of discrepancy control in Zhang et al. [2022]. Perhaps provide more discussion.
Questions
1. Maybe I am misunderstanding, how is using a NN to parameterize $\mathcal{A}[\rho^f]$ helping in computation? $\rho_t$ is still computed via push forward of $\rho_0$. Would using the true $\mathcal{A}[\rho]$ be better? Or is the goal to first compute $f$ as a surrogate to $\mathcal{A}[\rho]$ then compute $\rho^f$ from the PDE with $f$ as a surrogate for $\mathcal{A}[\rho]$? 2. Why is $K_c$ set to $0$ instead of clipping some large value (line 173, equation (19))?
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
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
1. Addressed in Section 3, approximation error of neural network that if a neural network can represent the true $\mathcal{A}[\rho]$, EINN loss is zero. Perhaps reference some papers on neural networks designed to represent $f$ for MVE.
Summary
This paper proposes a new method for learning solutions to partial differential equations of the McKean-Vlasov type, and presents a theoretical analysis of the proposed method. In particular, for a certain class of equations, it is shown that the error bound does not increase exponentially in time. This is a significant result that takes advantage of the characteristics of the equations.
Strengths
The uniform bound presented by the authors is a very strong theoretical result. Although I do not understand the details of the proofs, the proofs seem to be reliable because several references are provided. In my opinion, this paper is an excellent paper as a paper in the research area of numerical analysis.
Weaknesses
The most important weakness is that it is not clear what practical applications the McKean-Vlasov type of partial differential equations have. Unlike numerical analysis, which is a research area in mathematics, there are many people who are interested in applications in the machine learning community. So these results may not be of much interest to them. Another significant weakness is that the adjoint method and the Monte Carlo method are required in the proposed approach, which are computationally expensive. Considering the computational cost, the proposed method is not considered superior to classical numerical methods, such as the finite difference method. In addition, the uniform bound seems to hold only for a class of equations, but this is not clearly written in Introduction.
Questions
(1) What are the practical applications of the McKean-Vlasov-type partial differential equations? (2) Because computationally expensive methods such as the adjoint method and the Monte Carlo method are required, the computational cost of the proposed method may be larger than the classical approaches such as the finite difference method. So, what are the advantages of the proposed neural-network solver compared to the classical numerical methods (e.g, the finite difference method)?
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
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
No potential negative societal impact is expected.
Summary
The authors propose a novel method for solving the McKean-Vlasov Equation (MVE), a Partial Differential Equation (PDE) that models the dynamics of stochastic particle systems with mean-field interactions. The new approach, named Entropy-dissipation Informed Neural Network (EINN), leverages neural networks to estimate the velocity field underlying the MVE and minimize a proposed potential function. This method is particularly designed to tackle singular interaction kernels, where traditional methods struggle. The authors provide empirical comparisons with state-of-the-art (SOTA) neural network-based MVE solvers and demonstrate the superior performance of the EINN framework. Despite the advancements, a noticeable weakness in this work is the omission of a discussion on the neural operator framework for solving PDEs, which represents a significant progression in the field. This lack of consideration limits the comprehensiveness and depth of the study.
Strengths
1. The paper provides a novel approach to solving MVEs. The idea of using entropy dissipation in the underlying system is inventive and beneficial to the community. 2. The research is solid and thorough. The authors have established theoretical guarantees for their EINN framework and provided clear derivations and proofs. 3. The empirical comparison of EINN with state-of-the-art methods and the positive results strengthen the claim of the proposed method's effectiveness. The authors ensured a fair comparison by using NNs with the same complexity for all methods tested.
Weaknesses
1. The lack of discussion on the neural operator framework for solving PDEs. This approach has been a considerable advancement in the field, and the authors' failure to address or refer to it is an oversight. Further exploration or comparison of the neural operator approach could enhance the comprehensiveness and relevance of the paper. 2. While the authors did a commendable job providing rigorous mathematical proofs, the paper could be difficult for those without a strong background in the relevant mathematical fields. Providing more intuitive explanations or graphical illustrations might have made the methodology more accessible. 3. The empirical testing seems to be limited in scope. More extensive experimentation, including larger-scale problems and more varied examples, would have strengthened the empirical evidence.
Questions
1. While the authors aim to utilize the Lyapunov function to study the deviation of the hypothesis solution from the true solution, they might have overlooked the possibility of overfitting. There seems to be no clear mention of measures taken to avoid overfitting in the model. How does the model ensure that it is generalizable and not just fitting the data it was trained on? 2. How can the EINN framework be adapted or expanded to deal with other types of PDEs? 3. How does the performance of EINN scale with the complexity of the problem, particularly in high dimensional cases? 4. Given the reliance on NNs, how sensitive is the performance of the EINN approach to the choice of NN architecture and parameters? 5. The authors should refer to and discuss the neural operator work for solving PDEs as related work. Here are some references: [Li, Zongyi, et al. "Fourier neural operator for parametric partial differential equations."]; [Kovachki, Nikola, et al. "Neural operator: Learning maps between function spaces."]; [Li, Zongyi, et al. "Neural operator: Graph kernel network for partial differential equations."]; [Gupta, Gaurav, Xiongye Xiao, and Paul Bogdan. "Multiwavelet-based operator learning for differential equations."]; [Xiao, Xiongye, et al. "Coupled Multiwavelet Neural Operator Learning for Coupled Partial Differential Equations."].
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
The authors have specifically designed this approach for solving MVEs, which might limit its general applicability for solving various types of PDEs.
Summary
The paper proposed an entropy-dissipation informed neural network for solving certain PDE systems. Using energy dissipation, the paper designed a special structure that propagates information from previous dynamics. They provided a theoretical guarantee for this framework, including the control of the KL divergence between the target and the estimator and a discussion of the approximation error of the NN. They also did some numerical experiments to verify their claims.
Strengths
1. The idea of using energy dissipation to design a framework is very interesting. EINN also takes the previous dynamics precisely and provides a uniform-in-time solution. 2. Mathematically speaking, the proofs are solid and well-written. 3. McKean-Vlasov type PDEs is an important class of PDEs for people from both math and physics communities. The presence of the singular kernel does bring challenges, which are handled by EINN.
Weaknesses
1. The authors specifically take advantage of the special structure of this type of PDEs. On one hand, it performs well for those PDEs; on the other hand, this limits the application of this method to other PDEs. 2. Numerical experiments are done with only comparison with PINN and DRVN.
Questions
1. In Figure 1, first row, at iteration 0, the value of the objective functions are not the same. 2. In Figure 1, it seems DRVN and EINN perform similarly. Please comment on this. 3. Please comment on how EINN compares with more recent neural-net-based PDE solvers, such as Fourier neural networks. 4. I wonder if more experiments can be done, other than 2D NSE and 3D MVE.
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
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
yes
I would like to thank the authors for carefully answering my many questions. I am positive towards this paper and am raising my score to 7. I have read all the other reviews and rebuttals.
Thank the authors for addressing all my questions. While some writing could be improved, I do think this work is very novel and interesting. I have raised the rating.
The authors addressed most of my concerns. I have raised my score to 6.
Thank you so much for answering my questions. Although some of my concerns have been addressed, it does not seem to me that the benefit explained by the authors is worth the computational cost. My rating is unchanged.
Response to Reviewer Jhg5 regarding the computational cost
We thank the reviewer for expressing again his concern on the computational cost. We agree that this is an important factor and would like to further clarify on this point. The following discussion will be included in our revision. Let's break down the computational costs of both EINN and the finite-difference (FD) method for comparison. For the simplicity of the analysis, we assume both EINN and FD are implemented **in a time-continuous fashion, i.e. we directly use their ODE formulation and solve the corresponding ODEs via numerical integrator**. In this way, given the same total evolving time, we can take the computational cost of a single numerical ODE step as the proxy of the total computational cost. Note that both approaches can be implemented in a discrete-time manner, but that would introduce an additional factor due to the time-discretization. We choose to discuss the time-continuous case for the sake of simplicity. Notations: Suppose that the ambient space of the PDE is $d$-dimensional. For FD, the computational domain of interest is a $d$-dimensional box, i.e. $[-L, L]^d$; the granularity of FD is $\eta$, i.e. in every dimension there are $L/\eta$ grid points and in total there are $(L/\eta)^d$ grid points. For EINN, let $p$ represent the number of parameters in the neural network, $b$ represent the SGD batch size, $m$ represent the number of particles, and $N$ represent the number of SGD steps. For the sake of concreteness, we will take the general MVE $$\frac{\partial }{\partial t}\bar \rho(t, x) + \mathrm{div}_x(\bar \rho(t, x) (-\nabla V(x) + K\ast \bar \rho(t, x)) = \Delta \bar \rho(t, x),$$ as example and analyze the overall computational cost of EINN and the finite difference method. - [Cost of FD] The solution to FD is represented using the density value over $(L/\epsilon)^d$ grid points. To evaluate the convolution, it takes $O((L/\epsilon)^{2d})$ operations, which dominates the cost of other computations. Consequently, in a single numerical ODE step, the cost of the FD method is $O((L/\epsilon)^{2d}). - [Cost of EINN] The cost of evolving the particles is $O(p(b+m))$ since all particles are driven by the hypothesis velocity field, represented using an NN with $p$ parameters and a single forward pass costs $O(p)$ for a single particle. The cost of evolving the score function (see Eq.18 in our paper) is $O(bdp)$ since the every function evaluation of the hypothesis velocity field costs $p$ and we need to perform this over every coordinate in the divergence and over every entry in the SGD mini-batch. The cost for evaluating the interaction is $O(d b m)$, where we recall that every point in the mini-batch of SGD interacts with $m$ particles (see section 2.1). Finally, the cost of computing the gradient using the adjoint method is $O(bdp)$ since the every function evaluation of the hypothesis velocity field costs $p$ and we need to perform this for every coordinate and every entry in the SGD mini-batch (recall that the cost of the back propagation as well as the Jacobian-vector-product of an NN is O(p)). In summary, in a single numerical ODE step in a single SGD step, the cost of EINN is $O(pb + pm + dbm + dbp)$. Additionally, since for EINN we need to run multiple SGD updates, the amortized cost of EINN overall is $O(Npb + Npm + Ndbm + Ndbp)$. Based on the above analysis of computational costs, we offer the following discussion: 1. The cost of finite-difference grows exponentially with the problem dimension $d$. Consequently, EINN has a clear advantage in this setting. 2. For the low-dimensional case, the comparison between the cost of EINN and finite difference heavily depends on the specific settings of the parameters and there is no unanimous decision on which method is more efficient. Let us just give one concrete example: In our 3D-MVE example, $p \simeq 3000$ (MLP with width 20 and depth 7), $m = 10000$, $b=100$, $N=10000$. So the overall cost of EINN is in the order of $10^{13}$. If we set for finite difference $L=0.1$ (this is a very small domain already), $\eta=1/1000$, we have the overall cost is of the order $10^{13}$. With this example, we do not mean to claim that EINN has the same or even strictly less cost than finite difference, but simply to emphasize finite-difference may not necessarily has a huge advantage over EINN in terms of computational cost even in this low-dimensional setting. Thank you again for your time and please let us know if this addresses your concerns.
Thank you very much again for the detailed response. Unfortunately, not all of my concerns have been addressed. 1. I agree with the authors that the proposed method is perhaps effective for high-dimensional cases; numerical experiments for such high-dimensional problems are preferably provided, though. 2. The discussion about the computational cost in low-dimensional cases seems to be misleading. I agree with the authors that "the comparison between the cost of EINN and finite difference heavily depends on the specific settings of the parameters and there is no unanimous decision on which method is more efficient." In addition, of course, the application of neural networks to physical simulations is a developing field, and it is not necessary to completely beat classical methods at the present time. However, the specific example provided for the 3D-MVE example seems to be misleading. - The authors first estimate the computational complexity for EINN, and then set up the finite difference method so that the computational complexity becomes comparable to that of EINN; however, this gives readers the impression that the finite difference method often requires a similar amount of computation as EINN. - The authors estimate the computational cost of the convolution as $O( (L/\varepsilon) ^{2d})$; however, if I understand correctly, this should be $O( d (L/\varepsilon) ^{d}(log (L/\varepsilon) ) )$ because the convolution can be computed by using FFT. So, if $\eta=1/1000$, the computational cost of the finite difference method could be around $10^{7}$ or $10^{8}$. - Perhaps, as the authors have noticed, the computational cost is closely related to **the tolerance of numerical errors.** In finite difference methods, at least the second-order method is typically used. If the mesh size is set to $\eta=/1000$, which means the number of nodes in each axis is 100, the error could be around $10^{-4}$, which seems to be more accurate than the results of the proposed method in the paper. In other words, the finite difference method seems to be able to find more accurate numerical solutions at a rate about, say, $10^5$ times faster. If the tolerance of numerical errors is the same as the numerical errors of the proposed method, the finite difference method is expected to be much faster, even if FFT is not employed for the computation of the convolution.
Response to Reviewer Jhg5 regarding the computational cost
We are glad that the reviewer agrees that in the high-dimensional case, EINN is potentially more efficient than the finite difference method. We also thank the reviewer for pointing out the strategies that to improve the performance of the finite-difference method. > numerical experiments for such high-dimensional problems are preferably provided, though. We will be happy to provide additional numerical results for high-dimensional tasks. However, at short notice, we could not finish the experiment before the author-reviewer discussion deadline. We will be working to provide this result in our revision. > The discussion about the computational cost in low-dimensional cases seems to be misleading. We are sorry that the reviewer find our discussion on the MVE case to be misleading and we acknowledge that there are many strategies to reduce the computational cost of finite difference method. Again, as stated in our previous response, we do not mean to claim that EINN has a computational advantage over the finite-difference methods in low dimensional cases. We simply want to point out that the computational complexity of both approaches are heavily influenced by the choice of parameters and there is no unanimous decision on which method is more efficient, which is acknowledge by the reviewer as well. Moreover, the discussion on whether NN based approach would be a better alternative to the classical method in low dimensional PDE is beyond the scope of our paper (this is basically one of the most important debates in the field). The goal of our paper is mainly theoretical oriented and is to provide rigorous error estimation for MVE with a singular interaction kernel. We thank the reviewer again for providing valuable comments to help us to improve the quality of the paper.
Decision
Accept (poster)