Fast and Scalable Optimal Transport for Brain Tractograms

We present a new multiscale algorithm for solving regularized Optimal Transport problems on the GPU, with a linear memory footprint. Relying on Sinkhorn divergences which are convex, smooth and positive definite loss functions, this method enables the computation of transport plans between millions of points in a matter of minutes. We show the effectiveness of this approach on brain tractograms modeled either as bundles of fibers or as track density maps. We use the resulting smooth assignments to perform label transfer for atlas-based segmentation of fiber tractograms. The parameters – blur and reach – of our method are meaningful, defining the minimum and maximum distance at which two fibers are compared with each other. They can be set according to anatomical knowledge. Furthermore, we also propose to estimate a probabilistic atlas of a population of track density maps as a Wasserstein barycenter. Our CUDA implementation is endowed with a user-friendly PyTorch interface, freely available on the PyPi repository (pip install geomloss) and at www.kernel-operations.io/geomloss.

[1]  Bernhard Schmitzer,et al.  Stabilized Sparse Scaling Algorithms for Entropy Regularized Transport Problems , 2016, SIAM J. Sci. Comput..

[2]  Alan L. Yuille,et al.  The invisible hand algorithm: Solving the assignment problem with statistical physics , 1994, Neural Networks.

[3]  Anand Rangarajan,et al.  A new point matching algorithm for non-rigid registration , 2003, Comput. Vis. Image Underst..

[4]  Rachid Deriche,et al.  Unsupervised white matter fiber clustering and tract probability map generation: Applications of a Gaussian process framework for white matter fibers , 2010, NeuroImage.

[5]  Martin Styner,et al.  Exploratory Population Analysis with Unbalanced Optimal Transport , 2018, MICCAI.

[6]  Marco Cuturi,et al.  Sinkhorn Distances: Lightspeed Computation of Optimal Transport , 2013, NIPS.

[7]  Jean Feydy,et al.  Sinkhorn Divergences for Unbalanced Optimal Transport , 2019, ArXiv.

[8]  Gabriel Peyré,et al.  Computational Optimal Transport , 2018, Found. Trends Mach. Learn..

[9]  Paolo Avesani,et al.  White Matter Tract Segmentation as Multiple Linear Assignment Problems , 2018, Front. Neurosci..

[10]  Marco Cuturi,et al.  On Wasserstein Two-Sample Testing and Related Families of Nonparametric Tests , 2015, Entropy.

[11]  Alain Trouvé,et al.  Interpolating between Optimal Transport and MMD using Sinkhorn Divergences , 2018, AISTATS.

[12]  Yogesh Rathi,et al.  An anatomically curated fiber clustering white matter atlas for consistent white matter tract parcellation across the lifespan , 2018, NeuroImage.

[13]  Jean Feydy,et al.  Kernel Operations on the GPU, with Autodiff, without Memory Overflows , 2020, J. Mach. Learn. Res..

[14]  Alain Trouvé,et al.  Global Divergences Between Measures: From Hausdorff Distance to Optimal Transport , 2018, ShapeMI@MICCAI.

[15]  Peter F. Neher,et al.  TractSeg - Fast and accurate white matter tract segmentation , 2018, NeuroImage.

[16]  François-Xavier Vialard,et al.  An Interpolating Distance Between Optimal Transport and Fisher–Rao Metrics , 2010, Foundations of Computational Mathematics.

[17]  Isabelle Bloch,et al.  White Matter Multi-Resolution Segmentation Using Fuzzy Set Theory , 2019, 2019 IEEE 16th International Symposium on Biomedical Imaging (ISBI 2019).

[18]  Maxime Descoteaux,et al.  Recognition of white matter bundles using local and global streamline-based registration and clustering , 2017, NeuroImage.