Fast and Scalable Optimal Transport for Brain Tractograms

Abstract

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 reachof 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.

Optimal transport: discretization and algorithms

“…The book of Cuturi and Peyré present these difficulties in more details and explain how to circumvent them [83]. In addition to the works already cited, we refer to the PhD work of Feydy [29,45], and especially to the implementation of regularized optimal transport in the library GeomLoss 2 .…”

A Stochastic Multi-layer Algorithm for Semi-discrete Optimal Transport with Applications to Texture Synthesis and Style Transfer

“…. , Y L ) and the clusters (C y ) 0 <L,y∈Y +1 satisfying (11). Let us also consider the multi-layer parameters…”

“…Transportation maps acting on images can then be used to perform image warping [1,22,43,35,12], which is relevant for images that reflect some kind of density of material. Notice that similar tools can also be used to process shapes identified as probability distributions which allows to perform shape interpolation [57] and shape registration [10] with applications in brain imaging [11].…”

“…In the case of large-scale discrete OT, the author of [53,54] proposed multiscale algorithms for unregularized and regularized OT problems that exploit a hierarchical representation of the supports of the source and target distributions in order to drastically diminish the number of variables in the dual formulation of OT. This coarse-to-fine strategy was also used in [11] to accelerate computations on very large datasets. However, for semi-discrete OT in a high-dimensional setting, these multiscale schemes are hardly tractable.…”

Nonlinear Alignment of Whole Tractograms with the Linear Assignment Problem

“…There, the building block of the methodology is to compute which streamline of the first subject corresponds to which streamline of the second subject, as in a combinatorial optimization problem. The principle of streamline correspondence has also been used for the problem of bundle segmentation [6,8,12,13,19].…”