SOE: SO(3)-Equivariant 3D MRI Encoding

TL;DR

This work proposes a novel method for encoding 3D MRIs that enforces equivariance with respect to all rotations in 3D space, in other words, SO(3)-equivariance (SOE), and evaluates SOE pretrained on the structural MRIs of two public data sets with respect to the downstream task of predicting age and diagnosing Alzheimer's Disease from T1-weighted brain scans of the ADNI data set.

Abstract

Representation learning has become increasingly important, especially as powerful models have shifted towards learning latent representations before fine-tuning for downstream tasks. This approach is particularly valuable in leveraging the structural information within brain anatomy. However, a common limitation of recent models developed for MRIs is their tendency to ignore or remove geometric information, such as translation and rotation, thereby creating invariance with respect to geometric operations. We contend that incorporating knowledge about these geometric transformations into the model can significantly enhance its ability to learn more detailed anatomical information within brain structures. As a result, we propose a novel method for encoding 3D MRIs that enforces equivariance with respect to all rotations in 3D space, in other words, SO(3)-equivariance (SOE). By explicitly modeling this geometric equivariance in the representation space, we ensure that any rotational operation applied to the input image space is also reflected in the embedding representation space. This approach requires moving beyond traditional representation learning methods, as we need a representation vector space that allows for the application of the same SO(3) operation in that space. To facilitate this, we leverage the concept of vector neurons. The representation space formed by our method captures the brain's structural and anatomical information more effectively. We evaluate SOE pretrained on the structural MRIs of two public data sets with respect to the downstream task of predicting age and diagnosing Alzheimer's Disease from T1-weighted brain scans of the ADNI data set. We demonstrate that our approach not only outperforms other methods but is also robust against various degrees of rotation along different axes. The code is available at https://github.com/shizhehe/SOE-representation-learning.

Figures (2)

• Figure 1: SOE Model Architecture: The model takes in a pair of unrotated and rotated samples $x^1, x^2$ as input and maps them to representations preserving the rotational transformation $R_{rot}$. The Encoder encodes the input volume into a $d$ dimensional scalar array and the VN module maps that scalar array into a $d' \times 3$ dimensional vector embedding. $R_{rot}$ is applied on the input volume through the $Rot(\cdot)$ function.
• Figure 2: ADNI Classification: Comparison of rotation augmented SOE vs. no pretraining model. Displayed are the % increases in BACC (left) and F1 (right) from a model with no pretraining to a SOE pretrained model. The labels on each x axis shows the rotation degrees employed during training and the y axis rotation degrees employed during evaluation.