Deep Kinetic JKO schemes for Vlasov-Fokker-Planck Equations

Abstract

We introduce a deep neural network-based numerical method for solving kinetic Fokker Planck equations, including both linear and nonlinear cases. Building upon the conservative dissipative structure of Vlasov-type equations, we formulate a class of generalized minimizing movement schemes as iterative constrained minimization problems: the conservative part determines the constraint set, while the dissipative part defines the objective functional. This leads to an analog of the classical Jordan-Kinderlehrer-Otto (JKO) scheme for Wasserstein gradient flows, and we refer to it as the kinetic JKO scheme. To compute each step of the kinetic JKO iteration, we introduce a particle-based approximation in which the velocity field is parameterized by deep neural networks. The resulting algorithm can be interpreted as a kinetic-oriented neural differential equation that enables the representation of high-dimensional kinetic dynamics while preserving the essential variational and structural properties of the underlying PDE. We validate the method with extensive numerical experiments and demonstrate that the proposed kinetic JKO-neural ODE framework is effective for high-dimensional numerical simulations.

Figures (11)

• Figure 1: Drift error vs. time step size $\Delta t$ for Example 1 in the 1D (left) and 3D (right) cases. The plots demonstrate linear convergence behavior as the time step decreases. The dashed lines indicate a reference $\mathcal{O}(\Delta t)$ convergence rate.
• Figure 2: Drift error plots using \ref{['eq:drift-error']} compare the errors between our algorithm and the score-based model for Example 1 in the 1D case (left) and the 3D case (right). The 1D experiment was conducted over the time interval from 0 to 10 with a time step of 0.1, while the 3D experiment was run from 0 to 20 with the same time step. The same parameters, including the neural network architecture, were used for both experiments and both methods.
• Figure 3: Evolution of the distribution from $t=0$ to $t=2$ from Experiment 2 with step size $0.1$. The results show convergence toward the stationary distribution at $t=2$. Bright regions indicate areas of high particle density, darker regions indicate low density, and the contours represent the stationary solution.
• Figure 4: (A) The $x$-marginal and $v$-marginal at time $t=10$. (B) Convergence of the Kullback--Leibler divergence to the stationary distribution for $t \in [0,10]$ under the particle JKO dynamics from Experiment 2.
• Figure 5: Experiment from Example 3. Evolution of the distribution from $t=0$ to $t=2$ with step size $0.1$, using the sine potential. The results show convergence toward the stationary distribution at $t=2$. Bright regions indicate areas of high particle density, darker regions indicate low density, and the contours represent the stationary solution.

Theorems & Definitions (15)

• Proposition 1
• proof
• Proposition 2
• proof
• Proposition 3
• proof
• Proposition 4: Free energy dissipation for \ref{['JKO0']}
• proof
• Proposition 5: Kinetic neural ODEs
• proof