Score-based free-form architectures for high-dimensional Fokker-Planck equations
Authors: Feng Liu
Faguo Wu
Xiao Zhang
Venue: ICLR 2025 Poster (2025)
Keywords: Fokker-Planck Equations
Normalization Condition
Score Model
Physical Constraints.
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We propose a score PDE loss to decouple the normalization condition, allowing for free-form architectures in solving  high-dimensional Fokker-Planck equations.

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.

Published as a conference paper at ICLR 2025
SCORE-BASED FREE-FORM ARCHITECTURES FOR HIGH-DIMENSIONAL FOKKER-PLANCK EQUATIONS
Feng Liu 1,3,4,5 Faguo 44,66, Xiao Zhang 2,4,5,6,77+
School of Artificial Intelligence, Beihang University
2 School of Mathematical Sciences, Beihang University
3 NNational Superior College for Engineers, Beihang University
$ Keyy Laboratory of Mathematics, Informatics and Behavioral Semantics, MoE
"Zhongguancun Laboratory
6 "Bejjng Advanced Innovation Center for Future Blockchain and Privacy Computing
"Hangzhou International Innovation Institute, Beihang University
ABSTRACT
Deep learning methods incorporate PDE residuals 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 rep- resentation capacity under hard constraints. In this paper, we propose a novel framework: Fokker-Planck neural network (FPNN) that adopts 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 various high-dimensional steady-state Fokker-Planck (SFP) equations, achieving superior accuracy and over 20x speedup compared to state- of-the-art 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 Uni- modal 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 efficiency, accuracy, memory and computational resource usage.
INTRODUCTION
The Pokker-Planck (FP) equation governs the time-varying response probability density func- tion (PDF) of dynamical systems driven by stochastic processes (Risken & Caugheyz] 1991). It finds wide applications in statistical physics, chemistry, biology, mathematical finance, and structural dynamics (De Decker & Nicolis 2020
Jsoft
NN
strong
soft
(a) Real NVP
Representation
Solution
Constraints
Tu et al. 2020 Hu et al.) [002]] [Boghosian et al.] (2017). Solving FP equations presents three main challenges: high-dimensional vari- ables, unbounded spatial domains, and the nor- malization condition (NC). Grid-based methods face the curse of dimensionality, while path in- tegral methods and Monte Carlo simulations are limited by noise and computational complexity Naess & Moe 2000; Elman et al.) 2014 Chen
weak
hard
KRnet, LNF, GMM
Jhard
(b) GMM
Figure 1: Left: Strengths and weaknesses for isting models. Right: Real NVP and GMM with 16 components fail to learn the ring function. The former consistently leaves gap, while the latter exhibits numerous irregular protrusions.
Majda 2017 Natarajan al. 2021).
"Corresponding author.
Email:(LLu fang. buaa..ddu.cn

---

Published as a conference paper at ICLR 2025
Deep learning methods, with their grid-free. causality-free and feature-learingg capabilities, have demonstrated potential in addressing high- nigh-dimensional partial differential equations (PDEs) (Han et 2018 Yu et al., 2018 (Sirignano & Spiliopoulos 2018 Liu al.] (024). Among these promis- ing methods, physics-informed neural networks (PINNs) (Raissi et al., 2019) leverage automatic differentiation to enforce the physical constraints of underlying PDEs and achieve great success in various problems, such as the Navier-Stokes equation, Burgers' equation, Schrödinger equation, etc. However, for steady-state Fokker-Plaane (SFP) equations, PINN faces a new challenge: the zero so- lution also satisfies this equation and directly minimizing the plain PDE loss will quickly collapse the model a trivial solution.
To mitigate this issue, introducing labeled data or Kullback-Leibbe divergence terms (Zhai et al. 2022 Chen al. 2021) ). For 4D Ring problem, FP solver (Zhai et al.) (022) utilizes the direct Monte Carlo method with very large number of particles (10T0 sample points) to obtain 10 reference points. This method is computationally intensive and errors in the reference solutions may inversely limit the model's performance.
methods guide the neural network (NN) to the desired solution by
Table 1: Deep learning methods for SFP equations.
Arbitrary NN?
Whether strictly satisty NC?
Methods
Loss function
Model
Data-Driven
Jptsin + ATshed
pe(x) = ps)(2;
X
Jplin + ATToen
Po(2) Psx( 0)
X
Normalization
Condition
(al)
PXRaed (an: 0). PaMMx( (z;; 8), PINN(ar: 0)
/
Iasoee
Po (ar) = [an f),, Pa(x) Baluidyy
palm)
Another strategy that does not require labeled data is to impose normalization constraints in soft or hard manner. Soft constraints prevent zero solutions by adding normalization penalty term that enforces the density integral to unity (Alhussein et al.] 2023). While the normalization condition prevents the approximate solution from being zero across the entire domain, minimizing the plain PDE loss still has a tendency toward the trivial solution. These conflicting forces result in tortuous optimization dynamics (Wang et al. 20024] Al-Aradi et al. 2022), requiring delicate balancing of the multi-objective loss function. Specialized structures, such as normalizing flows (Dinh et al., 2016 Tang et al., 2022; Feng et al., 2022) and Gaussian mixture model (GMM) (Anderson & Farazmand, 2024), are developed to represent density functions and enforce hard constraints. These methods shrink the trial solution space and compromise the representation capability, leading to significant errors and computational burden when learning certain PDFs.
In summary, the former method requires manual balancing of multi-ohjective losses, while the latter may constrain the model's 's representation capacity. As shown in Figure the ideal method strictly satisfies normalization constraints without sacrificing representation capacity. To the best of our knowledge, there is no effective approach handle this conflict. In this paper, we propose novel framework: Fokker-Planck Neural Network (FPNN) to decouple the fitting and normalizing stages through score PDE loss. FPNN efficiently solves high-dimensional SFP equations and offers the following advantages:
Adaptive Domain for Complex Systems. High-dimensional state spaces often exhibit intricate dynamics and interaction patterns. To accurately identify critical features and high-prohd high-probability regions, we use the stochastic Runge- Kutta (SRK) method with strong order 1.5 to generate steady- state training data and adaptively determine the appropriate domain for SFP problems.
Flexible Scaling for Small Solutions. Since the probability density is non-negative and must integrate to unity over the spatial domain, the solution to high-dimensional problems is typically quite small and networks struggle to resolve such small yet meaningful scales directly. FPNN employs score PDE loss, allowing the network to freely choose appropriate scale and learn the unnormalized density. This effect of "data standardization' " is infeasible for existing methods without labeled data prior knowledge.
Efficient Enforcement of Normalization Constraints. The score loss allows free-form network architectures to strictly satisfy normalization constraints. FPNN avoids evaluating the density
2

---

Published as a conference paper at ICLR 2025
function integral during training process and instead calculates the normalizing constant by post- processing.
Scalability and Improvements in Optimization Dynamies. Without any labeled data or the enforcement of zero boundary conditions, we successfully solve various 4-20 dimensional SFP equations for complex physical systems. FPNN circumvents the normalization condition in loss function which causes difficulty in training, and substantially improves optimization dynamics. Therefore, a large learning rate (I 0.01) is allowed for fast descent and stable convergence, and FPNN achieves over 20x speedup compared to state-of-the- 11.36% 13.87% and 12.72% for 4D Ring, 6D Unimodal and 6D Multi-modal problems respec- tively, requiring only 256, 980, and 980 parameters. Our framework demonstrates substantial gains in efficiency, accuracy, and memory usage.
methods. We get the MAPE of
2 PRELIMINARIES
Fokker-Planck equation. Consider stochastic differential equations (SDEs) Oksendal 2013 of the form
dx = p (X)d oXX)dW
where the drift coefficient p Rd is vector field, o (XX E Rd is matrix-valuee function and W, is an M-dimensional standard Wiener process. The density function p(æ: of state variable X represents an invariant distribution, satisfying the steady-state Fokker-Planne (SFP) equation:
0p(x) Øt
A(p))
8 (Digp))) -V (ppL) V V. (Dp)), Un,Ox,
Lp Σ
Ox:
:=1==
(2)
f p(zr)da =1
i.e. Cp(x) O, p(x) 0
(||=|| o0):
where x Rd is the spatial variable, D(x) = (z) (x)² the diffusion matrix, L denotes the fokker-Plande differential operator and | indicates the &g norm of x.
Physics-informed neural networks. PINN is a deep neural network (x to approximate the solution p(x) with the universal approximation theorem (Hornik et al. 1989). The SFP equation in Eq. (2) can be expressed as,
Cp(x) 0 x E Ω
B,p(2,)) =0 0 x; E U2,
Ip
(3)
Here, B, are boundary conditions (BCs) for boundaries on, C Ω. and the integral operator I repre- sents the NC. These physical constraints are incorporated as regularization (typically mean-squared enror terms) in the loss function. To accelerate model training and reduce violations of physical laws, hard constraints can be used to omit part loss terms.
Score-based generative model. Score matching (Hyvärinen & Dayan 2005) is popular method
for learning unnormalized statistical models and score-based generative models have shown promis- ing performance in both sample quality and sample efficiency, which can be combined with numeri- cal SDE solvers or fast ODE solvers to generate samples (Song al. [202b)] Song Ermon 2020) Lu et al. 2022). Score loss usually optimized by minimizing the Fisher divergence between the gradients log-density functions (i.c., scores) and the ground truth scores from data, without handling the intractable partition functions (Song et al. 2020 primary inspiration for our score PDE loss.
Luo 2022
This idea serves as the
3 TOKKER-PLANCE NEURAL NETWORKS
Considering that optimization, data, and model form the cores deep learning approaches, we introduce the FPNN framework through these three components. First, define the score-based loss as the objective function to solve SFP equations. Then, steady-state data from
We mainly focus on the SFP equation, where the stochastic system reaches an invariant equilibrium dis- tribution and parameters μt, are independent of time f. Examples of PDEs, experimental settings, well as expressions for (x(z)) and (r(:) (z) can be found in Sec.

---

Published as a conference paper at ICLR 2025
Data
Model
Prediction
SDE = p(0)dt + oXDaW
MLP
TNN
Arbitrary NN
Madel
SRX
Estirate partiion
Zp
(x dx
Approddmat silitton
Perato: [-1,8,18]]
/2g
Score PDE Loss
Fl
FPAN
SFP
-P (pH) V-V (bD)
Minimize score lase: [1
Training process: (,.
Figure 2: FPNN framework for high high-dimensional SFP equations. We the SRK method to simu- late SDE for generating training dataset Drain and domain Q. Then, the data is fed into free-form architecture to produce the score. Next, train the model by minimizing the score-based FP loss. With the trained model Pa. we compute the partition function and perform normalization once to obtain the solution of SFP equation.
true distribution is required to evaluate the score loss, for which we simulate SDEs using the SRK method. Finally, we briefly review the tensor neural network (TNN) and the multi-layer perceptron (MLP), implementing appropriate modifications suit our task.
3.1 SCORE-BASED PDE LOSS
Score matching is originally designed for estimating unnormalized probability densities in machine learning, statisties, and signal processing. The score of density p(a:) is defined V 1(z). When the model learns a density (e()) with the partition function Ze. the approximate solution of SFP equation a normalized density denoted by.
(xx) Zo
po(#))
7
(o(x))d
(4)
Note that since log Po(x)) = loo (xx)) log Zo and (z = V log Po (zx), we immediately conclude that pe (x) and (x(x) share the same score, which does not depend on the intractable partition function Zn-
Inspired by this fact, we directly model the unnormalized density Po and the score function sp V log Po using neural network, and employ sy to approximate the true score. Once we reformulate the plain PDE loss of SFP equations into a score-based form, our model can learn the correct score function and focus the shape of PDF without handling NC. With the trained model po (xx), we compute the partition function only once by proc-ppog and obtain the approximate solution pe(x:) via Eq. (4)
Theorem 1 (Score-based FP loss). Assume the approximate solution (x( is differentiable and positive, satisfying regularity condition P for a fixed set of (p(xx) D())), [() finite for any O. Denote the model score function as sa = V log (ar) to approximate V log p(x:) and p(x) the true solution of SFP equation. The plain PDE loss can be expressed as the following up constant factor
Ep(x; [IF (x)]] = En (1me (xxx+ V [(x))]
(5)
where ji(x)) (((x) V (xx) - (sx)ss(a)) does not explicitly involve the density. F is named as the score-based FP operator, and the proposed PDE loss solely depends on the score 88 (x))).
Sinne solving FP equations with PINN is well well-posed problem, this result can be directly inferred from the loss function setup. Additionally, the loss functions E.[ [x)))]] and Ep(ae) FFso [=) share the same minimum value (zero) and uniqueness conditions for the solution under NC.
4

---

Published as a conference paper at ICLR 2025
Proof in Sec. A and further clarify the connection between the score PDE loss and score match- ing. By modeling the unnormalized density, we avoid the computational burden caused by frequently estimating Zp at each iteration, and also eliminate the interference of NC during training process. In this way, FPNN decouples the fitting and normalizing stages, ensuring continuous and efficient training process.
3.2 STEADY-STATE DATA GENERATION
To evaluate the score loss, we need training data from the true distribution p(x) of SEP equation, which can be sampled through the SDE simulation.
Dataset generation. Since the SDE is described by steady equation, the stochastic system al- ways reaches an invariant equilibrium distribution under the combined effects of drift and diffusion. Thus, freely choose the initial distribution po(#)) and sample an initial point set. These par- ticles then evolve according to the SDE dynamics over a long time 7 to produce a training dataset D train from the target distribution p pr-
fing N,
Unieoda: It..5
Algorithm SRK method for steady- state data generation
1: Input: terminal time T, number of time steps N, initial distribution Pu, drift term (zx) and diffusion term (zz)
2: // Dataset generation phase
3: Select initial points 20 Po
4: Compute Al = T/N
5: for n = 0, O......,N' - do
.N
Sample AW, (O, At)
7: Po = u({x), To = (Ix)
σ1 ( o.AW,, /2)
9: a2 = o + Mo(3
AW2)/44 0 AWA,, 22)
10:
T3 = o Ma(3A/ AW
12+
100 Hull mecet Do
Cacndia: D6 ml
11:
μ4 1( m3(3
AW 1/20,2W,
12:
F = (No 1111/2
13:
σ (oo 20 20 + /3)/6
14: In+1 = In + T + oAW,
15: end for
16: Generate training dataset Dinin = NN ~ PT
17: // Domain selection phase
18: Find the spatial bounds a; min (2x), by = max(:;)
d) for Durain
19: Determine the domain s X j=1 [a,, bi]
20: Output: D 52
Figure 3: Left: Cross- views of the training data distributions for SFP equations with com- plex patterns. Right: The SRK method consisting of two phases: dataset generation and domain selection.
Domain selection. Based on these samples, we determine an appropriate domain for SFP prob- lems, the regions not covered by dataset Draiin have very low near-zero probability density. This implies that the integral of the true density function outside 52 (i.c., 1- Jn pl x) (z)x)) is approximately zero. And this fact provides theoretical guarantee for calculating the partition function Zo over Ω. Additionally, training FPNN on Ω conserves computational resources and avoids the undefined value of log Pol (=(x)) at Pa (a O. There is no need add gated functions or enforce BCs, because the correct score model will ensure that pa is close to 0 at 0Ω.
We use SRK methods for the strong approximation of SDE, and generate trajectories and samples with the strong order 1.5 Newton scheme (Newton 1991). The training dataset Duain and problem domain Q are adaptively provided by Algorithm| Other numerical approximations for SDE simu-
Speciffically, S [d) needs to he sufficiently large to ensure that the particles thoroughly explore the spatial space and eventually reach steady state.
5

---

Published as a conference paper at ICLR 2025
lation alternatives, such the Suler-Maruyamai method and the Milstein method (Bayram al. 2018). For FPNN, we do not require the data points to exactly follow the true distribution pl ()), and our loss function offers degree of tolerance and flexibility. The SRK method and its generated data for different SFP equations are given in Figure 3
3.3 NETWORK ARCHITECTURE
FPNN separates the score model learning from NC, allowing for unrestricted network architectures. In our implementation, we focus on two general approximators: TNN and MLP to model the un- normalized density Pe Pe(x). The former is considered for its significant reduction in computational complexity when performing high-dimensional numerical integration for Ze. The latter is widely used deep learning tasks and the integral over domain 52 can be estimated by Monte Carlo sam- pling. Both networks are modified to better accommodate complex SFP equations.
3.3.1 TENSOR NEURAL NETWORKS
Inspired by SPINN (Cho et al. 2024) and TNN (Wang et al. 2022), we parametrize each spatial component using separated neural network R R"", which takes the coordinates of i-th axis as input and produces a r-dimensional feature representation. We propose the power embedding (PE) enhance the representation and better capture high-order drift dynamics. In each dimension, the input z; is transformed into [x;, "?..,.... ....
"",""" before being fed into the network Sa (m a man- ually specified typer-parameter This embedding enables the model to learn complex PDFs more efficiently and accurately (see Tables 4). The final prediction is obtained by taking an clement-wise
product of these r-dimensional features and summing the result:
d EII fa (4);;) ==1==1
po(
We employ the activation function tanh for all hidden layers and add an additional softplus activa- tion function before the output layer of the sub-networks f; to make density po (xx) strictly positive. A k-dimensional TNN can be built with & one-dimensional TNNs and approximate any k-dimensional vector functions u Rd R. For details on the approximation properties, refer Theorem 2.1 of TNN (Wang et al. 2022 and Appendix D.4 of SPINN (Cho et al. (024)..
Due to the low-rank structure of density representation in TNN, an efficient and accurate quadra- ture scheme can be designed for high-dimensional integration. Theorem P decomposes the high- dimensional integral of density Po (z: into a series of mieedimessional integrals, and we efficiently compute it using the piece-wise Gauss-Legende quadrature rule within 10 subintervals.
Theorem 2 (Partition function calculation). Given the density (x(x) parameterized by TNN, the partition function Zo is computed as,
Ze
Σ n1=1
wi
(ni) for (
.e.)).)))))))
where the nodes {u) (n.) dimension (i- ,,.
N: =n==1 and weights tw (uo(ni)) Jm-=1 are used for numerical integration in the i-th
3.3.2 MULTI-LAYER PERCEPTRON
To illustrate the applicability of our score PDE loss to different network architectures, we parame- terize the d-dimensional spatial variables using MLP fo RRd - Rr. The tanh function used in hidden layers and the softplus function is employed to ensure non-negativity in the output layer. The r-dimensional features are summed up to produce the final output:
Pa (x)
Σ (z
(8)
In our experiments, find that hidden neurons with small activation values overshadowed and inadequately trained in the original setup. And the modified MLP demonstrates improved learning performance.

---

Published as a conference paper at ICLR 2025
Efficient numerical integration in Eq. essentially forms high- dighessieeni grids using tensor prod- ucts, which imposes certain requirements on the network architecture. More generally, we provide the Monte Carlo estimation of partition function:
Po( az)
Zo po (2: By(w))
(((z)
For simplicity, we set ((x)) as an uniform distribution on Ω, i.e., ((x)) = x(5 for x E Q where denotes the Lebesgue measure. Since the optimization process does not involve normalization operations, we are free to choose the quadrature scheme to compute Zy by post-processing.
EXPERIMENT
Consider the challenges of solving high-dimensional SFP equations, we evaluate on several bench- mark problems with analytical solutions. The 4D Ring is tested in FP solver (Zhai al. while the 6D problems and 10D Multi-modal problem are adapted from TFFN (Wang et al. 2024 We further construct the 1OD Gaussian mixture distribution (Tang et al., 2022) and 20D Gaussian function to comprehensively test the applicability of our FPNN. Our PDE examples span various 4-20 dimensional steady-state solutions, including ring-shape density, arbitrary potential function, and Gaussian mixture distribution, with complicated interactions among spatial coordinates.
2022
Notably, few recent works can effectively solve such high-dimensional, challenging and different types of SFP equations, without the limitations in Table [ We utilize TENN as the baseline for comparison. The models are implemented in PyTorch framework and trained on NVIDIA Quadro RTX 8000 GPU with 48GB memory. Codes are provided in the supplementary material. We set consistent seed across all experiments to ensure the reproducibility our results.
6.0
shpe
(a) 4D Ring
(TANL
(b) 6D Unimodal
(c) 6D Multi-modal
Figure 4: Comparison of efficiency between TEEN and FPNN. The first column displays the true solution, while columns 2-5 and columns 6-9 show the predicted solutions (lop) and corresponding absolute errors (boltom) for TFFN and FPNN during training. Both models utilize the same error colorbar for fair comparison.
7

---

Published as a conference paper at ICLR 2025
4.1 EVALUATION METRICS
Due to the small solutions to high-dimensional problems, mean absolute error (MAE) is insufficient to accurately evaluate the performance models in different SFP problems. Even if the network learns zero solution, MAE is not noticeable and merely equal the average true density dataset, making difficult to identify this error. However, mean absolute percentage error (MAPE) results in 100%, giving great indicator of performance. We incorporate MAPE to measure the consistency between the predicted and true solutions.
the test
For high-dimensional problems, are limited to visualizing the results using selected slices. However, only testing errors on cross-sectional data is insufficient to measure the performance of high-dimensional solutions due their multi-modal complexity. Thus, we generate test dataset Dest = {p(x) > e, x E R") to globally evaluate error metrics. To avoid the ineffective MAPE, Dist is generated by the gradient ascent method on analytical solutions, with a threshold to reject extremely small probability densities. This approach is more efficient than the traditional method of randomly sampling spatial points and filtering out those with densities below The same test dataset and evaluation metrics are used for FPNN and TFFN, ensuring fair comparison.
Table 2: Experimental results of TFFN and FPNN on 4-6 dimensional SFP equations.
TFFN
FPNN (Ours)
SFP equations
Domain Ω
MAE
MАPE
MAE
MAPE
4D Ring
1.8, 1.8]*
3.61 x 10 3
49.25%
5.56 x 10
3.84%
6D Unimodal
1.22
4.00 x 10
293%
1.48 x 10 3
4.33%
6D Multi-
|
1.84 x 10-3
92.90%
1.98 x 10-1
12.18%
4.2 EXPERIMENTAL ANALYSIS
Score PDE loss. Without the interference from NC, score loss sig- nificantly improves the optimization dynamics, enabling FPNN quickly learn the correct scores and successfully capture the shape PDF. This crucial for the
of calculating normaliz- ing constant and is a key factor in the success of our framework. In we plot the training loss Jazone of FPNN for various SFP equations. Unlike the vanilla PINN, the magnitude of Jscore hardly change with increasing dimensionality or decreasing solutions. The score PDE loss consistently remains within stable range of 102 to 10 demonstrating numerical stability across dimensions. Empir- ical evidence shows that when Jscore falls below 1, FPNN generally learns the density function well. Thus, score loss also serves as an indicator of training progress and guides the network training re- gardless of dimensionality.
Storee
Figure 5
Figure 5: Plot of score PDE loss for 4-20 dimensional SFP equations during training.
Noting that both FPNN and TFFN strictly satisfy NC, we compare two PDE losses Jscoe and Jphain for 4D Ring problem in Figure 6 To balance the numerical scale of loss functions, we evaluate the plain PDE loss for FPNN and TFFN using 10k samples. These points are uniformly sampled from and not used for training both models. FPNN calculates Zo for normalization based on the current model each step. Under the same evaluation, we observe that Jplain of FPNN drops faster and lower, indicating that our model trains more efficiently and stably using the score PDE loss. Moreover, although Jplain of TFFN also decreases to small values, the shape of approximate solution does not match the true density in Figure 4a This deviation grows larger in higher-dimensional problems, leading to the failure or inapplicability of existing methods such as PINN and TFFN in high-dimensional SFP problems.
POS
Steps
Figure 6: Plot of score PDE loss and plain PDE loss for 4D Ring.
Network Architecture. Increasing network capacity and width enhances model performance, re- sulting in fewer training steps and more accurate solutions. In Table 4 and Table 5 PE shows great
8

---

Published as a conference paper at ICLR 2025
improvements in learning drift functions with high-order polynomial forms. Surprisingly, MLP- based FPNN utilizes 256 parameters to handle the complex 4D Ring problem and MAPE is reduced to 11.36% in only 4.8 minutes. In contrast, TFFN with 33,792 parameters is trained for 27.6 minutes and still fails to achieve comparable performance, as illustrated in Figure 4a Given the flexibility in model, ResNet (He et al. (0T6) is also alternative for deeper networks. In Figure [
we observe that both TNN-based and MLP-based FPNN correctly capture the two peaks along the 2g axis in 10D Multi-modal problem, requiring only 200 and 1k steps, respectively.
Table 3: Comparison of network structure for FPNN on high-dimensional SFP equations.
SFP equations
Domain ?
Model
Parameters
MAE
MAPE
TNN
84,480
1.26 x 10 4
32.00%
10D Multi-modal
[-1.2, 1.2
MLP
13,184
6.53 x 10-5
18.30%
TNN
126,080
1.84 x 10-?
24.50%
10D Gaussian mixture
[-5, 5,5
MLP
1,480
1.16 10
18.38%
20D Gaussian
2
TNN
252.160
1.22 10-8
12.79%
Partition Function. For MLP-based FPNN, we employ Monte Carlo sampling to draw a dataset Dz from the distribution g(ær) and estimate the partition function Ze in (9) The size of Dz. denoted by (Dzl. determines the accuracy of unbiased estimation of Zo. which improves as [Dx] increases. Figure 7 iallluustttees how the number of MC samples affects the estimation of Zo. and subse- quently, the prediction error of our solution.
MAKE
137%
Due to the randomness, errors fluctuates significantly with small dataset Dz, especially in the range [Dz] = 5k ~ 20k. As [Dz] increases, the prediction stabilizes. Higher-dimensional densities require more samples to accurately estimate the partition function. For all SFP problems, we set [Dzl 20k, which our experi- ments show is sufficient for the accuracy of FPNN. However, for 10D Gaussian mixture problem, the expanded domain resulted in a MAPE of 58.66% with this setting. After increasing [Dz] to 100k, we get better solution with a MAPE of 18.38%.
[0al
(a) MLP-based FPNN
For TNN-based FPNN, we also compare the Gauss-Legendes quadrature scheme in Eq.(7 with the MC sampling method in Eq. As shown in Figure 7b with more samples in Dz, the MC estimation ZMc gradually approaches the numerical integration ZGL (the rightmost column), and MAPE steadily decreases. More com- parisons and limitations are provided in Sec. B
[Dzl
(b) TNN- FPNN
Figure 7: 10D Gaussian mix- ture: MAE, MAPE and Zo at different [Dzl for FPNN.
(TAM)
Model
Layers
Parameters
MAE
MAPE
[vn. hidden layers. r]
[1. 64, 128] [3. 20, 20, 20, 20]
33.792
7.27 x 10-3
99.82%
TNN
5,360
4.25 x 10-3
65.41%
[3. 64, 64, 64)
34,304
6.91 x 10
7.58%
[S. 64. 128]
34,816
8.20 x 10
5.66%
64. 128]
35,584
x 10
3.84%
(d. hidden layers]
[4, 8, 8, 8]
256
1.26 x 10-3
11.36%
MLP
[4, 20, 20, 20, 20]
1,360
9.74 x 10
8.74%
[4, 64, 64, 64]
8,640
8.81 x 10
6.48%
[4, 128. 128)
17.152
6,21 x 10
5.39%
Table 4: 4D Ring: comparison of FPNN with different networks (10k steps).
Figure 8: 10D Multi-modal: predicted
solution of FPNN during training.

---

Published as a conference paper at ICLR 2025
Casc caltos
PPAN (TAN))
saps
raps
Model
Layens
Steps
[ms hädden layers, r]
20, 20]
7,800
102
10"22
30.60%
[1, 64, 64, 64]
50,688
28
1.33
x 10-3
31.84%
TNN
64, 64]
51,456
2k
2.94
10
9.49%
[5, 64, 64]
52,224
24
2:44
10"
8.00%
64, 128]
53,376
24
2.18
10
7.03%
Cac
PPNN
GR rg!
128)
102
[dl hiddin layers]
MLP
2010
44
10-2
41.64%
[6, 20, 20) 64, 64, 64]
980
6k
5.40 x 10-3
13.87%
8,768
104
2.94
10
8.61%
Table 5: 6D Unimodal: comparison of FPNN with different networks.
Figure 9: 10D Gaussian mixture: pre- dicted solution of FPNN during training.
Computational Efficiency. Our computational complexity is significantly reduced compared to existing methods, because the normalization condition does not need be explicitly considered during training. For PDE residuals, we only perform an additional logarithmic operation on the output layer, and the computational costs for Jplain- Iscore and their computational graphs remain roughly comparable. With the improved optimization dynamics and post-process of normalization, FPNN reduces both the total number of training epochs and the computational cost per iteration, achieving faster efficiency and enhanced performance.
Exact Souzton
hocts
uart
pm
1000
22.8
Figure 10: 20D Gaussian: predicted solution of FPNN. The first column shows the exact solution and the changes of MAPE during training. Due to symmetry, the cross sections of any two dimensions near the origin exhibit the same shape. The remaining parts display the predicted densities for pairs of adjacent dimensions, such (x1, 22,, 0, o,. 0), (0, 0, 23, v4, -., 0), and on.
5 LIMITATIONS AND FUTURE WORK
Despite comprehensive experimental results show the effectiveness of our FPNN framework on challenging high-dimension steady-state FP equations, there remain challenges solving time- dependent FP equations. Time-varying probability densities are often non-localized and it difficult to identify a suitable integration domain at any given time and compute the time-varying normalizing function. Our exploration merely promising start, and we look forward to its applications in broader fields, such as physics, finance, diffusion models, and mean-field games. In the future, aim to develop an efficient solver for general Fokker- Planckk equations.
6 CONCLUSION
We propose a novel framework, FPNN, that decouples the normalization condition through score PDE loss for solving high-dimensional SFP equations. FPNN eliminates the computation and influ- ence of the partition function for more efficient and stable training. In addition, arbitrary network architectures allowed to strictly enforce normalization constraints, without sacrificing represen- tation power. Our method outperforms existing methods in both training efficiency and prediction accuracy, demonstrating strong potential and general applicability.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by the National Science and Technology Major Project under Grant 2022ZD01 16401, the National Natural Science Foundation of China under Grant 62 62141605, and the Fundamental Research Funds for the Central Universities, China.
DATA AVAILABILITY
We provide the code for generating test datasets and fix the random seed for reproducibility, enabling future research to perform comparisons with FPNN under consistent evaluation metrics. Our code available at https//githuh..om//niffff//
REFERENCES
Ali Al-Aradi, Adolfo Correia, Gabriel Jardim, Danilo de Freitas Naiff, and Yuri Saporito. Extensions of the deep galerkin method. Applied Mathematics and Computation, 430:127287, 2022.
Hussam Alhussein, Mohammed Khasawneh, and Mohammed F Daqaq. Physics-informed solution of the stationary fokker-plank equation for class of nonlinear dynamical systems: evaluation study. arXiv preprint arXiv:2309. /6725, 2023.
Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Ap-
plications, 13(3):313-326, 1982.
William Anderson and Mohammad Farazmand. Fisher information and shape-morphing modes for solving the fokker-planck equation in higher dimensions. Applied Mathematics and Computation, 467:128489, 2024.
Mustafa Bayram, Tugcem Partal, and Gulsen Orucova Buyukoz. Numerical methods for simulation of stochastic differential equations. Advances Difference Equations, 2018:1-10, 2018.
Bruce M Boghosian, Adrian Devitt-Lee, Merek Johnson, Jie Li, Jeremy A Mareq, and Hongyan Wang. Oligarchy as phase transition: The effect of wealth-attained advantage in fokker- planck description asset exchange. Physica A: Statistical Mechanics and its Applications, 476: 15-37, 2017.
Nan Chen and Andrew J Majda. Beating the of dimension with accurate statistics for the fokker-planck equation in complex turbulent systems. Proceedings of the National Academy of Sciences, 1284499:1864422866 2017.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary
differential equations. Advances in neural information processing systems, 31, 2018.
Xiaoli Chen, Liu Yang, Jinqiao Duan, and George Em Kariaadkis. Solving inverse stochastic prob- lems from discrete particle observations using the fokker-planck equation and physics-informed neural networks. SIAM Journal on Scientific Computing, (3(3):B811-B88) 2021.
Junwoo Cho, Seungtae Nam, Hyunmo Yang, Seok-Bae Yun, Youngjoon Hong, and Eunbyung Park. Separable physics-infeeee informed neural networks. Advances in Neural Information Processing Sys- fems, 36, 2024.
Yannick De Decker and Grégoire Nicolis. On the fokker-planck approach to the stochastic thermo- dynamics of reactive systems. Physica A: Statistical Mechanics and its Applications, 553: 124669, 2020.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. arXiv preprint (605.06600008800 2016.
Howard Elman, David Silvester, and Andy Wathen. Finite elements and fast iterative solvers: with applications in incompressible fluid dynamics. OUP Oxford, 2014.
Xiaodong Feng. Li Zeng, and Tao Zhou. Solving time dependent fokker-planck equations via tem- poral normalizing flow. Communications in Computational Physics, 32(2):401-44) 2022.

---

Published as a conference paper at ICLR 2025
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Fait Sardee- Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014.
Jiequn Han, Arnulf Jentzen, and Weinan E. Solving high-dimensional partial differential equations using deep learning. Proceedings of the National Academy of Sciences, 115(34):8505-8510, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks uni- versal approximators. Neural networks, 2(5):359-366, 1989.
Jingwei Hu, JianGuo Liu, Yantong Xie, and Zhennan Zhou. A structure preserving numerical
scheme for fokker-planck equations neuron networks: Numerical analysis and exploration. Journal of Computational Physics, 433:1 2021.
Yiwen Huang, Aaron Gokaslan, Volodymyr Kuleshov, and James Tompkin. The gan is dead; long live the gan! modern gan baseline. arXiv preprini arXiv:2501.05441, 2025.
Aapo Hyvärinen and Peter Dayan, Estimation of non-normalized statistical models by score match- ing. Journal of Machine Learning Research, 6(4), 2005.
Tero Karras, Miik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in neural information processing systems, 55:265655-265777 2022.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hame Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. arXiv preprint arXiv:2209903000 2022.
Ziming Liu, Yixuan Wang, Sachin Vaidya, Fabian Ruchle, James Halverson, Marin Soljačić, Thomas Y Hou, and Max Tegmark. Kan: Kolmogorov-arnold networks. arXiv preprint arXiv:2404. 19756, 2024.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: fast ode solver for diffusson probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5787, 2022.
Calvin Luo. Understanding diffusion models: A unified perspective.
arXiv preprint
arXiv:2208.1 / 1970. 2022.
Naess and V Moe. Efficient path integration methods for nonlinear systems.. s Probabilistic engineering mechanics, 15(2):221-231, 2000.
Natarajan, PP Popov, and GB Jacobs. high-order semi-lagrangian method for the consistent monte-carlo solution of stochastic lagrangian drift-diffusse models coupled with eulerian dis- continuous spectral element method. Computer Methods in Applied Mechanics and Engineering, 384: 18001, 2021.
Nigel Newton. Asymptotically efficient runge-kutta methods for class of ito and stratonovich equations. SIAM Journal Applied Mathematics, 51(2):542-567, 1991.
Bernt Oksendal. Stochastic differential equations: an introduction with applications. Springer Science & Business Media, 2013.
Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686-707, 2019.
12

---

Published as a conference paper at ICLR 2025
H Risken and TK Caugheyz. The fokker-planck equation: Methods of solution and application. Journal of Applied Mechanics, 58(3):860, 1991.
Justin Sirignano and Konstantinos Spiliopoulos. Dgm: A deep learning algorithm for solving partial differential equations. Journal of computational physics, 375:1339-1364, 2018.
Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems,
2020.
Yang Song. Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. In Uncertainty in Artificial Intelligence, pp. 574-584. PMLR, 2020.
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in neural information processing systems, 34:1415- 1428, 2021a.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- tional Conference on Learning Representations, 2021b.
Kejun Tang, Xiaoliang Wan, and Qifeng Liao. Adaptive deep density approximation for fokker- planck equations. Journal Computational Physics, 457:111080 2022.
Tongbi Tu, Ali Ercan, and M Levent Kavvas. Probabilistic solution to two-dimensional stochas- tic solute transport model by the fokker-planck equation approach. Journal of Hydrology. 580: 124250, 2020.
Taorui Wang, Zheyuan Hu, Kenji Kawaguchi, Zhongqiang Zhang, and George Em Karni- adakis. Tensor neural networks for high-dimensional fokker-planck equations. arXiv preprint arXiv:2404.05665, 2024.
Yifan Wang, Pengzhan Jin, and Hehu Xie. Tensor neural network and its numerical integration. arXiv preprint arXiv:2207.027555 2022.
Bing Yu al. The deep ritz method: deep learning- based numerical algorithm for solving varia- tional problems. Communications in Mathematics and Statistics, 6(1):1-12, 2018.
Jiayu Zhai, Matthew Dobson, and Yao Li. A deep learning method for solving fokker-planck equa- tions. In Mathematical and scientific machine learning, pp. 568-597, PMLR, 2022.
13

---

Published as a conference paper at ICLR 2025
A LOSS FUNCTIONS FOR SFP EQUATIONS
We first introduce the loss functions and notations listed in Table [ to provide the background for SFP equations. Then, we present the derivation of the score-based FP loss, which constitutes the core innovation and contribution of our work. Furthermore, we explain the relationship between the score PDE loss of SFP equations and the score-matching loss of generative models.
A.1 RELATED WORKS
PINN is 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, PINN faces challenges in Fokker- -lankkk equations, with the plain PDE loss:
Jplain (pe) )+ = E [IV (paft) - V (V ((pp))
(10)
is evident that the zero solution py = 0 also satisfies this equation. Therefore, directly minimizing Jplain often leads to network collapse, where pe converges to trivial solution. Data-driven methods rely reference solutions ((zr..P))) and introduce the regression loss:
N
Tuabel
Σ (po(z,) - pp;}²
(11)
N
Without labeled data, existing strategies generally fall into two categories involving normalization constraints in soft or hard manner. Soft constraints add a normalization penalty term to enforce Zg = 1, thereby enabling the of (q. This involves sampling dataset Dnom from the distribution ((x)), with the number of samples given by (Droml our experiments, f2 0²=1'a (ai, 0,] is hypercube with the volume of v(22) II 1 |b, a.l. The normalization loss is defined as:
((2) DE
Jnorm
po(x) -
(12)
[D [Pnorm]
Hard constraints, on the other hand, develop specific structures to represent the density function, such as normalizing flows and their variants. The flow-based generative modeling is to seek an invertible mapping fx) and the PDF of follows the change of variables formula strictly adhere NC.
(x(x) pz( del [=)(x))
(13)
However, the former method requires manual balancing of multi-objective losses, while the latter may constrain the model's representation capacity. Furthermore, in terms of computational effi ciency, TFFN requires estimating Ze every iteration, "soft" PINN necessitates computing Joom at each step, and density estimation with normalizing flows involves simultancous tracking the changes of both and log-density.
FPNN deviates from all existing approaches by changing the loss function directly. Our score PDE loss allows model the unnormalized density Pe and decouples NC from the training process. Specifically, we replace Tplnin with the following score PDE loss:
Jscore (B) = E | log Po
x((x) + V ((x)]],
(14)
(xx)) = (ax))
D(x) Vog
In this form, the trivial solution pa no longer satisfies the equation. The score PDE loss maintains equivalence with the plain PDE loss while avoiding a zero solution and circumventing Tubel: Jnve: or the normalization condition loss function which causes difficulty in training.
A.2 DERIVATION OF SCORE PDE LOSS
Theorem 1 (Score-based FP loss). Assume the approximate solution pe po(x)) is differentiable and positive, satisfying regularity conditions: for a fixed set of (Hıx),, D(x)). ). Ei [[[[p( is finite for any O. Denote the model score function as so(x) V log Po( ) to approximate V p(x)) and
14

---

Published as a conference paper at ICLR 2025
p(x) is the true solution of SFP equation. The plain PDE loss be expressed as the following up to a constant factor
(mmp(a))(s(()) (e (((x) +
V-
(15)
where ju(x=) == p V (())
((x)) - ()())) does not explicitly involve the density. F is named as the score-bused FP operator, and the proposed PDE loss solely depends the score 80 s)(x)
Proof. First, for any positive probability density p(z), note that Lp(x) can be expanded to
Lp(x) = V (p(x))((=)) + V [V (D())p(z))
= V (p(x))(()) - V (D(z)p(=)))
V (p(x))(()) (V D(z))p(a) D(z))p(())
= V (p(x)(((()) V D(x) - (ogg)ooaa))
= (((x))((aa)
(16)
d
dz. [x())
where we define
xax)) == w(== V D(x) Vozg)oaaa
(17)
For an approximate solution pa(x). the plain FP loss measures the expectation of PDE residuals Jplain = gx)) [lCpo [x) where ((x²) can be selected as the uniform distribution on Q2 Based on Eq.(T6) we can rewrite this loss to obtain
Iplaia = En((r) (po(())) / v(2)) [pp (x) da
(18)
1
/ Vpo(z>) (7(x)) po (z) d(=
v(it) /, [o(o=)| VVogro( TT(x) (( dr
v(22) ( ||logoo(ar) x(x) x())
(it) E Ex(a) [[+))()]]]
v(2) Iscoro (80)
(19)
Our goal is to approximate the solution p(x) and score Vlog using po(x)) and sz Thus, we replace the expectation of pe with the ground truth p, allowing to train the score model better. For SFP equations, we easily obtain samples from the invariant distribution p(æ) through SDE simulation. When pa is close to p. we observe that, up constant factor of 1/v(Ω). the score-based FP loss is equivalent to the plain FP loss of PINN. But this form circumvents the need to handle the partition function and normalization condition.
If we sample training data over large spatial domain without prior knowledge, regions with ex- tremely low density may fall below machine precision, resulting in pe 0 and "Not Number" (NaN) errors for log po- But it has no impact on the prediction stage, as there is no logarithmic operation. By using the SRK method to generate high-probability training data of the steady-state distribution, we define the domain of SFP problems and avoid this issue.
To address numerical issues, another effective approach is to directly model the log-density log using neural network, which can be implemented with the score PDE loss. Assuming that the computer exactly represents positive values in the range [c, K]. this gives the maximum range of density the output. When we model log p, the range is K, K] and the corresponding density [e k e which is significantly larger than the previous range. This approach helps to mitigate the occurrence of NaN values during training.
15

---

Published as a conference paper at ICLR 2025
Sometimes, enhance training by adding weights w(x) according the probability density, and generalize the objective function as Pp(a) [2) (0 < a I 1). We set a = in our experiments, i.., norm for score loss function. Here is the following derivation:
E pp)))| [F)())))) F(a)))Fso(2;) da
v(2)) / v(12) p(x)) - [
y()) (((x) 1- [x) (po())))
= ( (p)
(20)
A.3 CONNECTION BETWEEN SCORE PDE LOSS AND SCORE MATCHING
To illustrate this connection, we consider a special case where D(x) = Ia. If minimizing the score loss Jacore drives (x) to zero, the training process becomes equivalent to optimizing the objective function:
E Eaa ((x) V log po [
(21)
Thus, FPNN inherently integrates score matching while preserving the original Fokker-Plaane equa- tion and corresponding physical laws.
(a) Score V log Paata (x)
Score matching
Langevin dynamics
Data samples (X Nx) (x)
Scores
New samples
*x()) Posa)
(b) Drift (xx)
Dx Ky.
De,, Ko. 00
6D Multi-modal
4D Ring
Figure 11: (a) the score in generative models and (b) the drift in SFP equations. Both represent the gradient information of target distributions and indicate regions of high probability.
Score matching. In Song's work (Song et al. 2021b), the Noise Conditional Score Network (NCSN) uses a weighted sum of denoising score matching objectives:
N
8" = arg mim
of Eunm
[m) s0 (z) V log ( [
(22)
and Denoising Diffusion Probabilistic Model (DDPM) leverages re-weighted variant of the evi- dence lower bound (ELBO);
0r arg min E 04)) Ep Ep, ( (l V P (T(ra))
(23)
It can be seen that equivalent score-matching loss:
8* arg min E E (p uxxx)
(24)
16

---

Published as a conference paper at ICLR 2025
shares a similar form with score-based generative models. These generative models gradually learn the score of the perturbed data distribution (i.e., conditional Gaussian distribution), while our FPNN directly learns the known drift uux)) for solving FP equations.
A.4 GENERATIVE MODELS AND FPNN
For score-based generative modeling, we consider three distinct dynamical equations:
Forward SDE.
da = fax, ()dl + g(]]d
(25)
The probability density function p(æx, t) satisfies the FP equation:
dp ôt - - (fp)+V.V.
(1 "996/
(26)
Probability flow ODE.
dæ = (
1 "0())0" ( logr) dt
(27)
Substituting the drift and zero-diffusion term yields the FP equation:
Op - [ (flostt) 1 (0(0)=((VVogp)
(ogg)))
(28)
Ot
-V- (090"D))
(29)
which is identical to 25. Here, g(t) is independent of x and it is important to empha- size that V Plogp must correspond the score of density p(æ, t) induced by the forward SDE to ensure consistency in simplifying (2) Since there is no diffusion term, the probability flow (PF) ODE reversible, allowing for efficient denoising process.
Reverse SDE.
dx (f(x,t) ((t)g" log dt gt)
(30)
Compared to the PF ODE, the reverse SDE introduces additional randomness to denoising processes, thereby increasing the diversity of samples. [Anderson (1982) provided the de- tailed derivation. Here, we give a less rigorous but intuitive explanation for clarity. Let the drift term of the PF ODE and the reverse SDE be fi t) and f2 t). We consider time reversion T- J and obtain: dæ FP equation is:
f2(2 T - r)dr g(T r)dw. The corresponding
0p v (-Far)
"99"'")
Ør
V (f gg" v Vlggg) +
(,900
1
Vp
v (T(
(31)
When expressed in terms of I, this becomes consistent with Eq.
dp -- (F((z,t)))
(32)
Øt
For generative models, both the drift term f and the diffusion term 9 are typically predefined. EDM (Karras et al. 2022) replaces the functions f(t) and g(t) with the geometrically meaningful scale schedule s(t) and noise schedule att). In summary, the forward SDE is only applicable for progressively adding noise to a dataset and cannot be used for sampling images. To generate high- quality samples, need adopt reversible PF ODE or reverse-time SDE, both of which require
17

---

Published as a conference paper at ICLR 2025
pre-estimation of the unknown scores via score matching. When the score of the target distribution is not analytically available, sliced score matching (Song et al. 2020 2021a) can be used:
(x) [|| [2,, , 1) - V log (xx)
= (m) 5 +) - 2s (z, t) V (o) + V logp:(z))
(33)
(m)) [I| (z) (x. } + 2V sa t]] +
= arg min (.r) [1| t) + 2V s0(x,
(34)
The relative abundance of samples contains the gradient information of the log-density in the data space. This approach only uses data samples to estimate the correct score. The ultimate goal of both reverse-time PF ODE and SDE is map from the prior Gaussian distribution (x( to the target data distribution (z).
Flow matching, Different from score matching, another approach learn the vector field Φo (x, t) Rd-1 - Rd that helps transform the prior distribution PG to the data distribution Pdata This method bypasses the unknown V log and directly learns the complete drift term Φ9 = 99 V log p. resulting in the following equation:
dæ dt (x.
(35)
We can train the network (y in various ways. For instance, continuous normalizing flows (Chen
et al. 2018) match the distributions at the endpoint and optimize the network parameters by mini- mizing the Kullback-Leibler (KL) divergence:
d log p(x(())))) - (0(
(36)
dt
0 arg min KL (x))(x))
(37)
0
where (((rr) (ar) is transformed from Paama by the instantaneous change of variables. Alternatively, we can match trajectories and vector fields at intermediate times, such as flow matching (Lipman al. 2022 Liu et 2022):
0" arg min -e--
do (2r, 1)
(38)
where un is the desired transport velocity and usually produces the straight paths.
Generative models. We provide new perspective on the relationship between diffusion models and GANs. the denoising process of image generation treated as gradient ascent along log Pt (ar), then score- based diffusion models move x: to increase while gradually reducing time f, and finally samples under the distribution Po == Pdata- For generative adversarial networks (GANs) (Goodtellow et al. 2014| Huang al. 2225), the discriminator directly models the data distribution Pdata := Po. and the generator training can be regarded as performing gradient ascent Pu in the data space to generate high-probatee probabbbbyylltyy samples.
generate
These methods generate samples from (z), but optimization spaces are different, resulting in different paths. Diffusion models uses trajectories induced by {pr(x), t: T + 0}, while GANs are guided by (z:).
FPNN.
The SFP equation is related to known SDE da (((x) + adw. Traditional PINNs directly minimize the PDE residual to train an approximate solution. Our approach reformulates the forward SDE into the PF ODE:
dæ
(1()) - -000
o T V logp(z)) dt f(x)ddt
(39)
and introduces the score function associated with the solution of FP equation. Due to the stationary invariant property, we can set f(x) to train the network and this reduces to Eq. 21 for J
T
In the PF ODEs, the score of generative models and the drift of FP equations share formal con- sistency. The drift term f(x,, t) non-zero in the former and drives towards the data distribution, while f(x) = maintains a steady state in the latter. The former is data-driven, whereas the latter guided by physical laws.
18

---

Published as a conference paper at ICLR 2025
(a) Score-based generative modeling
(b) Fokker-Planck neural network
Forward SDE
forward Addrigrciin
SDE dx (xxxtt (())d
fxx, r)dt + (()
Reversible PF ODE
PF ODE
dx [(xxt) 2 (e)olo de
boclovent
dx =
(xx 2 (x)loggp))) - (207(0)
-V-(20(00
Reverse SDE
lankonend
Score PDE loss
dx (x,, - (e)Vlog p(x, c]]dt +g(+)dw Dessidee
Eaco (Fogg(()))))) f(x) V -(0)
Data-driven: (r) [(x:1)
Pilnta
Physical law: f(x)
Poots
Figure 12: The tasks of image generation and SFP equations.
B COMPUTATION OF PARTITION FUNCTION
The main proof is derived from Section 3 of Wang al. (2022). Here, we demonstrate the simplified integral expression relevant to our work, and the conclusion remains valid for the power embedding.
Theorem 2 (Partition function calculation). Given the density pa( parameterized by TNN, the partition function Zo computed as,
Zg
Σ mT,
wi
(ni) fas
(2)
(elma);
(40)
where the nodes {a} (mi) N; dimension (i 1....,d) d).
In; and weights {w:
(n Ina=1 are used for numerical integration in the i-th
Proof. For the unnormalized density Po (x)) represented by TNN, the computation of Zo can be decomposed into d one-dimensional integrals:
Zo
po(x)dx
52
/
II
J (a(2x)da
II
(41)
m=
Σ L
1 bi
....
f15(21) fas dra
"(("" f (1. fds(daadda
(42)
For simplicity, we omit the parameters f in network f.. (q.(4I) holds due to the linearity of integral and Eq. (2) follows the fact that the sub-network solely depends only æ; (i 1,..., d). With- out loss of generality, for the i-th dimension, we use a piece-wise Gauss-Legendrre quadrature rule to compute the one-dimensional integral fh: (x;(4;) and then compute Zo in Eq. Specifi-
[ No (or(me) NN =m.=1 and corresponding weights (m) =n.=1 introducing cally, we select N; Gauss points {a.
the index (M1,**. na) E N = (I....,Ni N,) x x {1,..., .2). Hence, Gauss points and
19

---

Published as a conference paper at ICLR 2025
weights for numerical integration over the high-dimensional domain !2 are represented as:
JnEN
JEEN
(11, ;aa)" I, (n4) =1,,,.... I....,
(43)
(wo
IdEN
wi w4 = win
(n.) , I.....d
The partition function computed by numerical integration:
Zo
po(a))d
w w
(z)
2
REN
Σ
wi
wa
(na) II fi (e((t))
(mir" ma)EN
j=li=1
Σ
(m))) (winiifs fig (2) (ef(u))) ... ( fai (al
F (Ne maaeN
Ni
Na
w (31)
fiej (offul))
Σ m1=1
(ma) annnaaas sai (2ndee))
(44)
Since the cost of each one-dimensional integral independent of dimensionality, the overall computational complexity for the nigh-dimensional integration is the linear scale of dimension d. And we only calculate the partition function once in FPNN.
The advantages and differences between two computations of Ze are summarized follows:
MLP-hased FPNN. The advantages of MLP lie in simpler feature fusion of spatial vari- ables and faster computation of partition functions Za using MC sampling:
(
Zo
po(x)
(45)
ID
# Doum
However, in particularly nigh-dimensional settings or the large range of each interval in Ω, the volume of Ω2 i.e., v(2) II 16, a,l may exceed machine limits. This issue can be addressed by narrowing the intervals of interest or using higher numerical precision, such 'float64'.
TNN-based FPNN. TNN achieves high precision in numerical integration for estimating Zg. making it more suitable for higher-dimensional problems. Unlike MLP, TNN can maintain appropriate values for the integral in each spatial dimension (i.e., fhh fig (x;))d, regardless of the dimensionality or the interval length. Introducing regularization terms to guide the learning of smaller unnormalized density is an effective strategy to alleviate the numerical difficulties associated with the partition function.
C EXPERIMENTAL DETAILS OF PDES
We remark that FPNN is adaptable to wide range of SFP equations, but most problems difficult to obtain analytical solutions for performance evaluation. Inspired by Wang et al. (2024). we use gradient systems to construct test PDE cases with exact solutions. Given any potential function H (x) and diffusion term o we calculate the diffusion matrix D(x) (x) and set the drift term follows:
[(xx) D(x) D(x))
(46)
20

---

Published as a conference paper at ICLR 2025
Table 6: Network structure settings used for visualizing experimental results.
SFP equations
Model
Layers
Parameters
Steps
TFFN
4x[1, 64, 64, 64]
33,792
20k
4D Ring
TNN
4x(8, 64, 128]
35,584
10k
MLP
20, 20, 20, 20]
1,360
10k
TFFN
6x[1, 64, 64, 64]
50,688
20k
6D Unimodal
TNN
6x[5, 64, 64, 64]
52,224
5k
MLP
[6, 20, 20, 20]
980
20k
TFFN
6x[1, 64, 64, 64]
50,688
20k
6D Multi-modal
TNN
6x[5, 64, 64, 64]
52,224
5k
MLP
[6, 32, 32, 32]
2,336
5k
TNN
10x [1, 64, 64, 64]
84,480
1k
10D Multi-modal
MLP
[10, 64, 64, 64, 64]
13,184
5k
TNN
10x [1, 64, 64, 64, 64]
126,080
5k
Gaussian mixture
MLP
[10, 20, 20, 20, 20]
1,480
5k
Gaussian
TNN 20x [1, 64, 64, 64, 64]
252,160
3k
Then the solution to the SFP equation Lp(x)) 0 with respect to (p, D) reads
p(x)= = Z
1 (())
Z
H(()) dx
(47)
It can be easily verified that
Cp(x) = (p(x))u(()) V V (D(())))((=))
V (-p(x)) (V D(=))p(a) V(z))(()))
(48)
V (- + (V ((=)) - D(x))H(())(()))
0
The stochastic system is given by dX (X)dtt (X)dW, We adopt this form design SFP For exact solutions in our experiments, we use the SymPy library to compute the partition function in Eq. (47). SymPy is powerful and versatile tool for symbolic mathematics and provides computer algebra system (CAS) capabilities directly in Python. We point out that appropriate simplifications (e.g., completing the square and variable substitution) and the choice of integration order significantly impact the accuracy of results.
equations C.2 C.4
In Table we list the network settings used for plotting experimental results, where TNN and MLP
represent the architectures in FPNN framework using score PDE loss.
C.1 4D RING
Consider a stochastic gradient system four dimensional state space from Zhai et al. (2022) The
deterministic part is gradient system plus perpendicular rotation term r (x) [z2, where the potential function the gradient flow is H (x) = 2[[|=] - 1) We use the SFP equation 2 with the following drift ((x: and diffusion D:
31, 0, 0)T,
((x) (()) r(x),
D
Lt a= a
(49)
Since rxx) is orthogonal to the equipotential lines of (xx), the rotation does not change the invariant probability density function. The invariant probability measure has density function
1 p(x) = H(w))o²
== 1T (t 1)e -22/ /02
(50)
%
21

---

Published as a conference paper at ICLR 2025
004
004
00000
622
80
64
-13
90
13
(TNN)
-13
-13
00
-15
12
-13
80
15
(MLP)
064
0000
0009
3015
-15
-15
-15
008
0000
uce
-15
-15
30
-15
25
-15
00
Figure 13: 4D Ring: the exact solution and predicted solutions of different models. First row: the exact solution. Second row: TNN-based FPNN. Third row: MLP-based FPNN. Last TFFN.
FPNN (TNN) XE8-
FPNN (MLP)
TFFN X30-2
Xx, 0,01
1.20
0.55
SE
28
0.60
24
10
15
0.30
0.45
05
0.15
-15
0.0
0.
-15
3.0
0.0
0.
0.00
(X,,X,, 0.5. 0.5)
k30
025
0.60
0.45
13
035
-15
15
0.0
0.
15
0.0
0.0
0.0
0.00
Figure 14: 4D Ring: absolute errors between the exact solution and predicted solutions of different models. First column: the exact solution. Second column: TNN-based FPNN. Third column: MLP. based FPNN. Last column: TFEN.
Here, utilize the polar coordinates, variable substitution, and the surface formula of 4D hypersphere S 2n 2,3 We set = and calculate the infinite integral for Z using SymPy library.
We generated training dataset Durai with 20k samples using the 1.5-order SRK method, setting the terminal time = 1, the number time steps = 500 and a dimmenssion standard Gaussian dis- tribution (o Based on the range of the dataset Drainn the domain is selected Ω = [-1.8, 1.8)
For the test dataset generation, we sample 10k initial points uniformly within Q. Then we perform the gradient ascent using the true solution p(ar), with learning rate of 10-3 and the threshold
22

---

Published as a conference paper at ICLR 2025
== x=xx1 10-3. For fair comparison, we use the same test dataset Diest for TENN and FPNN with TNN/MLP architectures.
For equations in our experiments, FPNNs are trained under consistent settings: we use Adam optimizer with a learning rate of 0.01 and a batch size of 2k, resulting in 10 iterations per epoch. The network structure and prediction performance sub-networks of 3 hidden layers with 64 hidden feature size, updated for 20k steps using the Adam optimizer with a learning rate of 0.01, as shown in Table 6 Training data are uniformly sampled from Z, with 2k points resampled per iteration.
detailed in Table 4 TFFN consists of four
C.2 6D UNIMODAL
In this case, the diffusion D- = Ig and the drift uu(x)) = (a) take the form in Eq. (6). where the potential function (z) is written as:
H(x) = 3(( - x²)² + 2n + (23 - xA)² +
+ 2x + (ag x²)² + 2x
(51)
The variables (#1 22). (.), and (I5, have distinct interactions in the form of (x²² 24+1 ) + 2, i = 1,3, 5. The specific term results in u(x:) being a polynomial function of degree up to 7, introducing significant nonlinearity and challenges for training model. Our proposed PE is inspired by this complexity, and the goal is to alleviate the difficulty of fitting such high-order functions. The density function is more concentrated with steep gradients, and decays rapidly away from the origin. Although the density is unimodal, it remains a challenging problem for solving this SFP equation.
Uog. 0,0.0,0)
(t. 0.xxs 0, 07
(0.X,, Xy. 0,0,0) D)
Kg. 0.X,
Eolaion
035
175
0155
675
116
F1
FPAIN (TNN))
150
FRRN (MLP))
175
230
TFRN
Figure 15: 6D Unimodal: the exact solution and predicted solutions of different models. First row: the exact solution. Second row: TNN-based FPNN. Third row: MLP-based FPNN Last row: TFFN.
We set the terminal time T = and the number of time steps N = 500. The initial distribution is a 6-dimensional Gaussian distribution with zero mean and covariance matrix 0.01 And we generate 20k training points and select the domain Ω = [-1.2, 1.2 The test dataset is created by applying gradient ascent to 10k initial points from uniform distribution within Ω with a learning rate of 1 1x 1x1003 and a threshold e = 1x100 5 We use the same training settings for FPNN, and comparisons of different network architectures are shown in Table 5 Notably, PE significantly accelerates the
23

---

Published as a conference paper at ICLR 2025
learning of higher-order densities for TNNs. With appropriate embedding size m, TNN achieves MAPE below 10% with fewer steps, while MLP demonstrates its advantage by solving the 6D Unimodal problem with only 980 parameters.
FPNN (TNN)
(MLP)
(Kg, Kg. 0,
110
3.35
x20
0.40
330
0.35
830
30
0.30
025
24
02
643
18
0.35
030
0.30
0.25
0.50
(x. 0.Kx 0, 0)
0.36
0.22
36
634
0
R
020
0.16
0.58
0.24
Figure 16: 6D Unimodal: absolute errors between the exact solution and predicted solutions of different models. First column: the exact solution. Second column: TNN-based FPNN. Third column: MLP-based FPNN. Last column: TFFN.
C.3 6D MULTI-MODAL
We consider the potential (x) with a quadratic form, where the SFP equation given by the diffusion D Is and the drift μux) = (x) in [4.4 (x) involves interactions among (I1, T2, 23) and (I 2z, 26), which possesses the following form:
H (x) = + z + zड 0.5 21 x203)) In(x} 0.02)
(52)
In 0.02) + 0.5 + = 0.2 2446 + x5*6))
The density function has four modes and we plot the slices of high-density regions from different dimensions to illustrate the performance.
Table 7: 6D Multi-modal: comparison of FPNN with different networks (5k steps).
Model
Layers
Parameters
MAE
MAPE
[yn, hidden layers, r]
[1, 64, 64, 64] [5, 64, 64, 64]
50,688
4.20 x 10
35.76%
TNN
52,224
3.25 x 10
23.11%
[8. 64, 128]
53,376
2.30
10-4
21.59%
[d, hidden layers]
[6, 20, 20, 20] [6, 32, 32, 32]
980
2.00 10
12.72%
MLP
2,336
2.05 x 10
12.72%
[6, 64, 64]
4,608
1.98
10
12.18%
We set the terminal time T = 1 and the number of time steps N = 200. The initial distribution a 6-dimensional Gaussian distribution with zero mean and covariance matrix 0. 0.1 We generate 20k training points and determine domain Ω [-2, 2/6 for 6D Multi-modal problem. We learning rate of x 10-3 and threshold = 1 10 to generate the test dataset with 10k samples by gradient ascent method.
From the numerical results Table 7 MLP outperforms TNN terms of parameter efficiency, training speed and prediction accuracy, achieving MAPE of 12%. This may be attributed to the architectural characteristics of TNNs, which is more suitable for the solutions align with the variable separation form and quickly learns discrete low-rank features (see Table 5). Conversely, MI.P is
24

---

Published as a conference paper at ICLR 2025
better at fusing features across spatial coordinates in hidden layers. FPNN allows the of free- form network architectures, providing greater flexibility in handling SFP equations with diverse characteristics.
0, 0,
[x,, 0, O,
0,
10, xx 0,
Elsd
0033
300
800
(TNNN)
615
504
130
18
IMLP!
212
60112
301
306
504
15
15
000
TFPN
3300
00036
15
15
40
Figure 17: 6D Multi-modal: the exact solution and predicted solutions of different models. First row: the exact solution. Second row: TNN-based FPNN. Third row: MLP-based FPNN. Last row: TFFN.
FPNN (TNN)
FPNN (MLP)
TFFN
(K) X2. 0, 01
220
VID
1.5
15
3.5
10
3.0
1.0
-15
-1.
-35
06
825
052
-15
0.
-19
3.5
82
XL
(X.. 0, Xs, 0, 0,
1.3
27
3.5
+10
155
100
3.05
00000
675
1927
1.5
12
19
0.0
0.
8.15
4
-1.
1.5
0.0
3.5
0.2
-15
3.5
Figure 18: 6D Multi-modal: absolute errors between the exact solution and predicted solutions of different models. First column: the exact solution. Second column: TNN-based FPNN. Third column: MLP-based FPNN. Last column: TFFN.
C.4 10D MULTI-MODAL.
To evaluate the performance in high dimensions, we conduct experiments using a 10-dimensional probability density function with two modes. The SFP equation has the diffusion D T1o, and the
25

---

Published as a conference paper at ICLR 2025
potential function (x) for the drift ju(x)) as follows:
(()) =2. + 2 + z + 0.1 + 2143 + x203)) 2 + + 25 + 2 0.2 + 1446
+ x5556) + 3(0 + z 0.01 + 3 + =10 - 0.01..9310)) - In(2ng 0.02)
(53)
[0. 0. X1o
(0, Xs, 0.5, Xs0
10. 0, 03
(0. O. 0.5.
Evsict Boluton
11
035
0.5
0.5
FFNN (TNN)
0.3
224
618
212
036
FPNN (MLP)
833
412
612
Figure 19: 10D Multi-modal: the exact solution and predicted solutions of FPNN. Top: the exact solution. Middle: TNN-based FPNN. Bottom: MLP-based FPNN.
FPNN (TNN)
FPNN (MLP)
(0, O, Xg, X10)
-10-2
X10
1-
1.6
1:35
1.4
120
0.20
1.0
0.10
0
0.75
0.05
C.
0.00
0.45
0.30
-1
X50
C2
-1
8.0
0.00
X
(0. 0, Xg. 0.5, X10)
x10
x10
6.0
5.6
5.4
48
0.15
4.0
3.6
42
0.10
32
300
0.
2.4
24
1.5
1.
12
810
Xg
-1
Figure 20: 10D Multi-modal: absolute errors between the exact solution and predicted solutions of different models. Left: the exact solution. Middle: TNN-based FPNN. Right: MLP-based FPNN.
Here, (xx exhibits complex interactions across the coordinates (21, 1.)), (ma. 25, T6), (#7, 23) and (xy, 210) Due to the component (n(2 + 0.02), there are two peaks along variable Dg. By plotting the cross-sectional view of the last two dimensions, we can observe this shape.
We generate a training dataset of 20k points using the SRK method, with a terminal time T = and N 100 time steps. The initial distribution po(x:) is 10-dimensional standard Gaussian
26

---

Published as a conference paper at ICLR 2025
distribution and the domain is defined as [-1.2. The test dataset is created using learning rate of X 10 and a threshold of e = 1x10 ensure that densities of all 10k points exceed e. In FPNN, both TNN and MLP architecture successfully learn the two peaks in the 2g axis, while PINN and TFFN fail to capture this feature. It is evidenced by the divergence between model predictions and true solutions for TFFN in Figure The normalization condition influence the representational capacity and optimization process of models, as discussed in Sec. This results in poor performance of existing methods for 4-6 dimensional SFP equations and failure in higher-dimensiona problems.
C.5 10D GAUSSIAN MIXTURE
Following the test problem with two peaks Tang et al. 2022), we construct Gaussian mixture distribution p( (x) 31p1 (x) + (2)2(x), comprised of two I0-dimensional Gaussian distributions Pı and p2. For & I, 2, each pk (x) defined by probability density function with mean pa and covariance matrix Σμ:
2111
E12
81
0.55,
#I
[-.,5, 1.3,0.2 -.1,0000000, ΣΙ
E1
Σ14
2.2
0
0
1.5
0.2
0.4
0.4
0.3
1
0
Σ1
1.2
Σ12 =
0.1
1.2
0.4
E1 =
Σ11
0.3
0.9
0
0
2
0.2
0.2
0.8
(54)
Σn1
En2
32 = 0.45,
P2 = [1.2, - 1,,,,,,,11,,,,,,,,00 0,,00,,,,,,0011,,,,,,0000
Σg
Σ23
Σn1
2.2
1.2
0.4
0.2
0.8
0
Σn1
1
0
Σn
-0.4
1.2
-0.3
Eg3
E24
0.3
0
1
1.5)
0.2
0011
1.2
(55)
The matrices Σ, and Eg are positive definite. FPNN remains applicable for more randomly con- structed Gaussian components. To improve visualization, we make some modifications the mean (11, (42- First, two Gaussian peaks are placed at a certain distance apart and have different relative positions on the (I1, x2) and (3) planes. Secondly, we the last seven components of the mean near zero, so that we examine the cross-sections (æ1, 22, O.... 0) and (I1, 0, T3, .,..., 0) to reflect the high-density regions of exact solution. is straightforward to verify that p(æ)) serves as the true solution to following SFP equation:
logz))Voo($sppec + (2p2(x))) + ppx(x==
(56)
where we get the corresponding drift (xx) (o) 3pp2(x)) and diffusion D I1o-
As the region expands, the stochastic system takes longer to reach a steady state. select terminal time of T and N 500 time steps. 20k data points are sampled from 10-dimensional standard Gaussian distribution serve as the initial set, followed by the SRK method to obtain the target distribution pr and domain = [-5, 5]10 The test dataset with 10k points sampled from two Gaussian distributions according weights.
Notably, even though the true solution has the order 10-5, simple network architectures such as TNN and MLP successfully learn the two Gaussian peaks located on the coordinate planes (I1, Z2, O, ...... 0) and (I1, 0, 23, .......0 0) at 600 and 400 steps, ultimately achieving a mean relative error of 1.86 x 10 and 1.16 10-? respectively. And MLP reaches MAPE of 18.38% with only 1,480 parameters. The final prediction results and absolute errors are presented Figure 21
and Figure
27

---

Published as a conference paper at ICLR 2025
Exact Solution
FPNN (TNN)
FPNN (MLP)
(k1, X2. O 0)
0
16
08
0.8
-4
(4s, 0,X3. 0. 0)
24
24
16
166
88
08
Figure 21: 10D Gaussian mixture: the exact solution and predicted solutions of FPNN. Left: the
exact solution. Middle: TNN-based FPNN. Right: MLP-baaed FPNN
FPNN (TNN)
FPNN (MLP)
[x1, X2. ....,
X10-1
4.0
3.5
1775
2.5
3.0
2.5
125
R
R
1.00
0.5
1.5
0.75
0.50
0.5
0.0
0.00
X1
X1
(X1, 0, O, 0)
x10
5.4
x10
16
42
3.6
12
R
R
2.4
0.8
12
04
0.6
02
43
X1
Figure 22: 10D Gaussian mixture: absolute errors between the exact solution and predicted solutions of different models. Left: the exact solution. Middle: TNN-based FPNN. Right: MLP-based FPNN.
C.6 20D GAUSSIAN
Finally, we consider a SEP equation with drift term μ(x) = -ax and diffusion matrix D The exact solution is expressed follows:
a²
Ia-
2
d/2
a a||=
p(x) = T
exp
(57)
In our experiments, we select d = 20, = 3 and σ 1.5 to test our FPNN with TNN network.
So the true PDF is a 20-dimensional Gaussian distribution with zero mean and covariance matrix
g² /20. The training dataset is generated by the SRK method, yielding 20k samples with T = 1 and N 100. Since the initial distribution is chosen as a 20-dimensional standard Gaussian distribution, the drift effeet towards the origin dominates, causing particles to converge towards ! [-2, 2) during the SDE simulation. The test dataset directly constructed by sampling 10k points from the true distribution.
28

---

Published as a conference paper at ICLR 2025
FPNN (TNN)
Exact solution
0.5k steps
1K steps
1.5k steps
16
00000
-15
80
15
09
15
-15
18
15
-13
80
15
15
33005
03
22
60
0.
Figure 23: 20D Gaussian: the predicted solution of FPNN during training.
After 1.5k update steps, FPNN successfully learns a 20-dimensional Gaussian distribution, achieving a MAPE 23.29%. shown in Figure I0 each dimension aligns well with the true solution, demonstrating both the effectiveness of our method and the validity of evaluation metrics.
15
40
013
42
Figure 24: 20D Gaussian: absolute emoos between the exact solution and predicted solution of FPNN. The first column shows the exact solution and the changes of MAPE during training. The remaining parts display errors for pairs of adjacent dimensions, such as (I1, 02,00............. .)... (0, 0, 23, F4, O, .,, 0), and
29


Reviews:
[{"review_id": "mUntBkG2td", "reviewer": "Reviewer_1JZg", "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.", "questions": "**Main comments**\n1. 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?\n2. 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?\n3. 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?\n4. 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?\n\n**Minor comments**\n1. While the paper is well written, some typos are remaining, e.g. \"Score-based generate model\" (line 129)\n2. 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.\n3. 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. \n\n\n**References**\n\n[1]\n@inproceedings{zhai2022deep,\n  title={A deep learning method for solving Fokker-Planck equations},\n  author={Zhai, Jiayu and Dobson, Matthew and Li, Yao},\n  booktitle={Mathematical and scientific machine learning},\n  pages={568--597},\n  year={2022},\n  organization={PMLR}\n}\n\n[2]\n@article{lu2021deepxde,\n  author  = {Lu, Lu and Meng, Xuhui and Mao, Zhiping and Karniadakis, George Em},\n  title   = {{DeepXDE}: A deep learning library for solving differential equations},\n  journal = {SIAM Review},\n  volume  = {63},\n  number  = {1},\n  pages   = {208-228},\n  year    = {2021},\n  doi     = {10.1137/19M1274067}\n}\n\n\n[3]\n@article{cho2024separable,\n  title={Separable physics-informed neural networks},\n  author={Cho, Junwoo and Nam, Seungtae and Yang, Hyunmo and Yun, Seok-Bae and Hong, Youngjoon and Park, Eunbyung},\n  journal={Advances in Neural Information Processing Systems},\n  volume={36},\n  year={2024}\n}", "limitations": "", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper is overall well written and easy to follow.\n2. 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.\n2. 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].\n3. 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]?\n\n(see refs in questions)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cYAfZN9KlT", "reviewer": "Reviewer_aoGx", "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.", "questions": "1. Could be the proposed method widely used for other PDEs?\n2. Please compare with recent developed methods.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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\n2. Lack of Scalability and Computational Complexity Analysis\n3. Lack of Comparison with Other Recent Advances", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SuAnjynuXg", "reviewer": "Reviewer_hrfk", "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.", "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.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\n(2) The performance results are promising, though further validation is needed to strengthen the findings.\n\n(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.\n\n(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.\n\n(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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rB3tNfr7b7", "reviewer": "Reviewer_eksf", "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.", "questions": "See weaknesses above. Also, when using MAPE for metric, how do you calculate $p(x)$ for test dataset?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\n**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:\n\n1. 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.\n\n2. 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? \n\nBased 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)