Score-Based Generative Neural Networks for Large-Scale Optimal Transport
Max Daniels, Massachusetts Institute of Technology
Given two datasets, it is often useful to compute a matching between their individual datapoints, which can then be used to make inferences and comparisons between the two datasets. One way to compute a matching is to apply the theory of Optimal Transportation (OT) of probability distributions. This theory can be used to write a variational problem whose solution, known as an “optimal coupling,” induces the matching between two given data-generating distributions. However, many existing numerical methods struggle to scale to large, high-dimensional datasets which are often used in machine learning and computational sciences. We develop an approximate scheme to compute optimal transport couplings which is based on (1) a form of entropic regularization of the OT variational problem, known as Sinkhorn regularization, and on (2) a neural network parametrization of the Sinkhorn problem. We demonstrate the empirical success of our method on a variety of large-scale OT problems. Some potential applications include lineage tracing in single-cell biology, domain adaptation for the correction of mis-specified statistical models and keyframe interpolation in computational graphics.
Abstract Author(s): Max Daniels, Tyler Maunu, Paul Hand