Learn to Evolve: Self-supervised Neural JKO Operator for Wasserstein Gradient Flow

TL;DR

The paper tackles the computational bottleneck of solving JKO subproblems in Wasserstein gradient flows by learning a single neural operator that maps a density to its JKO update, enabling efficient multistep evolution through repeated application. It introduces the Learn-to-Evolve framework, which jointly learns the JKO operator and evolving trajectories via a self-generated data loop, providing convergence guarantees under Lipschitz assumptions and serving as dynamic data augmentation. The approach uses a Transformer-based neural operator with particle-based discretization and is validated on Aggregation, Porous Medium, and Fokker–Planck equations, demonstrating high accuracy, stability, and strong generalization to unseen initial data and parameter regimes. The work offers a scalable paradigm for proximal-operator learning in high dimensions and suggests broader applicability to other implicit iterative mappings and time-evolution problems.

Abstract

The Jordan-Kinderlehrer-Otto (JKO) scheme provides a stable variational framework for computing Wasserstein gradient flows, but its practical use is often limited by the high computational cost of repeatedly solving the JKO subproblems. We propose a self-supervised approach for learning a JKO solution operator without requiring numerical solutions of any JKO trajectories. The learned operator maps an input density directly to the minimizer of the corresponding JKO subproblem, and can be iteratively applied to efficiently generate the gradient-flow evolution. A key challenge is that only a number of initial densities are typically available for training. To address this, we introduce a Learn-to-Evolve algorithm that jointly learns the JKO operator and its induced trajectories by alternating between trajectory generation and operator updates. As training progresses, the generated data increasingly approximates true JKO trajectories. Meanwhile, this Learn-to-Evolve strategy serves as a natural form of data augmentation, significantly enhancing the generalization ability of the learned operator. Numerical experiments demonstrate the accuracy, stability, and robustness of the proposed method across various choices of energies and initial conditions.

Figures (13)

• Figure 1: Evolution of the training dataset shown on the energy functional contours. Yellow area denotes the initial density family $\mathbb{X}$; red markers indicate equilibria; the gray region is the generated training set $\mathbb{D}_k$. For illustration, we select two landmark initial from $\mathbb{X}$ and plot the generated trajectory starting from them, shown as blue dots. Green arrows show the neural operator’s predicted directions when generating $\mathbb{D}_k$, while red arrows show the updated predictions after training on $\mathbb{D}_k$. During training, $\mathbb{D}_k$ gradually converges to $\mathbb{D}^*$.
• Figure 1: Illustration of the proposed NN architecture where the output $v = NN_{\theta}(\rho)(x)$ denotes the velocity at a query point $x$. Here, a pointwise Multi-Layer Perceptron (MLP) is applied to lift the sample point from the physical dimension $d$ to the embedding dimension $h$, or vice versa. "Self Attn" and "X-Attn" refer to self-attention and cross-attention blocks respectively. The upper panel represents the encoder, in which the density $\rho$, represented by the sampled points $\{x_i\}_{i=1}^m$, serves as the prompts. One may concatenate the density values or the energy function parameters to each of the sample points. The lower panel depicts the decoder. In the KL divergence case, the target distribution $\rho_{\text{target}}$, represented by $\{y_i\}_{i=1}^{m}$, is introduced as an additional input; it is processed by a separate MLP to lift its dimension before concatenation with $X'$.
• Figure 1: Visualization of $\mathbb{D}_k$: (Columns 1-5) Evolution of the training dataset $\mathbb{D}_k$ at selected training iterations $k =0,1,2,195,717$, displayed in the domain $[-1.1,1.1]^2$. The red circles indicate the equilibrium state. (Column 6) Normalized individual logarithmic loss $\log \ell(\rho^{k,t})$ versus outer iteration $k$: green curves show losses over all inner updates, and red markers indicate outer updates where the accumulated loss (defined in \ref{['equ:indicator']}) decreases.
• Figure 2: Comparison of predicted solutions of the aggregation equation ($p=0.5$, $q=3$) for three test cases. Each column shows results at successive time steps $t$, with the corresponding energy $\mathcal{E}(\rho^t)$ displayed below (smaller values indicate closer proximity to equilibrium). For each case, the first row shows our model’s predictions and the second row those of the baseline. Red circles mark the equilibrium state.
• Figure 3: Predicted solutions of the aggregation equation ($p=0.5$, $q=4$) at each time $t$ with input uniform distribution represented by 100 (first row), 1024 (second row) and 5000 (third row) sample points. Red circles denote the theoretical equilibrium.

Theorems & Definitions (7)

• Lemma 3.1
• Proof 1
• Theorem 3.4
• Proof 2
• Remark 3.5
• Lemma 4.1
• Proof 3