FlowKac: An Efficient Neural Fokker-Planck solver using Temporal Normalizing flows and the Feynman-Kac Formula
Authors: Naoufal EL BEKRI
Lucas Drumetz
Franck Vermet
Venue: Accepted by TMLR ()

Abstract: Solving the Fokker-Planck equation for high-dimensional complex dynamical systems remains a pivotal yet challenging task due to the intractability of analytical solutions and the limitations of traditional numerical methods. In this work, we present FlowKac, a novel approach that reformulates the Fokker-Planck equation using the Feynman-Kac formula, allowing to query the solution at a given point via the expected values of stochastic paths. A key innovation of FlowKac lies in its adaptive stochastic sampling scheme which significantly reduces the computational complexity while maintaining high accuracy. This sampling technique, coupled with a time-indexed normalizing flow, designed for capturing time-evolving probability densities, enables robust sampling of collocation points, resulting in a flexible and mesh-free solver. This formulation mitigates the curse of dimensionality and enhances computational efficiency and accuracy, which is particularly crucial for applications that inherently require dimensions beyond the conventional three. We validate the robustness and scalability of our method through various experiments on a range of stochastic differential equations, demonstrating significant improvements over existing techniques.

Published in Transactions on Machine Learning Research (08/2025)
FlowKac: An Efficient Neural Fokker-Planck solver using Temporal Normalizing flows and the Feynman-Kac Formula
Naoufal El Bekri
naoufal.cllekrr@@mmall.co
IMT Atlontique, Lab-STICC, UMR 6285, 29238, CNRS, Brest, France Unio Brest, CNRS, UMR 6205, Laboratoire de Mathémaliques de Brelagne Atlantique, France
Lucas Drumetz
lucas. wtw.ddmmmmmtliimmaaaannni.n.ppep
IMT Atlantique, Lab-STICC, UMR 6285, 29238, CNRS, Brest, France ODYSSEY Team- Projerooect, INRIA Ifvemer IMT- HITT-Atlattqqe, 35042, CNRS, Brest, France
Franck Vermet
franck. senukk.eereet@@uii-bbeer.
Unio Brest, CNRS, UMR 6205, Laboratoire de Mathématiques de Bretagne Atlantique, France
Reviewed on OpenReview: https: // openreview. net/ forum? id=paey(Faoor
Abstract
Solving the Fokker-Planck equation for high-dimensional complex dynamical systems re- mains a pivotal yet challenging task due to the intractability of analytical solutions and the limitations of traditional numerical methods. In this work, we present FlowKac, a novel approach that reformulates the Bokker-Plance equation using the Feynman-Ka formula, al- lowing to query the solution at given point via the expected values of stochastic paths. A key innovation of FlowKac lies in its adaptive stochastic sampling scheme which significantly reduces the computational complexity while maintaining high accuracy. This sampling tech- nique, coupled with a time-indexed normalizing flow, designed for capturing time-evolving probability densitics, enables robust sampling of collocation points, resulting in a flexible and mesh-free solver. This formulation mitigales the curse of dimensionality and enhances computational efficiency and accuracy, which particularly crucial for applications that herently require dimensions beyond the conventional three. We validate the robust and scalability of our method through various experiments range of stochastic differential equations, demonstrating significant improvements over existing techniques.
1 Introduction
The Fokker-Plande equation (FPE) Risken & Frank 1996) describes the time evolution of probability den- sity functions associated with diffusion processos, playing a central role in modeling stochastic dynamical systems across diverse domains, including physics (Lucia & Gervino 2015), finance (Sornette) 2001), en- gincering (Langtangen) 1991), and biology (Degond et al.) 2020). The interest the FPE lies in the fact that it represents an Fulerian view (i.e. location-based)) of the process rather than the usual Lagrangian view, based on trajectories (Batchelor) (200). However, except in very specific cases, analytical solutions to the FPE are out of reach for systems with high-dimensional or complex dynamics, necessitating efficient numerical methods.
Traditional numerical approaches, such as finite difference, finite element, and path integral methods, face the curse of dimensionality when applied to high-dinensional problems where the computational cost grows exponentially (Schuëller 1997). Since these methods rely on discretizing the solution space using a mesh, their scalability is inherently limitod, motivating the need for alternative, mesh-free approaches. Deep learning- based methods have emerged as promising alternatives in recent years, offering greater flexibility and scalability in high-dimensional settings (Raissi et al. 2019 Yin et al. 2023

---

Published in Transactions on Machine Learning Research (08/2025)
In this work, we introduce FlowKad novel generative model that combines the Feynman-Kad formula with normalizing flows to efficiently solve the Fokker-Planck equation (FPE). The key features and ad advantages of our approach are as follows:
Similar to Beck al. (2021b); Mandal & Apte 2024), method reformulates the FPE using the Feynman-Kac formula which expresses the solution an expectation over stochastic paths. This reformulation allows us to define an efficient regression loss function which is computed by sampling stochastic trajectories one of the main computational challenges of these methods. We significantly improve this process through the stochastic sampling trick, which leverages properties of stochastic flows and Taylor expansions to reduce the computational burden. Importantly, this sampling algorithm can be extended to other Feynman-Kace Caatmann-Ka-bbseee based approaches.
By leveraging normalizing flows, a class of flexible deep generative models, our approach is partic- ularly well-suited for handling probability densities, allowing to access a continuous representation of the solution. Once trained, FlowKac acts as a powerful sampler and provides density evaluations without requiring additional SDE discretization and enabling efficient generation of new samples from the learned probability distribution. This generative capability allows us, analogous to Gabriē ct al. (0022), to dynamically refine the training set by inverting the trained flow. By generating samples that increasingly align with the true probability density, this adaptive sampling strategy improves both training efficiency and the accuracy of the estimated density.
Unlike Mandal & Apte (2024). which relies on stationary densities infer time-dependent solutions, our approach remains effective even for stochastic differential equations (SDEs) with degenerate non-existent steady-state distributions, broadening its applicability to wider class of problems.
Finally, being a neural-based method, our approach can be scaled to higher dimensions.
More broadly, the Keynman-Kac framework extends beyond the Fokker-Planck equation, encompassing a wide class of parabolic partial differential equations (PDFs), including the heat equation, reaction-diffusion systems, and other Kolmogorov-type equations. This implies that our approach could be naturally extended to address a wide range of problems stochastic dynamical systems, statistical physics, and beyond, where such PDEs play a central role.
The remainder of this paper structured as follows: In Section
2 we review the relevant literature on deep learning-baset methods for solving partial differential equations, with focus on approaches for the Fokker- Planck equation. Section
g introduces the problem formulation and provides an overview of traditional non-ncural numerical methods. In Section 4 we present our proposed approach, detailing the reformulation of the FPE, the construction of the loss function, and the implementation of both the baseline and advanced versions of our method, which incorporates an efficient stochastic sampling algorithm. Section 5 contains series of numerical experiments comparing our method to a state-of-the-art neural approach, followed by coucluding remarks in Section D
2 Related work
The present section discusses the relevant. literature on deep learning approaches to solving partial differential equations (PDEs), with a particular focus on methods applicable to the Paknck-Planckk equation (FPE), and puts it in context with the proposed approach.
A broad class of neural methods for PDEs revolves around circumventing the curse of dimensionality and handling complex boundary conditions efficiently. Among these methods, the Deep Galerkin Method (DGM), introduced in Sirignano & Spiliopoulos (2018), extends the traditional Galerkin method by using deep neural networks to approximate solutions to a PDE as a linear combination of simpler functions. Similarly, the Deep Ritz Method (Yu et al., 2018) employs a variational formulation to approximate solutions for elliptic PDEs. Another approach is the Deep Splitting Method (Beck al. 2021a), designed for parabolic PDEs.
Code available httpp://gitthii http:/hhhb...moooooaaaa///ataaat
2

---

Published in Transactions on Machine Learning Research (08/2025)
This method subdivides the time domain into smaller intervals where the PDE is approximately linear, and deep-learning approximation is applied to each subinterval.
Backward Stochastic Differential Equations (BSDEs) have also been explored to solve PDEs (Han et al. 2017 2018 Raissi 2024p. The method employs coupled
SDEs, discretizing them to approximate the solution of the associated PDE. This method is particularly effective in dealing with high-dimensional PDEs, successfully scaling to PDEs in hundreds of dimensions.
One of the most prominent approaches in recent literature involves Physics- Informeed Neural Networks (PINNs), introduced in Raissi et al. 2019). PINNs solve PDEs by incorporating the governing equation into the loss function thereby minimizing the residual using a least-squares approach. PINNs have been applied to both Lime-dependent and steady-state PDEs. More recently, extensions of PINNs that leverage generative models, such as Normalizing Flows, have been explored (Feng & Zhou 2022 Tang et al. 2021). Such an extension uses generative models to adaptively sample collocation points, gradually enriching the training set and improving convergence while reducing the computational cost.
Reformulating the Fokker-Planck equation using the Feynman-Kac formula has proven to be powerful approach. Several works such as Beck et al. (2021b); Sabate Vidales et al. (2021) h adopt this strategy by discretizing the underlying stochastic differential equation (SDE) and solving a minimization problem on the expected value of the corresponding stochastic process. Another related strategy, described in Mandal & Apte (0224) exploits the steady-state solution of the FPE and the Feynman-Kat formula to derive a general time-dependent solution.
A recent paradigm explores solving the Fokker-Planck equation under the transportation of measure frame- work, which gives rise to probability flow equation (Boffi & Vanden-Eijnden 2023). In this approach, the initial probability density is pushed forward through a velocity field to yield the time-dependent density solution. This method links to recent advances in generative learning, such as score-based diffusion mod- cls (Song et al. 2021) and flow matching techniquos [Lipman et al. 2023), offering a novel perspective on modeling time-evolving probability densities.
For a broader overview of neural-based approximation for PDEs we refer the reader to the surveys Beck
ct al. 2023); Blechschmidt & Ernst (2021): E et al. (2021)
3
Problem statment
3.1
Fokker-Planck equation
We consider a filtered probability space (Q,F, P) and a time horizon T, and define a diffusion process X = (Xthejorr) by the Itô Stochastic Differential Equation (SDE):
dX, = =((,,, + (X,, f)dW,, t E
(1)
where W (We)ee the m-dimensional adapted standard Wiener process, μ x [0, Rdd the drift coofficient and a : Rd x [0, + d the diffusion matrix. Both functions are time-dependent and assumed to satisfy global Lipschitz conditions ensuring the existence and uniqueness of solutions to the SDE (Stroock & Varadhan I997). Moreover, these regularity conditions guarantee the existence and smoothness of the associated probability density function (PDF), which evolves according to the Fokker- Planck equation:
Op = =p== :
(up) + v (V. (V.TToT)))
(2)
Øt
where V- denotes the divergence operator. Specifically, when applied to a matrix, operates column- wise, -ossm,-wise, meaning the divergence is computed separately for each column of the matrix.
To explicitly represent Equation 2 let = [21, ra]T 1 [1i, 12, ........ ma]" and D too T where D is the diffusion matrix, thus the FPE becomes:
0p (z. 1) = -- d 0t
82
Iuta, ()p(2, +)]
Σ
Dis(2, ()]
x E ERR [0,
Oz,
dz, Ox,Oz,
1=1
1
3

---

Published in Transactions on Machine Learning Research (08/2025)
and the initial condition is specified as:
- (()) x ZER
The solution is subject to the normalization and non-negativity constraints:
Jec p(a, t)dx =1 and p(x,t) 2 0, At E [0,
(3)
and the additional boundedness condition:
p(a,t) - as ||0||2 + co,
(4)
ensuring is valid probability density.
Over the years, significant work has been done to develop numerical solutions for the Pokker-Planck equation, spanning classical discretization schemes to modern neural network- based approaches. The following sections provide a concise overview these approximation techniques.
3.2 Non-learning- based integration methods
Analytical solutions for the Fokker Planck equation have been developed for only a limited number of low-dimensional systems. This limitation has led to a large body approximation theory and tech- niques (Schuěller) 1997 Hundsdorfer & Verwer 2003), each leveraging different perspectives address the intractability of the exact solutions. These include:
1. Classical density estimation techniques: These methods approximate solutions through random sampling, estimating the probability density function empirically using techniques such as histograms or kernel density estimation (KDE).
approximates the prob- 2. Path Integral (PI) Approach: The PI method (Naess & Johnsen 1991
ability density function as Gaussian over an infinitesimal time interval, which is then integrated to obtain a global solution.
3. Maximum Entropy Methods (MEM): In the MEM framework (obczyk & Trebicki 1990:
Sobczyk & Trebicki, 1992). : the solution is approximated through the various moments and con- strained to maximize the informational Shannon entropy.
4. Grid-Based Numerical Solutions: Grid-based techniques discretize the FPE and solve it through numerical integration and linear algebra Pichler et al. 2013). These methods span a spectrum of formulations:
Finite Differences (FD): FD leads to an explicit scheme where the values can be calculated directly, eliminating the need for matrix inversion. Despite its simplicity, explicit finite difference methods often face stability issues, necessitating implicit formulations for more reliable solutions.
Alternating Direction Implicit (ADI): ADI is another significant and improved finite dif- ference scheme. In ADI, finite difference steps in each direction are resolved separately, treating one dimension implicitly and the others explicitly in each step. This approach results in sta- ble finite difference formulation. The primary advantages of ADT include tridiagonal matrix, which allows for efficient computation.
Finite Element Methods (FEM): FEM partitions the domain into smaller subdomains, allowing flexible handling of complex geometries and boundary conditions.
Each numerical method offers unique advantages in terms of stability, efficiency, and applicability to complex geometries. However, the main drawback that all these methods suffer from the curse of dimensionality arising when handling high dimensional FPE. As these methods rely mostly on discretized grid, the computational complexity scales exponentially with the grid dimension, while the grid precision heavily impacts the approximation accuracy.
4

---

Published in Transactions on Machine Learning Research (08/2025)
3.3 Neural network solutions: The Physics-informer approach
A notable paradigm that addresses the dimensionality constraint is machine learning, which offers mesh- free solutions. A particularly promising method for solving differential equations is the Physics-Infow Neural Network (PINN) approach. This tochnique integrates knowledge from physical laws defined by partial differential equations and boundary conditions. In the PINN framework, a neural network, paramcterized by set of weights and computing po, employed to approximate p", the true solution to the Fokker-Planck equation (FPE). The neural network achieves this approximation by minimizing combination of two loss functions: one related to the dynamics of the Fokker-Planch operator (Cx), where the derivatives of the neural-based probability density Pe are efficiently computed through automatic different iation, and another enforcing the initial condition (4,). The training set consists of Ny spatial points and My temporal points, sampled uniformly from a given domain. The total loss function is then defined as:
[()) - (0) + Cz
(5)
where
N₂ Mr
Opo (x* 18) - FDo(at (8)) )²
Cx( = No x M
and
1 N (, = NF (po(=*,0) 0) )x²1)) 2
PINNs have demonstrated remarkable efficacy in modeling and approximating solutions to a wide range of differential equations, as comprehensively reviewed in Cuomo al. (2022). However, despite their widespread success, recent work by Mandal & Apte (2024) reveals a fundamental limitation of the PINN framework when applied to solving the Fokker-Planck equation and deriving time-dependent solution. Specifically, if we assume that Equation 2 has a strong, unique solution p", one can show that there exists sequence of functions (fn meN that minimize the loss however the sequence fails to converge towards the true solution p' (Mandal & Apte. 2024 Proposition 4.1):
lim 12+00
1
Ns Mr
05n at (1
FSn Fon( - +
1 NS
Σ
NS
In (2 .0) 10 (x+t) (20)) = (6)
N₂ x Mj
yet.
lim fn + p" n-000
In other words, despite successfully minimizing the loss, the sequence (Ja)ner does not guarantee convergence to the accurate time-dependent solution of the Fokker-Planck equation, emphasizing a significant limitation of the PINN approach. This theoretical insight serves as basis for the development of the FlowKac which aims to address these limitations and provide more robust framework for solving time-dependent Fokker- Planck equations.
4 FlowKac
In this section, we introduce FlowKac, our approach for solving the Fokker-Planck equation (FPE) by leveraging Iwo key components: the Feynman-Kac formula and normalizing flows. Each of these elements addresses critical challenges forming a powerful framework for efficient and accurate solutions.
The Feynman- Kac formula establishes fundamental connection between stochastic processes and partial differential equations (PDEs), enabling express the solutions of the FPE as expectations over stochastic paths.
5

---

Published in Transactions on Machine Learning Research (08/2025)
Complementing this, we employ the flexible architecture of a temporal normalizing flow (Both & Kusters 2019) that is highly effective for modeling complex probability distributions and offers the advantage of efficient sampling. This capability allows us to compute pointwise values of the FPE solution at arbitrary spatial and temporal points, eliminating the need for fixed computational grids. Consequently, FlowKac enables efficient. training and accurate evaluation of the ground-truth probability density at new points of interest. By combining the Feynman-Kac formula to derive the unique solution with the expressive power of temporal normalizing flows to approximate it, FlowKac provides a robust, time-dependent, and mesh-free framework for solving the FPE. Furthermore, FlowKac addresses the limitations of numerical approximation methods (e.g., FEM, ADI) and neural network-based approaches like PINNs.
4.1 Temporal Normalizing flow
The temporal normalizing flow (Both & Kusters, 2019) is a flow-based model designed to approximate and model intricate time-dependent distributions. Let X E Rd be random variable and (x) its corresponding probability density function that we aim to model. A normalizing flow (NF) (Papamakarios et al. 2021
Kobyzev et al.) 2021; Grathwohl et al. 2019 is constructed via a simple base distribution pz(2) and a diffcrentiable bijective mapping f - Z such that X == f- (Z). This transformation allows for exact density estimation through the change of variable formula:
logpx(2)) logpz (f(z) log det vxf(x))
(7)
To address the dynamics of time-dependent distributions, the temporal normalizing flow (TNF) extends
this framework to the temporal domain by introducing a time-angmented variable X = (X,t) and its corresponding latent representation Z = (z,t'. The TNF then models the evolution of X through a
time-dependent transformation:
log (xx) log (f(=")) + log det V=-
where Vn- f denotes the Jacobian matrix of the time-augmented mapping [(, A critical aspect to highlight is the dependence between the latent and observable variables, expressed by the relationships 2 = 2(x, t) and t' = =1((x,) t). Moreover, the TNF enforces the invariance of the time variable, t' = t. This constraint arises because the bijective transformation that maps the simple latent density into a more complex one is applied exclusively to the spatial variable x, ensuring that the resulting function retains the essential normalization property of a probability density. Time merely serves as an index for the transformation. As a result, the Jacobian matrix simplifies to:
8z 8z
dz
V V
= V,
ot
St
0
The TNF f is constructed through a sequential chaining of multiple transformations:
= f(xt) = Ty) o T[r- o... o (x, and x" (x,t) = Tiu) Tpi c... Tin) (z,
Each of these transformations adheres to the KRnet architecture (Tang et al. 2021 which is detailed in Appendix The architecture comprises two fundamental layers: an Actnorm layer that scales the input and adds a bias, followed by an affine coupling layer adapted from real NVP (Dinh et al. (017) and responsible for mixing the spatial and temporal variables. KRnet adds a nonlinear transformation at the end of this soquence, helping to capture more complex and challenging dynamics, which constitutes an improvement over RealNVP. KRnet also offers key advantages over ODE-based normalizing flows, such as Continuous Nor- malizing Flows (CNFs) (Chen et al. 2018). In CNFs, transformations are governed by neural ODEs, where the mixing of spatial and temporal variables occurs implicitly during ODE integration. In contrast, KRnet explicitly controls the mixing of spatial and temporal variables through its discrete, structured transforma- tions, allowing for faster training while maintaining flexibility in modeling complex probability distributions. The Jacobian matrix of this sequential flow is derived via the chain rule:
dde V.fl = det V Va- Tal
6

---

Published in Transactions on Machine Learning Research (08/2025)
here 244 is the output of each transformation, with 210] and 2|r]
4.2 Training algorithm
Our purpose is to train the TNF solve the time-dependent Pokker-Plauee equation accurately. As discussed in Section 3.3 the Physics-Informed approach can suffer from convergence issues, thus simply minimizing the PINN loss given by Equation 5 is insufficient to guarantee accurate solution. Instead, following work in Beck et al. (2021b) we leverage the Feynman-Ko Kac formula (Oksendal, 2013 Theorem 8.2.1) to guide the
training process:
Theorem 4.1 (The Ray formula) Let (Xt) be an Ito process with drift it and diffusion ø. Let E C(R)) and v E C (R" x R) satisfy for all (r, E R x [0,
du der
1
+
(o0")
qu
(9)
0t.
2 (.j
dr.zy
with initial condition v(z,0) 0) = f(x) such that f E (R()). Then, the unique solution can expressed as:
u( 1) = E [exp (-/ a(K8,, / (Xi) Xo
Such a formula provides a powerful connection between parabolic partial differential equations (such as the FPE) and expectations of stochastic processes.
This representation allows to define more effective loss criterion for the training process based on path- wise expectations. To apply the Feynman-Kac Kac formula, first rewrite the Fokker-Plance equation 2 into form matching equation g isolating a term that corresponds to the function (sce Appendix B for derivation and more details):
0p 01
d 0Di dzj
Op Ori i==
0pp 0x,0x; = (V V (V-
44 + 2
Σ
Σ
Dis
(10)
j=1
We identify the scalar function g V 1n-V- (V D) the effective potential term in the transformed
PDE. Then, we apply the Feynman-Kac formula to the transformed equation 10 to obtain solution of the form:
(x,
cxp
(8),s) ds) vp (X.) X.
(11)
here (Xx)120, referred to as the FlowKac process, represents a stochastic process characterized by the drift and diffusion terms:
0
#4 +
Ox
(12)
a.
Finally, to optimize the parameters of the TNF, we propose the following loss function:
[()) = //'' [Ipo(..t) PFK(, pFk(,,t)lddt = f/' I (
(Ppt))pFKK(_ppp (ok(x,t)) dx dt
(13)
This regression loss function measures the discrepancy between the TNF's predicted density pe and the true solution derived via the Feynman-Kac formula prk, with the I.2 norm inducing strong form of convergence
7

---

Published in Transactions on Machine Learning Research (08/2025)
which also implies the convergence in distribution of the underlying stochastic process. Thus, the TNF is guided toward the correct solution by minimizing this loss.
However, directly discretizing the integral given by Equation 13 is computationally prohibitive for high- dimensional systems, as it requires covering a d-dimensional mesh, causing the computational complexity to scale exponentially.
To alleviate this constraint, we use a sampling- baase approximation that bypasses the mesh-based evaluation. Specifically, we uniformly sample n, points from spatial domain [a, bjd sufficiently large to encompass the support of the distribution for all t, and n, temporal points from []. This reduces the computational cost by replacing the uniform discretization of the integral with discrete sum over sampled points:
((0)
(pa(=",t) (PPpk(( e))²
(14)
The k=1j=1
We further enhance this approach by leveraging the generalive properties of normalizing flows. Instead of relying solely uniform sampling, we use an adaptive sampling mechanism that incorporates learned density information (Gabrié et al. 2022). Specifically, after an initial training phase, we refine the training set by inverting the normalizing flow, thereby generating n, new spatial samples that better reflect the underlying probability density. This adaptivity concentrates computational effort on regions of higher probability mass, improving both training efficiency and density estimation accuracy. Quantitative comparisons between uniform and adaptive sampling are provided in Appendix G
A critical advantage of this approach lies in the properties of the Feynman-Kan formula, which enables the computation of pointwisc solutions to the Fokker-Planck equation without relying on adjacent mesh points. By eliminating mesh dependencies, this sampling stralegy significantly mitigates the curse of dimensionality while maintaining high fidelity in the approximation of the loss function.
The proposed training procedure is detailed in Algorithm I we sample multiple stochastic paths starting from each training point, compute the Reynman-Kac-basee densily estimate, and evaluate the corresponding loss. The process is repeated across all training samples.
Algorithm 1: FlowKac (naive)
Input Maximum epochs Ne, number of sample paths nw, number of spatial points ne; mumber of temporal points T4
for 1 -= 1,..,, Ne do Sample uniformly Ctrain = (z+ ())) ())1skksn Initialize loss Le = 0; for h: = 1,,,..,, n2 do Sample nw paths of X, at time points (P)issn,, starting from x²;
4
5
6
Compute PFK (* 12) using empirical mean;
Compute TNF's output pe p= t));
Accumulate loss: L. = L. + E, (po(a*, tD) - (r 8))²
8
end for
10
Update model parameters 0 using the Adam optimizer;
11 end for
Output: The predicted solution po(x, t)
The major computational challenge in the training process lies in the repeated application of the Feynman- Kac formula for each training sample. Since sampling the stochastic process (X,) conditioned on the initial point z², a new set of stochastic paths must be generated for every individual data point, increasing computational complexity. This challenge is inherent all methods based on Feynman- Kac reformulation and path sampling, where the necessity of repeatedly simulating stochastic trajectories leads to increase
8

---

Published in Transactions on Machine Learning Research (08/2025)
in computational cost, particularly in high-dimensional settings. As result, this bottleneck highlights the critical need for more computationally efficient approaches.
In the following section, we introduce an enhanced version of the algorithm that siguificantly accelerates the training process by addressing this inefficiency.
4.3 Stochastic sampling trick
In this section, we introduce a novel and efficient algorithm for sampling the Flow Kac process by leveraging the properties of stochastic flows (Kunita Kunita 1990). The SDE governing (Xt))ejo,r) is given by:
dx, aa t)dt + t)d,,
(15)
We denote by (2,((z)= = x the solution to the SDE 15 starting from at time s. 4s defines stochastic flow.
An important property of stochastic flows is stated by the following theorem (Protter
2005
Theorem 40):
Theorem 4.2 Let drift js and diffuusion coefficients have locally Lipschitz derivatives up to order k. Then, for any fixed realization w E Ωt the mapping (,,((..w) IRd Rd is le times continouously differentiable.
This theorem establishes that the solution to Equation 15 is smooth with respect the initial condition, thus allowing to employ high-order Taylor expansions (Cartan 1967 Theorem 5.6.1 and 5.6.3) to the stochastic flow and expressing solutions to the SDE starting from arbitrary initial points xo + h.
leveraging a solution starting from a reference point Zroft
1 2o,c(x) (2,r( + 2b (Fret))) + 2 (z + + (k)) 2 (E))t) (I111*)..
1
(16)
2
here, (1) (Zrett represents the linear term involving the Jacobian matrix Ja while "i (zret) h) = hT ThTHa (Fret)) denotes the quadratic term, involving the Hessian tensor Hæ Higher- order terms provide additional precision for approximating the stochastic flow but are typically truncated for computational efficiency.
For dynamics that depend linearly on the initial condition, the stochastic sampling trick provides exact solutions, regardless of the choice of x and Trefs as illustrated in the 1-dimensional geometric Brownian motion (GBM) example in Section 5.1 In these cascs, all higher-order terms in the Taylor expansion vanish since they are independent of the initial condition Fref, resulting in an exact reconstruction of the solution.
This expansion forms the cornerstone of our novel sampling algorithm as we can efficiently generate samples of X, starting from perturbed initial conditions while maintaining high accuracy. This approach significantly reduces the computational burden of sampling the FlowKac SDE for each data point as summarized Algorithm 2
However, for dynamics with nonlinear dependencies with respect to the initial condition, maintaining curacy requires careful selection of the reference point 2zef in the Taylor expansion. To preserve the local validity of the expansion and avoid significant approximation errors, we propose a variant where Fret is dy- namically chosen as the centroid of each mini-batch. The resulting modified training algorithm is detailed in Appendix K as shown in Algorithm|
5 Numerical results
To assess the performance of our proposed model, we conducted a series of experiments across a range of stochastic differential equations. We begin with a simple 1-dimensional example validate the foundational aspects of our approach. Following this, we apply our method more complex processes then nonlinear SDFs, demonstrating its robustness in capturing intricate dynamics. Finally, extended the framework
9

---

Published in Transactions on Machine Learning Research (08/2025)
Algorithm 2: FlowKac (stochastic sampling trick)
Input Maximum epochs Ne, number sample paths nw, number of spatial points nay number of temporal points Re; reference point for Taylor expansion Iref
for 1 - 1, N. do Sample uniformly Ctrain = 11=) (ISSSs 18jsmz
Sample nw Brownian motion paths W;: // Fix realization across all training points as required by Theorem 4.2
4
Compute Jacobian (o.t(Fret) and Hessian Ho using automatic differentiation; Divide Ctrain into m batchos {C"mm Initialize loss L. = 0;
6
for b= 1,..., m do
Compute Taylor expansion for batch Chh starting from Zref (Equation 16 Compute PFK:
8
Compute batch loss Lbi Accumulate loss: Le Le + Lbi
11
12
end for
Update model parameters 0 using the Adam optimizer;
14 end for
Output: The predicted solution po(x,t) t)
to tackle higher- dinhee--dmennionnll stochastic processes, showcasing model's scalability and effectiveness in handling increasingly complex systems. Our model is compared to the PINN-TNF approach of Feng & Zhou (2022). The comparison is based on both qualitative and quantitative assessments. Qualitatively, we give, for each SDE, probability density and marginal distribution plots to visualize the fidelity of the learned solutions. Quantitatively, following Feng Zhou (2022) and [Tang et al. (2021), we compute metrics based on the relative L distance and the Kullback-Leible (KL) divergence between the ground truth probability density p" and the model's predicted density pe. All quantitative results then sunmarized in dedicated discussion subsection 5.6
(t) = IIp Po(.,t)11 t1,
= / (z, (2,t))oo log (x
p' t) dx.
To ensure accurate computation, both metrics are evaluated on a spatial grid, leveraging Henat evaluation points uniformly spaced across the support of p" (.,t) ensuring thorough coverage of the probability density for SDEs up four dimensions. For higeer-dimensional examples, where uniform grid becomes computa- tionally infeasible, instead draw fixed number of sample points directly from the ground-truth density to compute the metrics efficiently:
LT(()= Eneee (p" t) po(x*.1)) novdl pa (x (x²,t)
Devul
p" (z
DkL(t)
Σ (**, log 1oml
Teeal
po po
For all numorical experiments, we employ the Adam optimizer (Kingma & Ba 2014 with a learning rate of 1 = 0.001 and default momentum parameters to optimize the normalizing flow parameters. The stochastic differential equations are numerically sampled using the torchsde framework of Kidger et al. (2021): Li
ct al. 2020), which implements adaptive-stop solvers based on the Suler-Maruyann and higher-order strong schemes. Finally, for the stochastic sampling trick, the Jacobian and Hessian terms are computed efficiently using the automatic differentiation capabilities of pytorch (see details in Appendix
We also include Appendix J dedicated comparison between FlowKac and an MLP trained directly on the Feynman Kac loss, highlighting the benefits of incorporating generative model and isolating them from those of the Feynman-Kan formulation.
10

---

Published in Transactions on Machine Learning Research (08/2025)
5.1 Univariate Geometric Brownian Motion
We first apply our framework to the univariate Geometric Brownian Motion (GBM). This process serves as an ideal test case due to the existence of known analytical closed-form solution, providing rigorous benchmark for our methodology. The following SDE characterizes the GBM dynamics:
dX, = JXxt oX,WW,
where represents the constant drift coefficient, denotes the constant volatility, and W, is standard one-dimensional Wiener process. The corresponding FPE is described as follows:
dp -- pz + 200x)
Op 1 8.p
222220 dr²
x2
(17)
2
We choose the following initial condition:
1og(z)²
(xx) n22
exp
20
The exact solution for the process is given by:
1 1 &+oW.)
X; Xo exp
with the associated probability density function:
(log(x) -n- 202) 1)2
1 D" (x,,t) 21/2x(t + 1)o
exp
(18)
2(t + 1)02
We solve the FPE on the interval (0, using a training set of size |Crrail = 4.10* generated through a uniform distribution, and nu = 500 sample paths. The model is trained for Ne = 200 epochs, and for our normalizing flow, we use a sequence of depth L = 8. Notably, we employ our stochastic sampling trick by computing first-order Taylor expansion to efficiently sample the resulting FlowKac process (Fquation 12
dx, (- + 20 X + oX,dW,
(19)
To compute (o,, the stochastic flow solution to Equation 19 oriignatting from x at t = O, first derive the Jacobian with respect to the reference point (see details in Appendix E
J (Tree exp ((-++>) 30
1 oW,
In this example, first-order Taylor expansion turns out to be sufficient to completely capture the dynamics, all higher-order terms vanish due to their independence from the reference point, and yields:
X = (2,((x) = oo. + (x Ja.t (xFot)))))
(20)
cxx ((u++
302)
2 t BoWW,
The stochastic sampling trick represents significant speed-up our computational approach, allowing us to compute sample paths of Equation 19 with perfect accuracy while significantly reducing the computational cost. Table Sooleeprreeents comprchensive performance comparison between the standard FlowKac implementa- tion and the optimized version incorporating the stochastic sampling technique. The results reveal consistent and substantial speedups across different. training set sizes and sampling configurations. Additional runtime comparisons for other SDEs are provided in Appendix I
In practical implementations, our experiments were conducted using 300 to 500 stochastic paths per training point, which proved sufficient for producing accurate density estimations. Notably, as the dinensionality
11

---

Published in Transactions on Machine Learning Research (08/2025)
Configuration
Run Time (seconds)
Speedup
Training Size
Wi Samples
FlowKac (naive)
FlowKac (Sampling trick)
Factor
((Crrain))
(nw)
x 10
500
31.1
3.9
8x
x 104
1,000
57.7
5.2
11x
104
500
93.1
6.6
14x
x 10a
1000
171.5
10.5
16x
Table 1: Performance comparison between naive FlowKac and FlowKac with stochastic sampling trick, for single training epoch using fixed batch size of 2000. The experiments were conducted on an NVIDIA GPU with GB memory. The speedup factor demonstrates increasing efficiency gains with larger training sels, highlighting the technique's scalability benefits.
of the SDE increases, the required training size will also grow accordingly to maintain accuracy. In this context, the benefits of the stochastic sampling trick become even more pronounced, mitigating the otherwise prohibitive computational burden associated with high-dimensionall sampling.
The results, displayed in Figure 2, show the comparison between the predicted density and the exact solution from Equation 18 This comparison demonstrates strong alignments between the predicted density and the exact solution, validating the accuracy of our framework for this univariate case.
Expet
0.6
FlowKac P)
FowKae Pi
0.7
FlowKus Ps
0.
0.6
0.5
0.4
0.4
0.2
0.2
0.1
0.1
0.
0.0
0.0
(t=0)
(t=0.5)
(t=1)
Figure Comparison between FlowKac density po and the true density of the univariate GBM process, = 0.5 and é = 1.
5.2
Multivariate
We next explore the dynamics of a two-dimensional Ceattein-Uhlebeckk (OU) process governed by The following SDE:
dx,, = AX,di + EdW,
(21)
0.1 1
X,dt + (0.6 dw,
(22)
-1
-0.1
and the initial condition is given as:
(1) 1
.
(23)
12

---

Published in Transactions on Machine Learning Research (08/2025)
The unique strong solution to Equation 21 can be expressed as follows Gobet & She 2016
Proposition
1):
X, ee Xo + S --s
(24)
its mean my and covariance matrix 122 are given by:
m, [| [[xol
(25)
1 = E X, eAl (E[X,X)) I c-AAE e- e-dio ddss) ca
(26)
Therefore, the process (Xt) follows a multivariate normal distribution, and the exact solution of the corre- sponding FPE is:
(.,t)=((m,, Vi).
We solve the FPE on the interval [-5,5]] 5]2 using a training set of size |Ctrainl = 4.10 generated through the uniform distribution. The model trained for Ne = 250 epochs. For the normalizing flow, we use a sequence of depth =
The results prosented in Figure show the comparison of the predicted 2-dimensional density at different time points (L -0, 1, 2, and 3) for the FlowKac model, the PINN model, and the exact solution derived from Equation 24| The FlowKac model demonstrates a strong capability to accurately replicate, across time, the bohavior the process. In contrast, the PINN model shows a significant deviation from the true density, particularly at later time points.
Quantitative results, as summarized in Table 2 and further emphasized in Table 3 highlight a significant performance gap between the two approaches. The PINN model demonstrates a progressive deterioration in accuracy proven by the increasing error metrics over time (L2 error increasing from 5.7 x 10-3 1.11) reflecting divergent behavior. In contrast, FlowKac maintains stable and lower error values throughout the simulation.
This reinforces the robustness and precision of our framework in modeling complex stochastic processes.
5.3 Multivariate Geometric Brownian Motion
The third process is a two-dinensional extension of the Geometric Brownian motion (Bareera et al. 2022 and is described by the following SDE:
1 _ x,dt BX,dW, dx, A+
(27)
2
where A, 2x2 such that the cigenvalues of (A+ B²) have negative real part, and W2 is univariate Brownian motion. Under these conditions, Equation 27 admits closed-form solution:
= (IA + BW,) No
(28)
here A= = (o) 000 (0) B = Ch
0
:) 0)= = 00.5 (0: 9) and the initial state Xo is distributed
02 according to the multivariate log- normal distribution:
140 (no- = (0.5) Eg 2 12
00.5 0.7
+ = Log-N
The exact solution this instance can be expressed a closed- flosed-form formula (details Appendix D
p² (-,t) Log-A (-=
00.5 /0.5+a0
+ ail
E, (0.5
(..5 + (b1)2t bib2t
0.5 + (b2) )
b1
(29)
0.7 + azt
13

---

Published in Transactions on Machine Learning Research (08/2025)
brad
Bextip
Deackp
32
82
34
PlonCac
Ponfacoy Py
Fowkaa Ps
HuwKax
12
33
14
11
11
66
PM
PVW P.
PWV
PN
31
31
3
2
2
(t=0)
(t-1)
(t=3)
Figure Comparison density distributions for the 2D Ornstein- Uhlenbeck process at 0, t t= 2, and = 3. The exact solution is depicted in the top row, FlowKac predictions in the middle row, and PINN predictions in the bottom row.
this density converges to stationary Dirac distribution centered at as t -
To solve the associated FPE on the interval (0, 6]2, we employ training set of size |Ctrain| = 6.10 generated through the uniform distribution. The model is trained over N, = 300 epochs, using a normalizing flow of depth L = 14. An ablation study analyzing the impact of varying the flow's depth, layer width, and learning rate on performance presented in Appendix H
Results depicting the two-dimensional densitics are displayed in Figure The comparison shows a strong alignment between the true solution of the Fokker-Planck equation and the prediction generated by FlowKac. As the density evolves towards the stationary state consisting of Dirac distribution centered at 0, the support of the distribution becomes progressively concentrated. Such a configuration presents significant challenges for mesh-based approximation methods, which struggle to capture such sharply localized distri- bution. Additionally, alternative Feynman-Kac reformulations relying on the steady-state solution would be ineffective for this SDE due to the singular behavior of the Dirac distribution.
Quantitative results in Table 2 reinforce the qualitative observations. Specifically, FlowKac demonstrates lower 1,2 errors and DxL. values when compared to the PINN approach, further validating our model.
14

---

Published in Transactions on Machine Learning Research (08/2025)
Eeaclp
track
Evadd
40
0.3
0.-
35
84
14
16
63
Rewwbr PV
Finddor
Ficedart Ps
6.0
0.
15
4
2
1
13
0.5-
23-
23
33
1.
(t=0)
(t=0.5)
(t=1)
Figure 4: Comparison of 2D density distributions for the Multivariate GBM at t = 0, t t= 0.5, and f 1. The exact solution is depicted in the top row, FlowKae predictions the middle row, and PINN predictions in the bottom row.
5.4 Duffing oscillator
To further assess the ability of our model to capture complex chaotic nonlinear dynamics, we solve the FPE for the 2-dimensional Duffing oscillator (Pichler et al. 2013) described by:
X2e (X1x) = (-D.Aee + w² X11 0.1w²X1/ dX22
dt + (o V0.8 dW.
(30)
with initial condition vb =N N (N((()) 1/2
Since no closed-form solution exists for this system, the ground truth approximated numerically by solving the FPE using the ADI schemo, following the approach in Feng & Zhou| (2022). To ensure robustness and accuracy, the computation is performed on a finely discretized mesh over the domain [-10, 10j² (implemen- tation details provided in Appendix Ch. This higl-resolution numerical solution serves reliable reference, enabling rigorous evaluation of FlowKac model's output.
For FlowKac, we train the model to solve the FPE associated with Equation 30 for w = 1, on the same interval [-10, 10j2. The training dataset consists of |Ctrainl = 3 x 10a sample drawn uniformly across the domain. The model is trained over N₂ = 300 epochs, using a normalizing flow architecture of depth and by employing the nonlinear layer with 60 bins (details in Appendix AA.. Notably, for this example, we simulate the FlowKac SDE directly rather than utilizing the stochastic sampling trick. This decision
10
15

---

Published in Transactions on Machine Learning Research (08/2025)
stems from the need for high-order derivatives (higher than second-order) in the Taylor expansion for the sampling trick to work. These high-order derivatives are computationally intensive to estimate using sampling approximations.
Figure D compares the solutions obtained with FlowKac, PINNs, and the ground truth computed using the ADI scheme. We again observe strong alignment between the predicted densities by FlowKac and the exact solutions. Additionally, the quantitative results in Table 2 further highlight FlowKac's superior accuracy over the PINN approach, demonstrated by consistently lower L and Dkr. values across time.
110
Exasto
Dast
Euddp
15
02
32
-13
41
-35
Howiaa Pe
TowKac
NowRac
Powhlo
310
150
-33
-113
-104
-15
-5
09
36
15
20
-10
-006
-75
50
75
103
-108
-00
-5
-3
84
25
15
-000
-30
-15
-44
-25
66
25
59
75
16.
FBN
12
12
-20
13
-33
-313
+15
-25
26
+75
15
-102
+15
.336
-5
44
25
144
(t=0)
(t=0.5)
(t=0.75)
(t=1)
Figure Comparison density distributions for Duffing non-linear oscillator at 1 0, t 0.5, 0.75, and == 1. The exact solution, given by ADI, is depicted in the top row, FlowKac predictions in the middle row, and PINN predictions in the bottom row.
5.5 Higher dimensional examples
Finally, to assess the scalability our approach, we consider a relatively higher-dimension Ornstein- Uhlenbeck process:
dx, AX,t + EdW,
(31)
where aTa, Σ ola and W. a d-dimensional Brownian motion. The model's structure allows for closed-form solutions and exact marginal distributions in all dimensions, making it particularly well-suited for high-dimensional analysis. The associated FPE is given by:
Op = - Øt
Op Dx;
S2p
+ p)
(32)
2
Or,Ox;
i=
i-j=
16

---

Published in Transactions on Machine Learning Research (08/2025)
For the initial condition, assume multivariate normal distribution:
1
1
vb = mg
Vo
4 l7
1
The exact solution is given by:
p(. - (0m00mm=-- e0 (m:=-=mo m V2 = e + o²
(e - -1]ld 2a
(33)
Figure 0 presents comparative analysis between FlowKac and PINN models for the 4-dimensional sctting, demonstrating the superior performance of our model, which consistently achieves lower DKL and T,2 errors across various time values.
As dimensionality increases, computing the PINN loss becomes increasingly tedious. Consequently, we report performance metrics exclusively for FlowKac up to d = 12, while noting that our model runs out-of-memory on an A100 GPU with GB of memory at d 100. To further illustrate the accuracy of FlowKac in higher dimensions, Figure compares the learned 2D marginal distribution the marginal dimensions randomly with the ground truth at various time steps for the 8-dimensional setting of the OU process. The learned marginal distribution is obtained by sampling points from the temporal normalizing flow, and the results indicate strong agreement with the exact distribution, reinforcing the scalability of our approach and its ability to handle high-dimensional stochastic processes elficiently. Additional evaluation metrics across various dimensions are provided in Appendix
chosen
(d:4)
0.40
(d=4)
0.08
-0- PINN (smd)
0.35
PNNN (chn)
kd. divergence
0.04
0.20
0.15
0.00
0.0
0.2
0.6
0.8
0.8
1.0
0.0
0.2
0.4
0.5
0.8
1.0
Figure 6: Comparison of DEL and relative 2 errors between FlowKac and PINN models for 4-dimensional Ornstein- Uhlenbeck various time values.
5.6 Discussion
This section provides a quantitative evaluation of FlowKac's performance compared the PINN framework, using KI-divergence (DKI.) and relative I,2 errors metrics. While Table 2 overview across different SDEs and time points in t [0, Table 3 specifically focuses the 2D Ornstein-Uhlenbeed process over an extended time horizon to assess long-term stability. Our analysis reveals several key findings:
For low-dimensional processes (e.g., 1D GBM), both approaches achieve strong performance, with PINN showing marginally better accuracy (DRl. = 1.30x vs 2.60 10-3). However, the practical significance of this difference is minimal given that the error levels are well within acceptable bounds.
In higher-dimensional settings (2D GBM and OU), FlowKac maintains consistent convergence prop- erties and stable error metrics, whereas PINNs exhibit growing instability over time, particularly in the 2D OU process, where errors increase exponentially (from DKL(O) 10-1 to D3I(3) 1.68).
FlowKac consistently preserves error values within the 1 to 10-3 range across all time points, demonstrating superior stability compared to PINNs.
17

---

Published in Transactions on Machine Learning Research (08/2025)
0-
0
Exact
Exact
Exact
2-
1-
21
(t=0)
(t=0.5)
(t=1)
Figure Comparison of 2D marginal distributions between FlowKac and the ground truth at f 0, 0.5 and
The PINN's observed behavior aligns with the divergence phenomenon presented by Mandal & Apte 2024). This limitation manifests even when the training loss given by Equation appears decrease significantly. In contrast, FlowKac's consistent performance across all test cases demonstrates its robustness as general-purpose solver for stochastic processes, outperforming PINNs in challenging scenarios.
Finally, it is worth noting that although the Feynman Kac representation permits direct pointwise estimation of the solution, such evaluations require simulating many stochastic paths. Once trained, FlowKac provides density evaluations and samples at negligible marginal cost, thereby amortizing the initial training cost over large number of queries. A brief runtime comparison in Appendix L quantifies this efficiency advantage.
SDE
Model
Dka())
= 0.0
0.25
0.5
== 0.75
== 1.0
= 0.0
0.25
0.5
== 0.75
==
FlowKae
2.60e-3
1.90e3
1.80e-3
2.00e-3
2.60e-3
6.28e-2
3.56e-2
3.82e-2
1.61e-2
5.93e-2
GBM ld
PINN
3.00e-4
9.00e-4
1.50e-3
1.60e-3
1.30e-3
5.20e-3
1.94e-2
2.53e-2
3.08e-2
3.44e-2
Omstein-Uhleebeek 2d
FlowKac PINN
1.60e-2
1.36e-2
1.22e-2
1.13e-2
8.90e-3
6.07e:2
7.79e-2
5.73e-2
4.88c-2
4.70e-2
1.00-4
1.65e-2
3.25e-2
4.946-2
5.70c-3
1.136-1
1.58e-1
2.10o-1
FlowKae PINN
1.98-22
1.306-2
1.92e-2
3.24e-2
5.89e-2
1.24e-1
9.26e-2
9.67e-2
1.08e-1
1.38e-1
GBM 2d
5.10e-3
1.7ee-2
1.88e-2
9.60e-2
1.65e-1
4.18e-2
8.25e-2
1.45e-1
2.05e-1
2.59e-1
FlowKac
1.15e-2
9.80e-3
9.70e-3
7.20c-3
5.12e-2
4.52e-2
4.33e-2
4.21e-2
5.64e-2
9.90e-2
Duffing Oscillator
PINN
5.70e-3
1.446:2
3.61e:2
7.30c22
1.136.1
2.16e-2
5.09e
9.82e-2
1.29e
1.67c.1
Comparison of KL divergence and L enror metrics between FlowKac and PINN models across different SDEs al various time points. Bold values indicate the beller-performing model for each metric and time point.
Table 2:
6 Conclusion
Neural-based approaches to solving the Fokker-Planck equation have shown remarkable results, particularly when dealing with high-dimensional systems where classical numerical methods suffer from the curse of
18

---

Published in Transactions on Machine Learning Research (08/2025)
Model
Metric
t == 0
==
t= 2
t= 3
DRL.
1.60e-2
8.90e-3
5.20e-3
3.19e-2
FlowKac
L
6.07c-2
4.70c-2
4.74c-2
1.37c-1
DKL
1.00e-4
4.94e-2
2.94e-1
1.68
PINN
LA
5.70e-3
2.10e-1
5.18e-1
1.11
Table 3: Performance comparison of FlowKac and PINN models the 2d Ornstein-Uhlenbeee process using KT. divergence (DKL) and L2 error metrics.
dimensionality. The widely used Physic-informed approach can encounter convergence challenges as demon- strated through the 2-dimensional Ornstein-U Uhlenbeck process.
To address these limitations, we developed FlowKac, a novel framework combining the Foynman-Kac formula and normalizing flows which enables the use of more robust loss metric improving the stability and accuracy during training. We also developed stochastic sampling trick that exploits the smoothness of stochastic flows, significantly improving sampling efficiency and reducing computational complexity without sacrificing precision.
Numerical experiments confirm FlowKac's effectiveness, demonstrating strong qualitative and quantitative alignment with true solutions. While higher-order Taylor terms enhance accuracy, they also introduce computational overhead, which must be carefully managed for scalability.
Overall, FlowKae provides robust and scalable solution for solving the FPE, with promising applications in physics, finance, and other domains requiring accurate modeling of complex stochastic systems.
Acknowledgments
We sinccrely thank Xiaodong Feng, Li Zeng, and Tao Zhou for providing the code for temporal normalizing flows used in this paper. We also thank Pinak Mandal for helpful discussions. We thank the anonymous reviewers for their thorough evaluation and insightful recommendations, which helped improve the presentation and clarity of this work.
References
Gerardo Barrera, M Högele, and J Pardo. Cutoff stability of multivariate geometric brownian motion. arXiv preprint arXio:2207.01 01666, 2022.
George Keith Batchelor. An introduction to fluid dynamics. Cambridge university press, 2000.
Christian Beck, Sebastian Bocker, Patrick Cheridito, Armulf Jentzen, and Ariel Neufeld. Deep splitting method for parabolic pdes. SIAM Journal on Scientific Computing, 43(5) 2021a.
Christian Beck, Scbastian Bocker, Philipp Grohs, Nor Jaafari, and Armulf Jentzen. Solving the kolmogorov pde by means of deep learning. Journal of Scientific Computing, 88:1 28, 2021b.
Christian Beck, Martin Hutzenthaler, Arnulf Jentzen, and Benno Kuckuck. An overview on deep learning- based approximation methods for partial differential equations. Discrele and Continuous Dynamical Sys- tems B, (666))369773744555 2023.
Jan Blechschmidt and Oliver G Ernst. Three ways to solve partial differential equations with neural net- works a review. GA M.itM-Mitteilungn,, :0222000000 2021.
Nicholas M Bolfi and Eric Vanden-Eijnden. Probability flow solution of the fokker planck equation. Machine Learning: Science and Technology, 4(3):035012, jul 2023.
Gert-Jan Both and Remy Kusters. Temporal normalizing flows. Xiu preprint ar arx 2019.
19

---

Published in Transactions on Machine Learning Research (08/2025)
H. Cartan. Calcul Différentiel. Hermann, Paris, 1967.
Ricky T. Q Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, NTPS'18, pp- 6572 6583, Red Hook, NY, USA, 2018. Curran Associates Inc.
Salvatore Cuomo, Vincenzo Schiano Di Cola, Fabio Giampaolo, Gianluigi Rozza, Maziar Raissi, and Francesco Piccialli. Scientific machine learning through physics-informed neural networks: Where we are and what's next. Journal of Scientific Computing, 92(3):88, Jul 2022. ISSN 1573-7691.
Biswa Nath Datta. Numerical linear algebra and applications. SIAM, 2010.
Pierre Degond, Maxime Herda, and Sepideh Mirrahimi. A fokker-planck approach the study of robustness in gene expression. Mathematical Biosciences and Engineering, 17(6):6459-6486 2020.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In International Conference Learning Representations, 2017.
Weinan E, Jiequn Han, and Arnulf Jentzen. Algorithms for solving high dimensional pdes: from nonlinear monte carlo to machine learning. Nonlinearity, (5)::2 dec 2021.
Li Feng, XiaodongZZeg and Tao Zhou. Solving time dependent fokker-planck equations via temporal nor- malizing flow. Communications in Computational Physics, 32(2):401-423, 2022. ISSN 1991-7120.
Marylou Gabrié, Grant M. Rotskoff, and Eric Vanden-Eijnden Adaptive monte carlo augmented with normalizing flows. Proceedings of the National Academy of Sciences, 119(10):2119922011, 2022.
Emmanuel Gobet and Qihao She. Perturbation of Ornstein-UUhleebbeck statiee stationary distributions: expansion and simulation. preprint, July 2016. URL httpss//haal.cccencc//aaa1111333
Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Πla Sutskever, and David Duvenaud. FFJORD:
free-form continuous dynamics for scalable reversible generative models. In International Conference on Learning Representations, 2019.
Jiequn Han, Arnulf Jentzen, et al. Deep learning-based numerical methods for high-dimensional parabolic partial differential equations and backward stochastic differential equations. Commanications in mathe- matics and statistics, 5(4):349 0804499--00000 2017.
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
Zheyuan Hu, Khemraj Shukla, George Em Karniadakis, and Kenji Kawaguchi. Tackling the curse of dimen- sionality with mnysics-informed neural networks. Neural Networks, pp. 106369, 2024.
W. Hundsdorfer and J. Verwer. Numerical Solution of Time- Dependen aduection--iifussioon-eeeee Reacttion Equa- tions. Springer Berlin Heidelberg, 01 2003. ISBN 778-3-644-05777-6
Patrick Kidger, James Foster, Xucchen Li, Harald Oberhauser, and Terry Lyons. Neural SDEs as Infinite- Dimensional GANs. International Conference on Machine Learning, 2021.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR, abs/ Abs/1112.6698 2014.
Ivan Kobyzev, Simon J. D. Prince, and Marcus A. Brubaker. Normalizing flows: An introduction and review of current methods. IEEE Trans. Paltern Anol. Mach. Intell, 43, 2021.
Hiroshi Kunita and Hiroshi Kunita. Stochastic flows and stochastic differential equations, volume 24. Cam- bridge university press, 1990.
HP Langtangen. A general numerical solution method for Fokker-Plant equations with applications structural reliability. Probabilistic engineering mechanies, 6(1):33 48, 1991.
20

---

Published in Transactions on Machine Learning Research (08/2025)
Xuechen Li, Ting-Kam Leonard Wong, Ricky T. Q. Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. International Conference on Artificial Intelligence and Statislics, 2020.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference Learning Representations, 2023.
Umberto Lucia and Gianpiero Gervino. Pokker-plande equation and thermodynamic system analysis. En-
tropy, 17(2):763-77) 771, 2015.
Pinak Mandal and Amit Apte. Solving fokker baaker-planck equations using the zeros of fokker-planch operators and the feynman-kac formula. arXiv preprint 3rXiv:2401.011 2024.
A. Naess and J. M. Johnsen. The Path Integral Solution Technique Applied to the Random Vibration of Hysteretic Systems, pp. 279 291. Springer Netherlands, Dordrecht, 1991. ISBN 978-94-011-3692-11
Bernt Oksendal. Stochastic differential equations: an introduction with applications. Springer Science & Business Media, 2013.
Manolis Papadrakakis, George Stefanou, and Vissarion Papadopoulos. Computational methods in stochastic dynamics, volume 1. Springer, 2011.
George Papamakarios, Eric T. Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshmi- narayanan. Normalizing flows for probabilistic modcling and inference. J. Mach. Learn. Res., 22, 2021.
Lukas Pichler, Arif Masud, and Lawrence Bergman. Numerical solution of the fokker-planch equation by finite difference and finite element methods a comparative study. In Compulational methods in stochastic dynamics, volume 2. Springer, 2013.
P.E. Protter. Stochastic Integration and Differential Equations. Stochastic Modelling and Applied Probabil- ity. Springer Berlin Heidelberg, 2005. ISBN 97835400031377.
Maziar Raissi. Forward-backward stochastic neural networks: deep learning of high-dimensionn partial differential equations. Tn Peter Carr Gedenkschrift: Research Advances in Mathematical Finance, pp- 637 655. World Scientific, 2024.
Maziar Raissi, Paris Perdikaris, and George Karniadakis. Physics-informee neural networks: deep learn- ing framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-7017 707, 2019.
H. Risken and T. Frank. The Planck Equation: Methods of Solution and Applications. Springer Series Synergetics. Springer Berlin Heidelberg, 1996. ISBN 9783540615309.
Mare Sabate Vidales, David Siška, and Lukasz Szpruch. Unbiased deep solvers for linear parametric pdes. Applied Mathematical Finance, 28(4):299-3) 329, 2021.
G.I. Schuēller. A shate-of-the-art report computational stochastic mechanics. Probubilistic Engineering Mechanics, 12(4):197-8)) -211, 1997.
Justin Sirignano and Konstantinos Spiliopoulos. Dgm: A deep learning algorithm for solving partial differ- ential equations. Journal of compulational physics, 775:1339-1364, 2018.
K. Sobezyk and J. Trebicki. Analysis of stochastic systems via maximum entropy principle. In Nicola Bellomo and Fabio Casciati (eds.), Nonlinear Stochastic Mechanics, pp. 485-497, Berlin, Heidelberg, 1992. Springer Berlin Heidelberg. ISBN 978-3-642-84789-9.
K. Sobezyk and J. Trebicki. Maximum entropy principle in stochastic dynamics. Probubilistic Engineering Mechanics, 5(3):102-110, 1990. ISSN 0266-8920.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poolc. basre-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
21

---

Published in Transactions on Machine Learning Research (08/2025)
Didier Sornette. Fokker-planee planckk equation of distributions of financial returns and power laws. Physica A:
Stalistical Mechanics and its Applications, 290(1-2):211-2111 2001.
Daniel W Stroock and SR Srinivasa Varadhan. Multiddimensional diffusion processes, volume 233. Springer
Science & Business Media, 1997.
Keju Tang, Xiaoliang Wan, and Qifeng Liao. Adaptive deep density approximation for fokker-plancee equa- tions. J. Comput. Phys., 457:11 2021.
Surya T. Tokdar and Robert E. Kass. Importance sampling: a review. WIREs Computational Statistics, 2
(1):54 60, 2010. doi: https///doi.orr/0.1100000 wic..6..
Yuan Yin, Matthieu Kirchmeyer, Jean- Yves Franceschi, Alain Rakotomamoniy, and patrick gallinari. Con- timoous PDE dynamics forecasting with implicit neural representations. In The Eleventh International Conference on Learning Representations, 2023.
Bing Yu et al. The deep ritz method: deep learning-based numerical algorithm for solving variational problems. Communications in Mathematics and Statistics, 6(1):1- 12,111,,22 2018.
A KRnet
The KRnet architecture Tang et al. (2021) comprises two fundamental layers: an Actnorm layer that scales the input and adds bias, followed by an affine coupling layer responsible for mixing the spatial and temporal variables. The Actnorm layer applies the following transformation:
6 = @ I[(] + b;
Subsequently, this output is transformed through the affine coupling layer. Let Øis = [.]]11 be a partition composed of m-uplet T[H.1 R and d - m-uplet and 24J,2 IRd-m, the transformation is expressed as:
*1+11111 = x[1111 2|+11,,2 = 2 2 (1 tanh (si+1 (z)r,,, + e @ tanh (rit1 (21)11))
where 0 < a < 1 is hyperparameter while 3 is trainable parameter. The functions 8:41 and ri+1 are implemented using neural networks providing the necessary flexibility to capture complex, time-dependent relationships.
To further enhance the nonlinear characteristics of the temporal NF, an additional layer is incorporated and is based on a cumulative distribution function. Consider a probability density function p(z) and a partition of the interval [0, 1] given by 0 ho < b < ... < h We construct p(z) through a piece-wise linear
interpolation:
p(x) = Wi+1 Wi (xx h:) + Wi. Ve E h441
h;+1 h4
the corresponding cumulative distribution function F is then defined
F(x)
Wi+1 - (x (x)) w; (x h4) 2(h++1 h:)
We + wk+1 (hi+1 - h(), Va E (hihhsi)
k=U
This layer refines the model's ability capture and represent complex distributional properties, significantly augmenting the expressiveness and accuracy of the TNF.
22

---

Published in Transactions on Machine Learning Research (08/2025)
B Fokker-Planck equation transformation
In this appendix, we provide detailed derivation of the transformed Fokker-Planck equation. This trans- formation facilitates the use of the Feynman-Kac formula to derive a solution. The Fokker-Plance equation describing the time evolution of a probability density function p(a, 1) for a d-dimensional stochastic process with drift vector μ = [un,,.., Pd] and diffusion matrix D 2200 = [Day] given by:
a.p Σ
am. |uipl +
Σ s,==1
dz. [s,, (Dup)).
Expanding the derivative terms in the equation to separate contributions from p. μ4, and D, we have:
d
0pp= =
Oz M4
Σ
U4
Σ
Sz. (p Dig + Dig P2,p)
i=1
i=1
i-=
Oz, U4
Σ
1.O.pp p+
(0 22, Dis + + p0:n, Dis 0z, Dis D2,p+ Dis Dwanp)
i-=1
i M4 +p
O Oacn Dij
Σ U; p + Σ DijO
OziZ, P P+ Σ (0,, Dig + pu, Dij)
i,)=
Introducing notations for divergence and corresponding second-order operator for compactness, we rewrite the equation:
0pp p(- (V Σ udr Σ
Dis DisOr,a,Pt P
Σ
(o,, Dis + de, Dis)
1
1J=1
The symmetry property of the diffusion matrix Dij Dji further simplifies the last sum
d
d
Oo. Dig + p2, Di)) = 2 Σ
Dx.p& D2+2 +2
Σ
dz O2, Dis
1=1
d
2 Σ
dr 0z Dis +
dupP Σ
Dij
jæti=l
2
0z, p
Oz, Σ j=1
Dis
Combining everything, we obtain reformulated Fokker-Planne equation with reorganized drift and diffusion contributions:
0.p
+
Σ
u4 2 0z Dig 0xp Σ
DijOz,a,P -V. D)) p- 0.
i=1
1,==
C Alternating Direction Implicit scheme
We employ the Alternating Direction Implicit (ADI) Papadrakakis et al. 2001 method to solve the 2- dimensional Duffing oscillator. The computational grid is defined with a time step h tm+1 - tm and a uniform spatial step 8 = I1.i+1 - 21,4 = 22,411 -2.5. The governing PDE is given by:
0p -2 On
8p
Op
8p
(w2r) 0.4wmz o.1 Or + 0.4wp 0.4
(34)
0t
Ozb
23

---

Published in Transactions on Machine Learning Research (08/2025)
The ADI scheme splits the numerical solution into two sequential half-steps, alternating between the spatial dimensions 21 and 22. First, we discretize along the x1-axis using the finite difference method, yielding an intermediate half-step solution Pa.j m+ denoted as P.j. This intermediate solution is then used in the second half-step, where we discretize along the Fz-axis obtain the final time-step solution. The first half-step equation takes the form:
P Pi.
22
Pi Pi+1.
1.j
P;-
(wP 0.4w12, - 0.1w² Pt.i1 P
0.4p
26
28
P 2p
t
Pi.j-1
+ 0.4
82
Rearranging the equation, obtain:
Pi+1. P.-1. 1,5 h -1 0.4H)p (w~x1 0.Awaz2 x1,
P
P
Pig + F2
25
28
+ 0.4 Pijt1 -2p"" + Pi.j h
This equation takes the form of a tridiagonal system:
a.P5-1.5 + brtl, + capit1s di,
where
hw2.1
0; 28
b; 1, C; -
and
d. (1 0.4h))? (wP - 0.4waz2 (0.1)xx P
P
Pig+1
2p +
0.4
25
82
We solve this system using the Tridiagonal matrix algorithm (Thomas algorithm) the intermediate solution p".
Datta
2010 to obtain
Afterwards, we apply an explicit finite difference method along the x2-axis. time-step pd ml given by:
The equation for the next
m1 Pij
pi.
Pi+1, - Pi-1.j 102,,,5 0.4wx22 0.1w (
m+1 m+1 Pij+1
Pi.
F2
0.4p
26
26
p m+1
2ph
m+1 + Pi,j-1
pitt1
0.4
82
This equation also reformulated into a tridiagonal system:
m+1
+ bipl c.p.tt) = d.,
where
0.4h
0.4h
0; = -w-2r1, - 0.Awaa,,--000000 - 2x
6; = (w~x1,5 1 - 0.Awwa, -0.1
0.1
25
82
25
82
0.8h
b, =
1 +
82
and
Pi+1. p 1. h.
== = (1 0.4h)pi, - 22.1
20
The Thomas algorithm is again employed to solve for m
24

---

Published in Transactions on Machine Learning Research (08/2025)
D
Multivariate GBM fosed-form solution
The multivariate GBM process governed by the SDE:
(A+ (A+ = X,dt BX,dW,
dX,
(35)
b1 0 0
where A
B
are diagonal matrices enabling a closed-form solution. The initial
by condition Xg follows a log-norma distribution:
0
[401 (402/
Ouu 0
V Log-N
Eo
002
The exact solution the 2D GBM SDE is given by:
X; exp (tA + BW,) Xo
(x1.0 ex w.) 12.0 x2200epp exp(a2t b2
(log X1,e) log X2,1
Defining Z. log X,
we obtain Gaussian vector with:
Hoi ait)
14
Juz azt)
00i + bitt
bib2t
Σ,
by byt
oiz bzt
Finally, the probability density function X, is:
Ex² (log m.)).
p(z,t)=
exp
(
- (log na) T
21 / (2e)
25

---

Published in Transactions on Machine Learning Research (08/2025)
E
Path-wise Jacobian and Hessian computation through Automatic Differentiation
Algorithm 3: Sampling SDE Solutions and Computing Sensitivities
Input: SDE with drift and diffusion d, time interval [2), number of sample paths Reamples: state space dimension dstate; number of time points Atime, reference point Tref
Output: Sample paths Xg IR X Peamplls' detate, Jacobian Je(rr))) E MR Neampinn X Xchttato X thetate
Hessian (Tre)) E IR Risim X Teamples Xdatate Xdatate Xdstate
Initialize: Fix Brownian motion sample paths (torchsde. (
Define function for SDE sampling:
expand Fref match the batch size Dsamplesi
Compute SDE sample paths X, starting from Trefi
Compute Jacobian:
Use automatic differentiation (torch. func. func.jacrev) to compute the Jacobian matrix w.r.t the initial condition 2zefi
0x
Jacobian
Dxre
Compute Hessian:
Define function that returns the Jacobian for a given input Fref: Apply second-order automatic differentiation to compute the Hessian lensor;
10
11
OJacobian
8 8x
Hessian
Dxret
Ox
12
Return: Sample paths XTrr, Jacobian (x(rrot), and Hessian He He(veet);
F
High-dimensional SDE: quantitative metrics
6.1
aheergmnce
Reeetw wmer
80
Figure DEL and relative L errors FlowKac across different dimensional settings (4, 6, 8, 12) and over various time values for the Orustein-Uhlenbeed process.
Figure 8 presents the relative I error and Kullback-Leible divergence (DxL. across different dimensions 4, 6.8. 8, 12) for the Ornstein- Uhlenbeck process. These results are computed over various time steps and demonstrate the robustness of approach in accurately modeling high-dimensional stochastic dynamics.
26

---

Published in Transactions on Machine Learning Research (08/2025)
G
Adaptive sampling
We investigate the rolc of adaptive sampling, where training points are drawn from the learned density by inverting the normalizing flow, on both training convergence speed and final accuracy. Using the 4- dimensional Ornstein- Ornstein-Uhlenbecck Uhlenbeck process as a benchmark, we compare training with uniform sampling versus adaptive sampling by tracking the evolution the relative L error and Kullback Leibler divergence (DKL).
Uniforma
1.8
Uniform Adaptive
Adapiiiee
1.6
1.0
1.0
0.5
0.4
0.0
Epoch
Epoch
Figure 9: DRL and relative L errors for FlowKac trained with uniform vs. adaptive sampling the 4D Orn- stein Uhlenbeck process. Adaplive sampling leads to accelerated convergence and improved accuracy.
Although samples are drawn from the learned flow during training, we do not incorporate an explicit im- portance sampling correction in the loss (Tokdar & Kass 2010). While a theoretically principled objective would require reweighting by the flow's probability density account for the change of measure, find that directly optimizing the empirical loss in equation 4.2 yields more stable convergence and consistently lower errors across experiments.
H Ablation study
To better assess the relative influence of the normalizing flow archilecture the model's performance, conduct an ablation study using the 2-dimensional Geometric Brownian Motion (GBM) process described in subsection 5.3. We vary key architectural and optimization hyperparameters. namely the number of flow layers L, the width W of the neural network layers, and the learning rate. The base configuration uses neurons per layer, and a learning rate of 10-3, We report the mean KT. divergence and relative L errors across time steps E {0, 0.25, 0.5, 0.75, 1}, along with the final training loss, and the best achieved loss to assess convergence behavior.
I. 14 layers, W
Table 4: Ablation results 2D GBM model. Default configuration: 14, W 32, 1r 1e-3. Metrics are averaged across E (0,0.2 0.25, 0.5, 0.75,
Configuration
KL Mean +
1 Mean +
Final loss +
Best Loss |
Base (L
14, 32, 1r 10-3)
4.3le-2
1.16e-1
7.32e-5
4.93e-5
(Lr 10 10-4)
1.37c-1
1.67c-1
1.85c-4
1.78c-4
(Lr -
6.0le-2
1.20e-1
1.88e-4
6.66e-5
(L = 10)
4.32e-2
1.25e-1
9.23e-5
6.86e-5
(L 18)
2.75e-2
9.54e-2
4.26e-5
3.6le-5
(W 16)
6.05e-2
1.28c-1
7.36e-5
7.36c-5
(W 48)
4.42c-2
1.16e-1
5.45e-5
5.45e-5
We observe that deeper (L 18) and wider (W 48) architectures yield modest improvements across all metrics. However, these gains are incremental and come at the cost of increased computation and potential overfitting. In contrast, using fewer layers smaller widths leads to measurable drop in performance,
27

---

Published in Transactions on Machine Learning Research (08/2025)
but offers better trade-offs in terms of stability and computational efficiency. For this reason, retain the parsimonious base configuration, which achieves competitive performance across all evaluation metrics while maintaining a manageable training cost.
Interestingly, the learning rate shows a more pronounced effect: both higher and lower learning rates degrade performance, either due to optimization instability or insufficient convergence. This highlights the sensitivity of density estimation in generative SDE modeling learning dynamics.
Runtime comparison
To quantify the practical impact of the stochastic sampling trick, we report a runtime comparison per training epoch between the naive version of FlowKac and the variant leveraging the stochastic sampling trick. The table below summarizes results across various SDEs and dataset sizes. All experiments were using a batch size of 2000 and 500 Brownian sample paths per point. The results confirm that the stochastic sampling trick offers significant speedups, particularly as the training size and model complexity increase.
Table 5: Runtime comparison (in seconds) between FlowKac (naive) and FlowKac with Stochastic sampling trick, for single training epoch, using fixed batch size of 2000 and 500 Brownian sample paths. Experiments were conducted on an NVIDIA '4 GPU (16 GB). The speedup factor highlights increasing efficiency for larger training sets and more complex SDEs.
Configuration
FlowKac runtime (s)
Speedup
Training Size
Process
Naive
Stochastic Trick
Factor
2 x 10a
GBM 1d
31
8x
6 x 10a
GBM ld
93
14x
2 x 10a
OU 2d
148
5
32x
6 x 104
OU 2d
431
7
61x
2 104
GBM 2d
125
6
20x
6 10
GBM 2d
374
8
47x
J Feynman-Kac based MLP comparison
In this appendix, compare proposed approach-Floom FlowKac, which combines the Feynman-Kac formula with temporal normalizing flow and the stochastic sampling trick-with baseline model that directly ap- proximates the Feynman-Kaa and leverages a standard MTP. Both models are evaluated on the 2-dimensional Ornstein Uhlenbeck process, and we present both quantitative and qualitative results.
Quantitative Results. Table reports quantitative error metrics KL divergence and relative L errors various time points t 2; While the Feynman- Kac MLP model performs reasonably well at earlier time points, its accuracy deteriorates time. In contrast, FlowKac maintains lower errors across all time steps, demonstrating improved stability and fidelity to the target solution.
Table 6: Quantitative comparison between FlowKac and an MLP baseline on the 2D Ornstein-Ullenbed process.
Time t DKL (FlowKac)
DKL (MLP)
L? (FlowKac)
La (MLP)
0.0
1.60e-2
3.38o-2
6.076-2
9.53e-2
1.0
8.90e-3
3.93e-2
4.70e-2
9.20e-2
2.0
5.20e-3
2.86e-2
4.74e-2
9.72c-2
3.0
3.19e-2
1.1le-1
1.37e-1
1.91e-1
28

---

Published in Transactions on Machine Learning Research (08/2025)
Qualitative Results. Figure 10 presents density plots at t = 2 and t comparing the outputs of FlowKac, the Feynman-Kac ML.P and the ground truth. While the MLP provides good approximation, its outputs do not satisfy the properties of valid probability density function such as non-negativity and proper normalization. In contrast, FlowKac, being generative model based on normalizing flows, guarantees both by construction.
Furthermore, ullike FlowKac, the MLP does not offer a natural mechanism for sampling, nor can it support adaptive sampling strategies that exploit the learned density during training. This limits its flexibility and effectiveness.
Exact
FlowKac
MLP
22
02
3-
2-
92
51
63
31
2
-3
20
0.0
(t. 2)
Flowaac
3-
51
62
51
11
-2
63
10
-1
3
10
(t = 3)
Figure 10: Comparison of exact, FlowKac, and MLP densities at times and 3.
K Alternative dynamic reference point algorithm
As an alternative to using a fixed Taylor expansion reference point I we can dynamically compute it per batch during training. This variant of the FlowKac algorithm enhances accuracy and locality in the Taylor approximation by adjusting the expansion point to match the current batch. The detailed procedure is described in Algorithm
L Amortized runtime: Kaynman-Kac and TNF
We compare the computational cost evaluating Neval 10 solution queries, corresponding to 5000 spatial points along 20 temporal points, using two approaches:
direct application of the Feynman Kac formula via stochastic sampling, and
amortized inference with trained Temporal Normalixing Flow (TNF).
29

---

Published in Transactions on Machine Learning Research (08/2025)
Algorithm 4: FlowKac (dynamic Zret)
Input Maximum epochs Ne, number sample paths nw, number of spatial points nly number of temporal points n4
for 1 = 1, N. do Sample uniformly Ctrain (x*, ( (P)ssknn,.. 1SiSmm
Sample nw Brownian motion paths W
// Fix realization across all points as
required by Theorem 4.2
Divide Crrain into m batches {C)m Jb=1
Initialize loss L. = 0;
for b = 1,..., m do
Compute Zref = centroid(C*);
Compute Jacobian (aro()) and Hessian Hæ, (zrot) using automatic differentiation;
Compute Taylor expansion for batch O starting from Fref (Equation 16
10
Compute PFK:
11
Compute batch loss Lp: Accumulate loss: Le Le + Lbi
12
13
end for
14
Update model parameters using the Adam optimizer;
15
end for
Output: The predicted solution po(x,t) t)
The results in Table 7 show that the Feynman Kac incurs important computational cost. In contrast, once trained, the TNF provides essentially instantaneous evaluations illustrating the amortization advantage of neural inference.
Table Runtime (in seconds) evaluate neval 10 solution queries for benchmark processes. The Feyn- man Kac approach requires sampling at each query, while the TNF offers amortized fast inference.
Process
Kac (s)
TNF (s)
OU 2d
50.42
0.01
GBM 2d
30.32
0.01
30


Reviews:
[{"review_id": "Cm2faBO3Se", "reviewer": "Reviewer_pZmn", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strength:**\n1) Bypasses expensive grid discretization; complexity grows only with the number of sampled points, not grid resolution.\n\n2) The normalizing flow’s invertibility enables targeted sampling in high-density regions, improving both sample efficiency and final accuracy.\n\n3) Leverages the Feynman–Kac formula for an exact characterization of the solution, rather than approximating PDE residuals.\n\n4) Achieves consistent accuracy improvements( relative l2 error and KL divergence) over PINNs across multiple test SDEs.\n\n**Weakness:**\n1) The sampling trick relies on a local expansion around a reference point; for highly nonlinear or multimodal flows, the accuracy of truncated expansions may degrade unless reference points are re-selected continually.\n\n2) Computing higher-order derivatives via automatic differentiation can be memory- and compute-intensive, especially in high dimensions or with deep flow architectures.\n\n3) Comparison is only against PINN-TNF methods; evaluation versus other Feynman–Kac-based neural schemes or classical mesh-free solvers (e.g., FBSDE) would further contextualize gains.\n\n4) Although the method scales beyond 4D, experiments report OOM failures around 12–100 dimensions.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "**Summary:** FlowKac introduces a mesh‐free neural solver for the Fokker–Planck equation by recasting it via the Feynman–Kac formula and training a time-conditioned normalizing flow to approximate the evolving density; it leverages stochastic trajectory simulations and an adaptive, generative sampling strategy—accelerated by a local Taylor–expansion trick for Jacobians—to achieve high accuracy and speedups over PINN baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please consider addressing my weakness section.\n\n**Minor Comments:** The Codebase link is provided, but the file content is unable to load.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "WisYk5k5pg", "reviewer": "Reviewer_31sN", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n- The paper is well written. The related studies and the motivation of the work could be easily followed.\n- The method is technically reasonable. Fitting networks directly to the pointwise evaluation of the target solution (because it's available via the F-K formula) makes sense.\n- The proposed efficient computation for the F-K formula seems to be effective.\n\n### Weaknesses\n\n(1)\nThe authors point out that the limitation of the vanilla PINNs for F-P eqs is that a zero PINN loss does not imply the solution is true, referring to Mandal & Apte (2024). I am curious if any empirical results speak of this limitation. Does the inferior performance of the baseline method in Section 5 stem from it? If not, what are the potential causes instead? It seems an important point of discussion because the limitation is the claimed motivation to use the F-K formula for loss definition. I am not sure if the current results support such a claim, or the claim is inherently difficult to support, and thus the superiority of the proposed method should be attributed to other factors too.\n\n(2) In Section 4.2, the authors mention the \"dynamic selection of $x_0$,\" but in Algorithm 2, $x_0$ is simply given as one of the inputs. What is the way to dynamically select $x_0$?\n\nBelow are relatively minor points.\n\n(3)\nIn Theorem 4.1, $q$ seems to appear out of nowhere without any description. It is revealed afterward in the middle of the page, but defining it clearly within the theorem will make it more readable.\n\n(4)\nAs the pointwise evaluation of the solution is possible by the Feynman-Kac formula, my naive thought is that in practice we could just evaluate it using the proposed efficient computation method, without fitting any networks. I know the point of learning a neural net is amortization, but the need for amortization depends on the computation speed of the non-amortized approach. So I would suggest providing information on the computation time evaluating the solution on $n_\\text{eval}$ points, both by the proposed method (neural net) and the mere computation of the F-K formula.\n\n(5) A really minor thing; consider using \\eqref or \\cref commands to refer to equations.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "A method for obtaining solutions of the Fokker-Planck equation is proposed. The method fits a (time-dependent) normalizing flow to the solution directly evaluated pointwise using the Feynman-Kac formula. The paper's particular contribution lies mainly in 1) the combination of the use of two notions (normalizing flows and the Feynman-Kac formula), due to which the convergence issue of PINNs for F-P eqs is solved, and 2) the efficient method to compute the Feynman-Kac formula.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please consider adding descriptions regarding points (1) and (2) in the Weaknesses section.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "kjhyaoGTdt", "reviewer": "Reviewer_9H29", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n* The paper is well written. And the mathematical description is adequate.\n\n* The authors present the algorithm for both the \"naive\" version as well as the improved version (\"stochastic sampling trick\"), which helps in the understanding of the key contribution.\n\n* The improved method achieves between 8x and 16x speedup (on at least one example) over the \"naive\" sampling method.\n\nWeaknesses:\n\n* Since the main contribution of the method is the \"stochastic sampling trick\", I think, the authors should evaluate the benefit of this method for all the examples. Currently, the only evidence is provided in Table 1 for the univariate Brownian motion example.\n\n* Time interval for all examples is taken as [0,1]. Long term behavior not shown. Will the error be bounded in the long term?\n\nQuestions:\n\n* The method is a collocation based method. Are there any restriction from the stability point of view? If yes, how does this method compare against PINN in that regard?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents a neural method for solving the Fokker-Planck equation. The neural model used is a normalizing flow augmented with the time dimension, implemented with the KRnet architecture. The ground truth is computed using numerical evaluation of the probability density using the Feynman-Kac formula. The authors also employ a sampling method where new samples are generated in the neighborhood of an existing sample, assuming sufficient regularity. This reduces the sampling complexity of the algorithm. The authors test their method on various stochastic differential equations, and compare their method with physical informed neural network (PINN) solutions.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* I think the discussion in Sec 4.3 can be improved, and contrasted better wrt. Sec 4.2. For example, the implication of Theorem 4.2 should be explained better with suitable visual representations. And this must be related with the algorithm notations. For example, Algorithm 2 uses batches whereas Algorithm 1 does not.\n\n* Please provide some examples with long term behavior.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "jKh0AoDUj2", "reviewer": "Reviewer_Y4Qv", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "*Strengths.*\n\nThe paper is mostly clearly written and is well-organized with very few errors.\nThe method is an interesting combination of techniques for stable solution of Fokker-Planck equations.\nThe paper evaluates the approach on a number of examples, comparing against PINN models.\n\n*Weaknesses:*\n\nWhat is the significance of the temporal normalizing flow in contrast to a vanilla normalizing flow? What’s the significance of the statement regarding the comparison with neural ODEs and continuous normalizing flows? What is the advantage of maintaining the time invariance?\n\nThe modeling power of the normalizing flow depends on the dimension of the input, since each layer has to be a bijection. With low to medium dimensional data this seems like a severe restriction. So even, if one has a complex distribution in, say 3d space, the modeling power of the normalizing flow would be restricted since each layer would have limited width.\n\nWhat happens if the normalizing flow distribution does not match the true distribution well? Does the adaptive sampling still work or does it correct itself? How much does it depend on a good initial training phase?\n\nAblation: What is the relative influence of the normalizing flow in the experiments? I don’t think that the authors have made this comparison.\n\nThe stochastic sampling trick seems to have some limitations as to applicability. Since it requires the knowledge of the flow at time 0. In the cases where the authors apply this, the closed form solution is always known. The only case where it is not known (the Duffing oscillator), the trick is not applied.\n\nCan you explain the initial training phase? I don’t think this is explained in the paper.\n\n\n*Typos*\nTitle: Feynman Kac-Formula -> Feynman-Kac Formula", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes FlowKac, a method for solving the Fokker-Planck PDE for high dimensional complex dynamical systems. It proposes a generative model that combines the Feynman-Kac formula with temporal normalizing flows using physics-informed networks (PINNs).\n\nPINNs, however, have a limitation where minimizing the PINN loss can lead to solutions that do not converge in the case of time-dependent solutions. This limitation is fixed by Feynman-Kac.\n\n\nSimilar to prior work it reformulates the Fokker-Planck equation using Feynman-Kac Feynman-Kac formula which relates PDEs to stochastic processes.\nIt expresses the solution as an expectation over stochastic paths and sampling such stochastic trajectories. The formula allows a regression loss with the advantage that it allows pointwise evaluation without relying on adjacent points. However, the formula has to be computed for each training point which leads to high computational complexity.\n\nThe paper proposes a more efficient way of applying the formula by using the ‘stochastic sampling trick’ with stochastic flows when the flow is k-times differentiable. In that case the flow can be approximated as a Taylor series given the Jacobian and Hessian, and the stochastic paths can be estimated.\n\nIt also proposes using temporal normalizing flows as neural sampler and density estimators to adaptively refine the training set so that the procedure concentrates on higher density samples.\n\nAs far as I can tell the new contributions of the paper are the stochastic sampling trick and the temporal normalizing flow applied to Fokker-Planck.\n\nFinally the proposed algorithm is evaluated on one and two dimensional systems (univariate and 2d geometric brownian motion, 2d Ornstein-Uhlenbeck process, 2d Duffing oscillator) and higher-order OU process with up to 12 dimensions showing better convergence in the higher dimensional settings.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please see the weaknesses above. I think ablation experiments for the normalizing flow component are quite important, since that is a central contribution. Furthermore, please address the questions regarding the flexibility of normalizing flows given the relatively low dimensions.\n\nPlease also address the limitations of stochastic sampling, explaining whether the stochastic sampling trick can be used in the cases where a closed form is not known and whether this is a strong limitation.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)