E(3) x SO(3) - Equivariant Networks for Spherical Deconvolution in Diffusion MRI

Axel Elaldi, Guido Gerig, Neel Dey

Keywords: Equivariance, Diffusion, MRI, fODF, Geometric Deep Learning, Spherical Deep Learning

TL;DR: We develop E(3)xSO(3) equivariant convolutional networks for better deconvolution of Diffusion MRI.

Abstract: We present Roto-Translation Equivariant Spherical Deconvolution (RT-ESD), an $E(3)\times SO(3)$ equivariant framework for sparse deconvolution of volumes where each voxel contains a spherical signal. Such 6D data naturally arises in diffusion MRI (dMRI), a medical imaging modality widely used to measure microstructure and structural connectivity. As each dMRI voxel is typically a mixture of various overlapping structures, there is a need for blind deconvolution to recover crossing anatomical structures such as white matter tracts. Existing dMRI work takes either an iterative or deep learning approach to sparse spherical deconvolution, yet it typically does not account for relationships between neighboring measurements. This work constructs equivariant deep learning layers which respect to symmetries of spatial rotations, reflections, and translations, alongside the symmetries of voxelwise spherical rotations. As a result, RT-ESD improves on previous work across several tasks including fiber recovery on the DiSCo dataset, deconvolution-derived partial volume estimation on real-world in vivo human brain dMRI, and improved downstream reconstruction of fiber tractograms on the Tractometer dataset. Our implementation is available at \url{https://github.com/AxelElaldi/e3so3_conv}.