# Fast and Scalable Optimal Transport for Brain Tractograms

arXiv: https://arxiv.org/abs/2107.02010 · HTML: https://mlsanity.com/paper/2107.02010/

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 7 of its 18 references that are in the map.
- First posted on arXiv 5 Jul 2021.

## Builds on (most cited first, 7 of 7)

- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [On Wasserstein Two Sample Testing and Related Families of Nonparametric Tests](https://mlsanity.com/paper/1509.02237/) (2015)
- [Kernel Operations on the GPU, with Autodiff, without Memory Overflows](https://mlsanity.com/paper/2004.11127/) (2020)
- [Stabilized Sparse Scaling Algorithms for Entropy Regularized Transport Problems](https://mlsanity.com/paper/1610.06519/) (2016)
- [Sinkhorn Divergences for Unbalanced Optimal Transport](https://mlsanity.com/paper/1910.12958/) (2019)
- [TractSeg - Fast and accurate white matter tract segmentation](https://mlsanity.com/paper/1805.07103/) (2018)

## Built on by (most cited first, 12 of 12)

- [Unbalanced Optimal Transport, from Theory to Numerics](https://mlsanity.com/paper/2211.08775/) (2022)
- [Accurate Point Cloud Registration with Robust Optimal Transport](https://mlsanity.com/paper/2111.00648/) (2021)
- [Multiscale guidance of protein structure prediction with heterogeneous cryo-EM data](https://mlsanity.com/paper/2506.04490/) (2025)
- [Unbalanced Low-rank Optimal Transport Solvers](https://mlsanity.com/paper/2305.19727/) (2023)
- [A Latent Implicit 3D Shape Model for Multiple Levels of Detail](https://mlsanity.com/paper/2409.06231/) (2024)
- [Fused Gromov-Wasserstein Variance Decomposition with Linear Optimal Transport](https://mlsanity.com/paper/2411.10204/) (2024)
- [Quantization-based Bounds on the Wasserstein Metric](https://mlsanity.com/paper/2506.00976/) (2025)
- [Differentiable Expectation-Maximisation and Applications to Gaussian Mixture Model Optimal Transport](https://mlsanity.com/paper/2509.02109/) (2025)
- [Quantifying the noise sensitivity of the Wasserstein metric for images](https://mlsanity.com/paper/2510.01015/) (2025)
- [Modeling Temporal scRNA-seq Data with Latent Gaussian Process and Optimal Transport](https://mlsanity.com/paper/2605.20989/) (2026)
- [RADAR: Relative Angular Divergence Across Representations](https://mlsanity.com/paper/2605.23028/) (2026)
- [FlashSinkhorn 2: Block-Sparse Entropic Optimal Transport](https://mlsanity.com/paper/2610.02395/) (2026)

## Abstract (arXiv)

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.