Deep Kinetic JKO schemes for Vlasov-Fokker-Planck Equations

TL;DR

This paper develops a deep neural network-based numerical scheme inspired by the JKO approach to efficiently solve high-dimensional kinetic Fokker-Planck equations, preserving their variational structure.

Contribution

It introduces a novel deep neural network framework for kinetic JKO schemes that handle high-dimensional Vlasov-Fokker-Planck equations while maintaining their structural properties.

Findings

Effective for high-dimensional kinetic PDEs

Preserves variational and structural properties

Validated through extensive numerical experiments

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.