Organizations: School of Mathematics and Statistics, Shaanxi Normal University, Xi’an, 710119, China · School of Mathematics and Statistics, Northwestern Polytechnical University, Xi’an, 710129, China · MOE Key Laboratory for Complexity Science in Aerospace, Northwestern Polytechnical University, Xi’an, 710072, China · School of Science, Xi’an University of Posts and Telecommunications, Xi’an, 710121, China · Department of Systems and Control Engineering, Institute of Science Tokyo, Tokyo, 152-8552, Japan
Efficiently solving the Fokker-Planck equation (FPE) is central to analyzing complex parameterized stochastic systems. However, current numerical methods lack parallel computation capabilities across varying conditions, severely limiting comprehensive parameter exploration and transient analysis. This paper introduces a deep learning-based pseudo-analytical probability solution (PAPS) that, via a single training process, simultaneously resolves transient FPE solutions for arbitrary multi-modal initial distributions, system parameters, and time points. The core idea is to unify initial, transient, and stationary distributions via Gaussian mixture distributions (GMDs) and develop a constraint-preserving autoencoder that bijectively maps constrained GMD parameters to unconstrained, low-dimensional latent representations. In this representation space, the panoramic transient dynamics across varying initial conditions and system parameters can be modeled by a single evolution network. Extensive experiments on paradigmatic systems demonstrate that the proposed PAPS maintains high accuracy while achieving inference speeds four orders of magnitude faster than GPU-accelerated Monte Carlo simulations. This efficiency leap enables previously intractable real-time parameter sweeps and systematic investigations of stochastic bifurcations. By decoupling representation learning from physics-informed transient dynamics, our work establishes a scalable paradigm for probabilistic modeling of multi-dimensional, parameterized stochastic systems.
Figures & tables
Figure 1: The architecture of the TPAPS q(x,t;Θ,pI) for jointly solving parameterized FPEs with diverse initial distributions and control parameters.
ResNet ID
EMB-RN
REP-RN
EVO-RN
DEC-RN
Input
{(μij,σij)}i=1DSTA
Embedding h
(t,Θ,HI)
Transient Ht(Θ,HI)
DIn / DMid / DOut of RB#1
2DSTA /50/50
DEMB /100/100
1+DPAR+DREP /100/100
DREP /100/100
DIn / DMid / DOut of RB#2
50/50/50
100/100/100
100/100/100
100/100/100
DIn / DMid / DOut of RB#3
50/50/ DEMB
100/100/ DREP
100/100/100
100/100/100
DIn / DMid / DOut of RB#4
100/100/100
100/100/100
DIn / DMid / DOut of RB#5
100/100/100
100/100/100
Table 1: The architectures of the ResNets in the TPAPS
Figure 2: Visualization of the GMD autoencoder. (a) Six GMDs and their reconstructions. (b) The PCA and distance visualizations of the 50-D embedding space. Each point represents a 2-D Gaussian distribution in Eq. ( 45 ). The three lines consist of continuously changing distributions, where μ1=[−2,−2]⊤ , μ2=[2,2]⊤ , and I is the 2-D identity matrix.
Figure 3: The transient solutions of the 1-D system ( 49 ) by TPAPS (colored circles) and MCS (black lines) with different initial distributions and system parameters.
Figure 4: The training losses and test errors of the TPAPS on the 1-D system. The boxplots in (b) are summarized on 2,500 transient solutions with random initial distributions and system parameters.
Time
103 batches
104 batches
105 batches
106 batches
t=0.5
0.4257/0.3561/0.3161
0.2062/0.1906/0.1476
0.0595/0.0923/0.0345
0.0212/0.0477/0.0133
t=1.5
0.3504/0.3701/0.2056
0.1794/0.2354/0.0798
0.0607/0.1291/0.0232
0.0217/0.0645/0.0102
t=3.0
0.3637/0.3745/0.2050
0.1648/0.2340/0.0685
0.0619/0.1410/0.0230
0.0234/0.0717/0.0098
Table 2: The mean, SD, and median (MEAN/SD/MED) of the L1 errors between the TPAPS and MCS results across 2,500 random transient solutions of the 1-D system ( 49 ) at various training stages.
Figure 5: Bifurcation study of the 1-D system ( 49 ) by using TPAPS and MCS. The bifurcation parameters are (a) the combination weight λ in the initial distribution, (b) the system parameter a , and (c) the noise amplitude σ . The resolution of each figure is 100×100 .
Figure 6: Four transient solutions of the 2-D subcritical system ( 52 ) by TPAPS and MCS with different initial distributions and system parameters. The L1 errors are displayed on the residual images.
Figure 7: Visualization of transient solutions of the 2-D subcritical system ( 52 ) with t∈[0,20] in the 1st, 2nd, and 3rd PCs of the 50-D embedding space. The 9 transient solutions initialize from the distributions pIA , pIB , and pID in Fig. 6 with system parameters ΘA , ΘB , and ΘD and converge to three stationary distributions pSA , pSB , and pSD , which are attaching to the respective system parameters. For better illustration, similar to Fig. 2 , the points in the background are embedding features of random Gaussian distributions N([μ1,μ2]⊤,diag{σ12,σ22}) , and the two lines of distributions G1 and G2 are defined in Eq. ( 47 ). The grey dashed lines indicate the axes of the three leading PCs.
Figure 8: The training losses and test errors of the TPAPS on the 2-D subcritical system ( 52 ).
Time
103 batches
104 batches
105 batches
106 batches
t=0.5
0.3442/0.1681/0.3007
0.1501/0.0679/0.1348
0.0846/0.0365/0.0750
0.0689/0.0302/0.0610
t=1.5
0.3393/0.2051/0.2731
0.1535/0.0891/0.1273
0.0954/0.0451/0.0848
0.0661/0.0365/0.0558
t=3.5
0.3417/0.2314/0.2566
0.1511/0.1115/0.1119
0.0985/0.0535/0.0842
0.0670/0.0438/0.0547
t=6.5
0.2978/0.2333/0.1985
0.1444/0.1347/0.0912
0.1000/0.0712/0.0748
0.0662/0.0503/0.0503
t=10.5
0.2439/0.2101/0.1502
0.1431/0.1650/0.0801
0.1036/0.0926/0.0666
0.0649/0.0551/0.0462
Table 3: The mean, SD, and median (MEAN/SD/MED) of the L1 errors between TPAPS and MCS results across 2,500 random transient distributions of the 2-D subcritical system ( 52 ) at various training stages.
Figure 9: Efficient study of transient and near-stationary ( t=20 ) dynamics of the 2-D subcritcal system ( 52 ) using the TPAPS. (a) shows an example of marginal transient distributions of x at different times. (b) shows the Hopf bifurcation of stationary distributions triggered by the system parameter μ by visualizing the marginal stationary distributions of x . (c) investigates two-parameter bifurcation of stationary distributions depending on noise intensity D and system parameter μ in terms of the SD σx of the marginal distribution of x and (d) demonstrates the two-parameter bifurcation of the non-ergodic distributions depending on μ and the x-coordinate x0 of the center of initial distributions. (e) compares several groups of the SD σx of the distributions obtained by TPAPS and the MCS.
Figure 10: Four transient solutions of the 2-D Duffing system ( 53 ) by TPAPS and MCS with different initial distributions and system parameters. The L1 errors are displayed on the residual images.
Figure 11: The training losses and test errors of the TPAPS on the 2-D Duffing system ( 53 ).
Time
103 batches
104 batches
105 batches
106 batches
t=0.5
0.9115/0.2776/0.9150
0.5292/0.2185/0.5053
0.3340/0.1317/0.3168
0.2331/0.1051/0.2166
t=1.5
0.7175/0.3149/0.6793
0.4036/0.2224/0.3551
0.2549/0.1435/0.2194
0.1892/0.1140/0.1589
t=3.5
0.4466/0.2799/0.3853
0.2747/0.2043/0.2143
0.1967/0.1413/0.1585
0.1524/0.1159/0.1160
t=6.5
0.2624/0.2023/0.1971
0.1528/0.1415/0.1027
0.1139/0.0995/0.0810
0.0963/0.0950/0.0630
t=10.5
0.1693/0.1353/0.1152
0.0969/0.0860/0.0648
0.0666/0.0586/0.0457
0.0600/0.0623/0.0364
Table 4: The mean, SD, and median (MEAN/SD/MED) of the L1 errors between TPAPS and MCS results across 2,500 random transient distributions of the Duffing system ( 53 ) at various training stages.
Figure 12: The TPAPS results of the Duffing system ( 53 ) at early training stages. The initial distribution, system parameters, and the results at 106 training batches have been detailed in Fig. 10 (d).
Figure 13: Efficient study of transient and near-stationary ( t=20 ) dynamics of the Duffing system ( 53 ) using the TPAPS. (a) shows marginal transient distributions of x at different times. (b) shows the impact of the noise amplitude σ on the marginal stationary distribution of x at t=20 . (c) investigates the dynamics of stationary distributions depending on the damping parameters η and the stiffness parameter β . The true transient solutions are obtained by MCS in (a2) and the true stationary results in (b2) and (c2) are calculated by Eq. ( 54 ). The resolution of all subfigures is 200×100 .
Method
105 training batches
106 training batches
seconds/batch
t=0.5
t=1.5
t=3.5
t=6.5
t=10.5
t=0.5
t=1.5
t=3.5
t=6.5
t=10.5
Baseline
0.3340
0.2549
0.1967
0.1139
0.0666
0.2331
0.1892
0.1524
0.0963
0.0600
1.11
TPAPS-A
0.2296
0.2290
0.1768
0.1301
0.1098
0.1896
0.1994
0.1459
0.1204
0.1037
1.16
TPAPS-B
0.3263
0.2598
0.1968
0.1098
0.0613
0.2532
0.2035
0.1616
0.0952
0.0541
0.48
TPAPS-C
0.3402
0.2695
0.2114
0.1340
0.0781
0.2875
0.2275
0.1695
0.0996
0.0596
0.67
TPAPS-D
0.3583
0.2835
0.2007
0.1067
0.0547
0.2910
0.2261
0.1606
0.0845
0.0446
1.08
Table 5: The mean L1 errors between several TPAPS and MCS results across 2,500 random transient distributions of the Duffing system ( 53 ) at different times, while the last column shows the training speed. Best values are bold.
The Fokker-Planck equation (FPE) plays a pivotal role in describing the time evolution of probability density functions (PDFs) for systems governed by stochastic dynamics. In this work, we propose a conditional normalizing flow-based physics-informed neural network (PINN) framework for efficiently approximating the solution operator of the FPE for a whole range of initial conditions. Leveraging the Chapman-Kolmogorov equation for Markovian stochastic processes, the problem is reformulated into approximating a transition PDF starting at initial time from a Dirac mass centered at an arbitrary point. The PDF of an associated linearized stochastic differential equation (SDE) is employed as the base distribution for the normalizing flow, providing a good approximation of the target PDF, especially for small times, and thereby avoiding the singularity of the map associated with the Dirac delta initial distribution. Furthermore, a time-weighted loss function is introduced to mitigate numerical instabilities arising at small times, achieving a balance between causality and training difficulty as time progresses. A variety of numerical experiments are presented to illustrate the effectiveness and robustness of the proposed method.
Li Zeng, Xiaoliang Wan, Yaobin Wang +2
School of Mathematics and Statistics, Fuzhou University, Fuzhou, China · Department of Mathematics and Center for Computation and Technology, Louisiana State University, Baton Rouge 70803, USA · Faculty of Science and Technology, Beijing Normal-Hong Kong Baptist University, Zhuhai 519087, China. +2
Solving high-dimensional Fokker-Planck (FP) equations remains a challenging problem in computational physics and stochastic dynamics, due to the curse of dimensionality, unbounded domains, and complex probability landscapes. In this work, we propose an adaptive probability flow residual minimization (A-PFRM) method for this problem. The second-order FP equation is reformulated as an equivalent first-order continuity equation associated with the probability flow ordinary differential equation, based on which a loss function is constructed to train neural network approximations without Hessian computation. To further improve the computational efficiency, the Hutchinson trace estimator is applied to compute the divergence in the corresponding score function, such that the training time can be dimension-independent on GPUs. Adaptive sampling strategies are employed to generate the collocation points, and our analysis shows that the Kullback-Leibler divergence between our A-PFRM approximation and the exact solution is bounded by the residual loss weighted by the estimated density function. Numerical experiments are presented to demonstrate the performance of A-PFRM, which include Ornstein-Uhlenbeck (OU) processes problems, Brownian motions with time-varying diffusion, and Geometric OU processes featuring non-Gaussian solutions up to one hundred dimensions.
Xiaolong Wu, Qifeng Liao
School of Information Science and Technology, ShanghaiTech University, Shanghai, 201210, China
We propose a new Neural Galerkin Normalizing Flow framework to approximate the transition probability density function of a diffusion process by solving the corresponding Fokker-Planck equation with an atomic initial distribution, parametrically with respect to the location of the initial mass. By using Normalizing Flows, we look for the solution as a transformation of the transition probability density function of a reference stochastic process, ensuring that our approximation is structure-preserving and automatically satisfies positivity and mass conservation constraints. By extending Neural Galerkin schemes to the context of Normalizing Flows, we derive a system of ODEs for the time evolution of the Normalizing Flow's parameters. Adaptive sampling routines are used to evaluate the Fokker-Planck residual in meaningful locations, which is of vital importance to address high-dimensional PDEs. Numerical results show that this strategy captures key features of the true solution and enforces the causal relationship between the initial datum and the density function at subsequent times. After completing an offline training phase, online evaluation becomes significantly more cost-effective than solving the PDE from scratch. The proposed method serves as a promising surrogate model, which could be deployed in many-query problems associated with stochastic differential equations, like Bayesian inference, simulation, and diffusion bridge generation.
Riccardo Saporiti, Fabio Nobile
CSQI, École Polytechnique Féderale de Lausanne, Switzerland