学会演化：用于 Wasserstein 梯度流的自监督神经 JKO 算子
机器学习
2026-01-12 v1 数值分析
数值分析
摘要
Jordan-Kinderlehrer-Otto (JKO) 格式为计算 Wasserstein 梯度流提供了一个稳定的变分框架，但其实际应用常因反复求解 JKO 子问题的高昂计算成本而受限。我们提出了一种自监督方法，用于学习 JKO 解算子，而无需任何 JKO 轨迹的数值解。学习到的算子将输入密度直接映射到相应 JKO 子问题的极小值点，并可迭代应用以高效生成梯度流演化。一个关键挑战是，通常只有若干初始密度可用于训练。为解决此问题，我们引入了一种“学会演化”算法，通过在轨迹生成和算子更新之间交替进行，联合学习 JKO 算子及其诱导的轨迹。随着训练的进行，生成的数据越来越逼近真实的 JKO 轨迹。同时，这种“学会演化”策略作为一种自然的数据增强形式，显著增强了所学算子的泛化能力。数值实验证明了所提方法在各种能量和初始条件选择下的准确性、稳定性和鲁棒性。
引用
@article{arxiv.2601.05583,
title = {Learn to Evolve: Self-supervised Neural JKO Operator for Wasserstein Gradient Flow},
author = {Xue Feng and Li Wang and Deanna Needell and Rongjie Lai},
journal= {arXiv preprint arXiv:2601.05583},
year = {2026}
}