{"id":"arxiv/2601.05583","url":"https://arxiv.org/abs/2601.05583","tags":["cs.LG","math.NA"],"score":0,"title":"Learn to Evolve: Self-supervised Neural JKO Operator for Wasserstein Gradient Flow","domain":"cs.LG","source":"arxiv","authors":["Xue Feng","Li Wang","Deanna Needell","Rongjie Lai"],"reviews":0,"subject":"cs.LG","summary":"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…","category":"research","comments":0,"disputed":false,"division":"cs","mentions":0,"discussed":0,"published":"2026-01-09","reactions":{}}