Rotation-Equivariant Deep Learning for Diffusion MRI

Abstract

Convolutional networks are successful, but they have recently been outperformed by new neural networks that are equivariant under rotations and translations. These new networks work better because they do not struggle with learning each possible orientation of each image feature separately. So far, they have been proposed for 2D and 3D data. Here we generalize them to 6D diffusion MRI data, ensuring joint equivariance under 3D roto-translations in image space and the matching 3D rotations in q-space, as dictated by the image formation. Such equivariant deep learning is appropriate for diffusion MRI, because microstructural and macrostructural features such as neural fibers can appear at many different orientations, and because even non-rotation-equivariant deep learning has so far been the best method for many diffusion MRI tasks. We validate our equivariant method on multiple-sclerosis lesion segmentation. Our proposed neural networks yield better results and require fewer scans for training compared to non-rotation-equivariant deep learning. They also inherit all the advantages of deep learning over classical diffusion MRI methods. Our implementation is available at https://github. com/philip-mueller/equivariant-deep-dmri and can be used off the shelf without understanding the mathematical background.

Equivariance versus Augmentation for Spherical Images

“…Data augmentation in the context of symmetric tasks was studied previously in [13], where a method to align input data is presented and compared to data augmentation. Closest to our work is a comparison between an equivariant model and a non-equivariant model for reduced training data sizes for an MRI application in [22]. In contrast, we systematically compare data augmentation with increased training data sizes for different tasks, datasets and several non-equivariant models with an equivariant architecture.…”

Geometric Deep Learning and Equivariant Neural Networks

“…for any nonlinear function α : R → R. Examples for α used in the literature include sigmoid [10], relu [43,27], shifted soft plus [37] and swish [1]. In [35], norm nonlinearities are used for matrix valued feature maps with || • || = Re(tr(•)) and α = relu.…”

“…This is directly relevant for instance in the case of spherical signals where G is a rotation group. In practical applications, it was found that equivariance improves per-sample efficiency, reducing the need for data augmentation [1]. For linear models, this has been proven mathematically [2].…”

“…A further important case studied extensively in the literature are networks equivariant with respect to E(n) or SE(n). These networks have been applied with great success to 3d shape classification [10,37], protein structure classification [10], atomic potential prediction [38] and medical imaging [1,39,40].…”

Symmetry-aware Neural Architecture for Embodied Visual Navigation

“…Novel convolution layers with different equivariance have been proposed so far, group equivariant convolution networks [12,15], steerable convolution networks and harmonic networks for rotation equivariance, scale equivariance [23,47], and permutation equivariance [44]. Such convolution layers equipped with various types of equivariance have been proven to be beneficial to better performance in tracking [41], classification, trajectory prediction [46], segmentation [30], and image generation [14]. However, these methods have never been applied to visual navigation.…”