Generalized Wasserstein Flow Matching: Transport Plans, Everywhere, All at Once

TL;DR

This paper introduces a novel Wasserstein-on-Wasserstein framework for flow matching in probability measure spaces, enabling efficient and scalable generative modeling with theoretical guarantees.

Contribution

It extends flow matching to nested probability measures, proposing scalable approximations and unifying various generative modeling approaches.

Findings

Introduces Wasserstein-on-Wasserstein formulation for flow matching.

Develops scalable sliced and linear Wasserstein approximations.

Provides a unified, theoretically grounded generative modeling framework.

Abstract

Flow matching has recently emerged as a flexible and efficient framework for generative modelling by learning deterministic transport dynamics between probability measures. In this work, we extend flow matching to the space of probability measures over probability measures, introducing a Wasserstein-on-Wasserstein (WoW) formulation. Leveraging the nested Wasserstein geometry, we show that measures over transport plans naturally induce velocity fields that realize metameasure flows. This yields a principled generalization of Wasserstein flow matching via coupled outer and inner transport plans. To address the substantial computational cost of WoW transport, we propose scalable approximations based on sliced and linear Wasserstein distances, enabling efficient training while promoting numerically stable, near-straight trajectories. Our framework unifies and extends existing approaches to point cloud and set generation, providing a practical and theoretically grounded method for generative modelling in WoW spaces.