Paper "Scalable Supervised Optimal Transport of Gaussian Mixture Models" accepted at NeurIPS 2026

Our paper Scalable Supervised Optimal Transport of Gaussian Mixture Models was accepted for publication at NeurIPS 2026.

Optimal transport (OT) is a mathematical framework to measure distances between distributions (e.g. the cells of individual patients) as the minimal cost of moving one distribution onto the other. Defining the cost with a fixed ground metric such as the Euclidean distance often captures variation that is irrelevant for the task, so we propose a method to learn the ground metric such that the OT distance reflects class relations, such as disease states of patients. To scale this to large datasets, we represent each distribution as a Gaussian mixture model and learn the metric between mixture components instead of individual cells. This replaces large OT problems with small ones, where the cost between Gaussians has a closed form.