Solving Inverse Problems of Distributions with Wasserstein Flow Maps

Many scientific datasets consist of distributions rather than individual points, such as point clouds, cell populations or spatial tissue neighborhoods. Recovering populations from partial or aggregate observations, and steering their generation toward desired properties, call for generative priors that can adapt to different measurement operators and guidance objectives. Wasserstein flow matching provides such priors, but obtaining terminal populations requires numerical integration, making repeated endpoint prediction for inference-time guidance computationally expensive. We introduce Wasserstein flow maps, a framework for learning finite-time transport maps between distribution-valued states, enabling direct prediction of terminal populations. We formulate these maps on Wasserstein space and derive semigroup, Eulerian, and Lagrangian characterizations. The resulting maps support deterministic reward lookahead and flow map trajectory tilting, allowing one pretrained prior per domain to address multiple inverse problems and guidance objectives without task-specific retraining. Across ShapeNet, single-cell populations, and spatial transcriptomic niches, Wasserstein flow maps improve generation metrics over Wasserstein flow matching baselines using 8 rather than 100 function evaluations, and outperform inverse problems methods on inpainting, super-resolution, deconvolution, and guided generation. On 100 held-out patients, lookahead guidance reduces reconstruction MMD by 35% for deconvolution from mean single-cell RNAseq embeddings and masked MMD by 37% for cell-type inpainting, relative to task-specific fine-tuned Wasserstein flow matching.