Accelerated-JKO: Discretizing Wasserstein Accelerated-Information Gradient Flow

Wasserstein gradient flows (WGF), discretized by the JKO scheme, are widely adopted in optimization, sampling and generative model. JKO inherits the $O(e^{-\beta t})$ convergence rate of WGF, which degrades for densities that are not strongly log-concave. Accelerated Information Gradient Flow augments the density with a momentum for improved convergence rate, but existing realizations integrate the accelerated particle ODE explicitly which does not fully leverage the well-posed benefit of the JKO proximal step. We introduce A-JKO and its folding-free variant IA-JKO, momentum-augmented JKO schemes that discretize the AIGF momentum dynamics while keeping each step a well-posed variational problem, and prove that, as the step size vanishes, both schemes converge to the Wasserstein AIGF. Building on this scheme, we construct a neural sampler that transports a prior to the target through a flow-based velocity field, trained with a off-policy matching loss derived from the JKO optimality condition. On synthetic (DW-4, LJ-13, LJ-55) and molecular (alanine dipeptide) benchmarks, our sampler outperforms conventional JKO-based neural samplers and is competitive with several recent neural samplers.