Fully Differentiable dMRI Streamline Propagation in PyTorch

TL;DR

This paper tackles the challenge of integrating tractography into end-to-end learning by delivering a fully differentiable PyTorch implementation of the MRtrix SD_Stream streamline propagator. By translating the MRtrix C++ algorithm into a differentiable PyTorch module with batched, gradient-preserving control flow, the authors preserve numerical fidelity while enabling backpropagation through streamline propagation. The differentiable propagator closely matches MRtrix results (over 93% of 10,000 streamlines with Hausdorff distance $<1$ mm) but incurs substantial speed and memory costs, highlighting a tradeoff between differentiability and computational efficiency. This work lays a foundation for end-to-end macrostructural reasoning in diffusion MRI and seamless incorporation of tractography into deep learning pipelines.

Abstract

Diffusion MRI (dMRI) provides a distinctive means to probe the microstructural architecture of living tissue, facilitating applications such as brain connectivity analysis, modeling across multiple conditions, and the estimation of macrostructural features. Tractography, which emerged in the final years of the 20th century and accelerated in the early 21st century, is a technique for visualizing white matter pathways in the brain using dMRI. Most diffusion tractography methods rely on procedural streamline propagators or global energy minimization methods. Although recent advancements in deep learning have enabled tasks that were previously challenging, existing tractography approaches are often non-differentiable, limiting their integration in end-to-end learning frameworks. While progress has been made in representing streamlines in differentiable frameworks, no existing method offers fully differentiable propagation. In this work, we propose a fully differentiable solution that retains numerical fidelity with a leading streamline algorithm. The key is that our PyTorch-engineered streamline propagator has no components that block gradient flow, making it fully differentiable. We show that our method matches standard propagators while remaining differentiable. By translating streamline propagation into a differentiable PyTorch framework, we enable deeper integration of tractography into deep learning workflows, laying the foundation for a new category of macrostructural reasoning that is not only computationally robust but also scientifically rigorous.

Figures (4)

• Figure 1: MI-Brainrheault2016mi visualization of 10,000 streamlines from same FODs using SD_Stream from MRtrix (left) and PyTorch (right). Given the same seed points and initial propagation directions, differentiable PyTorch streamline propagator reproduces streamlines that are visually indistinguishable from MRtrix.
• Figure 2: Newton–Raphson peak-finding in MRtrix (Single Direction)
• Figure 3: MRtrix SD_Stream Propagation (Single Streamline)
• Figure 4: Violin plot of the log-scaled Hausdorff distances between 10,000 streamlines generated by MRtrix and PyTorch. Colored markers indicate the 10th, 50th, 95th, and 99th percentiles. The corresponding streamlines at each percentile are visualized on the right, comparing spatial alignment between the PyTorch (blue) and MRtrix (red) methods.