Score-based free-form architectures for high-dimensional Fokker-Planck equations

Authors: Feng Liu, Faguo Wu, Xiao Zhang
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: Fokker-Planck Equations, Normalization Condition, Score Model, Physical Constraints.

Abstract:
Deep learning methods incorporate PDE residuals as the loss function for solving Fokker-Planck equations, and usually impose the proper normalization condition to avoid a trivial solution. However, soft constraints require careful balancing of multi-objective loss functions, and specific network architectures may limit representation capacity under hard constraints. In this paper, we propose a novel framework: Fokker-Planck neural network (FPNN) that adopts a score PDE loss to decouple the score learning and the density normalization into two stages. Our method allows free-form network architectures to model the unnormalized density and strictly satisfy normalization constraints by post-processing. We demonstrate the effectiveness on various high-dimensional steady-state Fokker-Planck (SFP) equations, achieving superior accuracy and over a 20$\times$ speedup compared to state-of-the-art methods. Without any labeled data, FPNNs achieve the mean absolute percentage error (MAPE) of 11.36%, 13.87% and 12.72% for 4D Ring, 6D Unimodal and 6D Multi-modal problems respectively, requiring only 256, 980, and 980 parameters. Experimental results highlights the potential as a universal fast solver for handling more than 20-dimensional SFP equations, with great gains in efficiency, accuracy, memory and computational resource usage.

Decision: accept (poster)
Reviewer scores: 5, 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1729764714819

thread_type:
official_review

<thread_object>
creator: reviewer_eksf

<content>
summary:
This paper intends to solve high-dimensional Fokker-Planck equations which faces several challenges including curse of dimension, normalization constraint, etc. The solution proposed is to train neural network with a score-matching loss which bypasses normalization constraint by computing normalization constant as post-process. The method belongs to supervised learning, where training data is generated by stochastic Runge-Kutta method. The result seems effective and surpasses baseline model in accuracy.

strengths:
**Originality** The score-based loss is novel and seemly interesting. Given the close connection between Fokker-Planck (FP) equations and diffusion process and the noticeable succuss of diffusion model with score-matching loss, it is worthy trying to solve FP with score-based loss.

**Clarity** I find the paper very clear to read and well organized.

weaknesses:
At the first glance, it is a seemly natural and attractive idea to solve Fokker-Planck (FP) equations with the proposed score-matching loss, especially considering the success in training diffusion models and the close connection between stochastic process and FP equations. However, after more careful thoughts I find it hard to reason through the following questions:

1. If the proposed method needs a postprocess of calculating normalizing constant, why not treating PINN with the same postprocess? One of the major motivations of the paper is to deal with normalization condition (NC), which the authors criticized PINN being hard to satisfy with soft constraints. However, if PINN is obtained without NC and is normalized with the same quadrature technique afterwards, this motivation is weakened.

2. The score-based FP loss (equation (6)) is derived for static FP equations. What is the difficulty with non-stationary FP? It seems to me the residual loss can be transformed similarly and just one more term is needed in equation (15), which is $\partial_t \log p_{\theta}(x)$. Even if we only consider SFP, it is clear now that score-based FP loss is essentially equivalent to residual loss of PINN. Therefore, I wonder what benefit score-based FP loss can introduce? 

Based on these questions, I suggest the authors do an ablation study of replacing score-based FP loss with PINN loss (residual loss) without normalization constraint. Otherwise, the improvement of FPNN over TFFN may be purely due to removing normalization constraint from loss function.

questions:
See weaknesses above. Also, when using MAPE for metric, how do you calculate $p(x)$ for test dataset?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730084922814

thread_type:
official_review

<thread_object>
creator: reviewer_hrfk

<content>
summary:
The paper introduces the Fokker-Planck Neural Network (FPNN) framework, a novel approach to solving high-dimensional steady-state Fokker-Planck equations by leveraging a score-based PDE loss that decouples density normalization from score learning. This decoupling allows FPNN to avoid continuous computation of the partition function, enhancing computational efficiency and stability, particularly for complex high-dimensional problems. Experimental results demonstrate FPNN’s superior performance over existing methods, achieving high accuracy and significant speedups, making it a promising candidate for scalable and efficient solutions to Fokker-Planck equations.

strengths:
(1) This paper proposes a novel network for solving high-dimensional steady-state Fokker-Planck equations. By utilizing a score-based PDE loss that decouples score learning from density normalization, the network achieves an effective balance between representational capacity and the constraints required for accurate solutions.

(2) The performance results are promising, though further validation is needed to strengthen the findings.

(3) The presentation of this paper is good.

weaknesses:
(1) In this paper, using SRK as part of the data generation process for steady-state Fokker-Planck equations is effective, but it might lead to an "unfair advantage" in comparisons if other baseline methods do not leverage a similar approach for handling randomness or approximating steady states. Thus, it will be meaningful to test the performance of the proposed method with the same training data as proposed in TFFN paper.

(2) Score-based methods have shown strong performance for in-distribution problems, but they often suffer from significantly reduced effectiveness on out-of-distribution (OOD) tasks. I am curious whether the authors evaluated the OOD performance for this model. Additionally, as mentioned previously, I wonder if the data generation approach used in this study simplified the distribution, making it easier for the network to learn. If that is the case, the improvements might be attributed more to the engineering aspects of data preparation rather than advancements in the network architecture itself.

(3) Additional comparisons would strengthen this paper. Although FPNN outperforms TFFN in the results presented, I could not find any published reference for TFFN, suggesting it may only be available on arXiv and has not yet undergone peer review. To more convincingly demonstrate the effectiveness of the proposed method, I recommend that the authors include additional, widely recognized baseline methods. This would provide a more comprehensive evaluation of FPNN’s efficiency and robustness.

questions:
The questions are addressed in the weaknesses section. Although this research is not fully aligned with my area of expertise, I will follow the rebuttal process and hope that my suggestions help improve the clarity and presentation of the paper.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730642601734

thread_type:
official_review

<thread_object>
creator: reviewer_aogx

<content>
summary:
This paper introduces the Fokker-Planck Neural Network (FPNN), a novel framework for solving high-dimensional steady-state Fokker-Planck (SFP) equations. Traditional deep learning approaches to these equations face challenges with representation capacity, loss function balancing, and maintaining normalization constraints. The proposed FPNN addresses these issues by decoupling score learning from density normalization, using a score PDE loss that enables strict adherence to normalization constraints while allowing flexible, mesh-free architectures. FPNN achieves significant computational efficiency, with over a 20x speedup compared to state-of-the-art methods, and requires only minimal parameters for high accuracy. The authors demonstrate its effectiveness on 4D, 6D, and even 20D problems, achieving low relative errors without labeled data. The FPNN framework contributes a fast, efficient, and accurate solution method for high-dimensional Fokker-Planck equations, with applications in computational physics and related fields.

strengths:
The paper presents a novel Fokker-Planck Neural Network (FPNN) that innovatively decouples score learning and density normalization for high-dimensional steady-state Fokker-Planck (SFP) equations. Traditional deep learning methods for Fokker-Planck equations generally incorporate the PDE residual as part of the loss function but often encounter issues with representation capacity, balancing multi-objective loss functions, and satisfying normalization constraints. FPNN addresses these challenges by using a score PDE loss, which separates score learning from normalization, allowing for a flexible, mesh-free network architecture. This approach is original as it removes the dependency on specific architectures, enables strict normalization through a single computation of the partition function, and offers substantial computational gains over existing methods. By rethinking how neural networks approach Fokker-Planck equations, the authors introduce a framework that improves upon limitations of previous methods, potentially broadening the scope of high-dimensional PDE applications in machine learning.

weaknesses:
1. Limited Discussion on Practical Constraints and Limitations
2. Lack of Scalability and Computational Complexity Analysis
3. Lack of Comparison with Other Recent Advances

questions:
1. Could be the proposed method widely used for other PDEs?
2. Please compare with recent developed methods.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730677093928

thread_type:
official_review

<thread_object>
creator: reviewer_1jzg

<content>
summary:
In this paper, the authors focus on steady state Fokker-Planck (SFP) equations and propose novel score-based PDE loss. The proposed loss does only depend on the score function $s_\theta$, avoiding the necessity to compute the normalization constant of the probability distribution. Furthermore, the authors propose to investigate the proposed loss on two types of architectures, namely tensor neural networks (TNNs) and MLPs. Experiments on several PDE examples are performed, showing the good performance of the proposed method.

strengths:
1. The paper is overall well written and easy to follow.
2. Experimental results are interesting and support the authors' claims.

weaknesses:
1. There is a potential clash with the notion of score in the ML/DL community, see questions.
2. It is difficult to evaluate the novelty of the work. For example, authors only compare the proposed approach to TFFN but no other baselines are provided. I would strongly suggest that the authors add more baselines to allow for an easier comparison with existing approaches, e.g. [1,3].
3. Similarly, it is not clear whether the set of chosen experiments are commonly used in the  PINN community. For example, is the dataset of the authors present in [2]?

(see refs in questions)

questions:
**Main comments**
1. I am afraid that the notion of "score" that is central to the paper is fairly different to the one usually referred to in deep learning, where the score is implicitly learned through denoising and the associated distribution is never explicitely computed. Here, the authors rather use a differentiable network and use this property within the loss. Can the authors comment on that?
2. Why is the computation of $Z_\theta$ important in the context of Fokker-Planck (FP)? If I understand well, the authors parametrize FP with a neural network, with a direct access to both $\nabla \log p$ and $p$. Is the summation approach of (8) and (10) tractable as dimension increase?
3. I struggle to understand the point of the authors lines 453-473. Firstly, the notation $|D_{norm}|$ is slightly confusing. Secondly, what is the source of randomness in $D_{norm}$  mentionned by the authors line 465 ? If the dataset for $\mathcal{D}_{norm}$ is simulated, could the authors generate more samples?
4. Fig. 9 is slightly unclear. The plots on top and bottom show two different things (top: MAP and MAPE, bottom: MAPE and Z) for two different architectures. In particular, why should Z and the MAPE be related? Could the authors comment on that?

**Minor comments**
1. While the paper is well written, some typos are remaining, e.g. "Score-based generate model" (line 129)
2. The color scheme from Fig. 4 (c) (bottom) is not clear, it seems most of the maps are identically 0. The authors may want to reduce the threshold.
3. The message would be more striking in Fig. 5 and 6 if experiment was run multiple time with different random seeds. This would allow the authors to provide smoother curves with mean and error bars. 


**References**

[1]
@inproceedings{zhai2022deep,
  title={A deep learning method for solving Fokker-Planck equations},
  author={Zhai, Jiayu and Dobson, Matthew and Li, Yao},
  booktitle={Mathematical and scientific machine learning},
  pages={568--597},
  year={2022},
  organization={PMLR}
}

[2]
@article{lu2021deepxde,
  author  = {Lu, Lu and Meng, Xuhui and Mao, Zhiping and Karniadakis, George Em},
  title   = {{DeepXDE}: A deep learning library for solving differential equations},
  journal = {SIAM Review},
  volume  = {63},
  number  = {1},
  pages   = {208-228},
  year    = {2021},
  doi     = {10.1137/19M1274067}
}


[3]
@article{cho2024separable,
  title={Separable physics-informed neural networks},
  author={Cho, Junwoo and Nam, Seungtae and Yang, Hyunmo and Yun, Seok-Bae and Hong, Youngjoon and Park, Eunbyung},
  journal={Advances in Neural Information Processing Systems},
  volume={36},
  year={2024}
}
</content>
</thread_object>

</thread>
<thread>
creation date:
1731831793937

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We greatly appreciate all the reviewers’ comments and constructive suggestions, which are extremely helpful in better organizing and presenting our work. Please allow us to first provide the background of our study, reiterate our motivation and improvements. Then we will respond to the concerns about comparisons and experiments with related work. This will help reviewers understand our algorithm and better evaluate our contributions.

$\textbf{Background}$

PINN is a general deep-learning framework for solving PDEs and has achieved significant success in various problems, such as the Navier-Stokes equations, Allen-Cahn equation, Schrödinger equation, etc. However, PINNs face challenges in Fokker-Planck equations, with the objective function: (plain PDE loss)

$\mathcal{J}\_{\text{plain}}(p_\theta)=\mathbb{E}\left[|\nabla\cdot(p_\theta\mu) - \nabla \cdot(\nabla\cdot(Dp_\theta))|\right]$

It is evident that the zero solution $p_\theta = 0$ also satisfies this equation. Therefore, directly minimizing the loss function above often leads to network collapse, where $p_\theta$ converges to a trivial solution. Existing strategies generally fall into two categories involving normalization constraints in soft or hard manner:  

1. Adding a penalty term $\mathcal{J}\_{\text{norm}}=\left(\frac{|\Omega|}{|\mathcal{D}|}\sum_{x\in\mathcal{D}} p_\theta(x) - 1\right)^2$, where the dataset $\mathcal{D}$ is sampled uniformly from $\Omega$ [4].  
2. Developing specialized structures to represent the density function [5-9].

However, the former method requires manual balancing of multi-objective losses, while the latter may limit the model's representation capacity (Figure 1).

Another earlier data-driven method guides the network to the desired solution by introducing reference solutions and the regression loss $\mathcal{J}\_{\text{label}}=\frac{1}{N^Y}\sum_{j=1}^{N^Y}(p_\theta(x_j)-p_j)^2$, i.e. **FP solver** [1] suggested by Reviewer 1JZg. Our 4D Ring problem is directly from this work.

For this case, FP solver utilizes a very large number of particles ($10^{10}$ sample points) for SDE simulations and constructs frequency histograms on a grid as reference solutions ($10^4$ reference points). This technique is computationally intensive and lacks scalability to higher dimensions.

As seen in Figure 6 of [1], despite using such a large amount of points, the $L_2$ error of FP solver is around $10^{-2}$. In contrast, our FPNN achieves a MAE of $5.56 \times 10^{-4}$ without any labeled data (resulting in an even lower $L_2$ error). Furthermore, FPNN reaches this superior performance with fewer parameters and only 2,000 unlabeled data points per iteration. FPNN requires fewer than 1,000 iterations to surpass this baseline performance (see Figure 4(a) in our work and Figure 6 in [1]).

$\textbf{Our work}$

Our method deviates from all existing approaches by changing the loss function directly. The improvements we have achieved can be attributed to a fundamental modification of the loss function, which decouples the normalization condition from the training process. Specifically, we replace the original loss term $\mathcal{J}_{\text{plain}}$ with the following:

$\mathcal{J}\_{\text{score}}(s_\theta)=\mathbb{E}\left[|\nabla\log \widetilde{p_\theta} \cdot \widetilde{\mu}(x) + \nabla\cdot\widetilde{\mu}(x)|\right], \quad \widetilde{\mu}(x) = \mu(x) - \nabla\cdot D(x) - D(x)\nabla\log \widetilde{p_\theta}$

In this form, the trivial solution $\widetilde{p_\theta}=0$ no longer satisfies the equation. The *score PDE loss* maintains equivalence with the *plain PDE loss* while avoiding a zero solution and eliminating the need for handling normalization conditions in training process.

To further illustrate, consider a special case where $D(x) = I_d$. If minimizing the score loss $\mathcal{J}_{\text{score}}$ drives $\widetilde{\mu}(x)$ to zero, the training process becomes equivalent to optimizing the objective function of *flow matching*:

$\mathbb{E}\left[\|\mu(x) - \nabla\log \widetilde{p_\theta}\|^2\right].
$

Thus, our method inherently integrates score matching while preserving the original Fokker-Planck equation and corresponding physical laws.

**Our main target** is to fully decouple the normalization condition during training, allowing the network to freely adjust the magnitude to learn the unnormalized density $\widetilde{p_\theta}$ (note that our network directly models $\widetilde{p_\theta}$ rather than $p_\theta(x)$ and appropriate scale is really necessary for learning high-dimensional densities). Subsequently, it is natural to perform a postprocess of calculating the partition function $Z_\theta$ under the score PDE loss. We derive the approximate solution as $p_\theta(x) = \widetilde{p_\theta}(x) / Z_\theta$. The **advantages** are also clear: it reduces computational costs, eliminates the interference caused by normalization constraints on optimization dynamics, and ensures a more efficient and coherent training process.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a "score Fokker-Planck neural net" to address the issue in trying to solve Fokker-Planck equations by PINNs. The idea is to first write the dynamics of the score function under Fokker-Planck and first solve this score Fokker-Planck equation. Next, normalization is computed numerically in an efficient way that derives from a tensorial parameterization of the neural net. I think this is a nice contribution to ML stochastic dynamical systems that warrants publication. There are some omissions in the bibliography (for example the score evolution from here: https://arxiv.org/abs/2210.04296), but nothing damning.