Gradient Flow Algorithms for Density Propagation in Stochastic Systems

We develop a new computational framework to solve the partial differential equations (PDEs) governing the flow of the joint probability density functions (PDFs) in continuous-time stochastic nonlinear systems. The need for computing the transient joint PDFs subject to prior dynamics arises in uncertainty propagation, nonlinear filtering and stochastic control. Our methodology breaks away from the traditional approach of spatial discretization or function approximation – both of which, in general, suffer from the “curse-of-dimensionality”. In the proposed framework, we discretize time but not the state space. We solve infinite dimensional proximal recursions in the manifold of joint PDFs, which in the small time-step limit, is theoretically equivalent to solving the underlying transport PDEs. The resulting computation has the geometric interpretation of gradient flow of certain free energy functional with respect to the Wasserstein metric arising from the theory of optimal mass transport. We show that dualization along with an entropic regularization, leads to a cone-preserving fixed point recursion that is proved to be contractive in Thompson metric. A block co-ordinate iteration scheme is proposed to solve the resulting nonlinear recursions with guaranteed convergence. This approach enables remarkably fast computation for non-parametric transient joint PDF propagation. Numerical examples and various extensions are provided to illustrate the scope and efficacy of the proposed approach.

💡 Research Summary

The paper introduces a novel computational framework for propagating joint probability density functions (PDFs) of continuous‑time stochastic nonlinear systems without discretizing the state space. Traditional approaches rely on spatial discretization (finite‑difference, finite‑element, spectral methods) or function approximation (Gaussian mixtures, polynomial chaos), both of which suffer severely from the curse of dimensionality. The authors instead discretize only time and treat the PDF as an element of an infinite‑dimensional manifold equipped with the 2‑Wasserstein metric.

The core idea is to view the Fokker‑Planck‑Kolmogorov (FPK) PDE governing the PDF evolution as a gradient flow of a free‑energy functional Φ(ρ) with respect to the Wasserstein‑2 distance. For systems whose drift is the gradient of a potential ψ and whose diffusion is isotropic (the “JKO canonical form”), the free energy is
Φ(ρ)=∫ψ(x)ρ(x)dx + β⁻¹∫ρ(x)logρ(x)dx,
which is a Lyapunov functional decreasing along the true PDF trajectory. The proximal recursion
ρₖ = arg min_{ρ∈𝔇₂} ½ W²(ρ,ρₖ₋₁) + h Φ(ρ)
is precisely the Jordan‑Kinderlehrer‑Otto (JKO) scheme; as the time step h→0 the iterates converge to the solution of the FPK PDE.

Directly solving the proximal step is infeasible because the Wasserstein distance itself is defined via an optimal‑transport linear program. To obtain a tractable algorithm, the authors add an entropic regularization term ε H(M) (where H is the matrix entropy) to the transport problem and then dualize. The regularized problem becomes a strictly convex optimization over transport plans M with a simple closed‑form dual that leads to a fixed‑point equation of the Sinkhorn‑type:
M_{ij}=exp