Diffusion MRI Transformer with a Diffusion Space Rotary Positional Embedding (D-RoPE)

TL;DR

This paper introduces D-RoPE, a diffusion space rotary positional embedding for diffusion MRI transformers, enabling robust, transferable representations across diverse protocols and improving downstream task performance.

Contribution

The novel D-RoPE embedding captures spatial and directional properties of dMRI, enhancing model transferability and performance across varied acquisition settings.

Findings

Pretrained D-RoPE model outperforms baselines in downstream tasks.

6% higher accuracy in classifying mild cognitive impairment.

0.05 increase in correlation coefficient for cognitive score prediction.

Abstract

Diffusion Magnetic Resonance Imaging (dMRI) plays a critical role in studying microstructural changes in the brain. It is, therefore, widely used in clinical practice; yet progress in learning general-purpose representations from dMRI has been limited. A key challenge is that existing deep learning approaches are not well-suited to capture the unique properties of diffusion signals. Brain dMRI is normally composed of several brain volumes, each with different attenuation characteristics dependent on the direction and strength of the diffusion-sensitized gradients. Thus, there is a need to jointly model spatial, diffusion-weighting, and directional dependencies in dMRI. Furthermore, varying acquisition protocols (e.g., differing numbers of directions) further limit traditional models. To address these gaps, we introduce a diffusion space rotatory positional embedding (D-RoPE) plugged into our dMRI transformer to capture both the spatial structure and directional characteristics of diffusion data, enabling robust and transferable representations across diverse acquisition settings and an arbitrary number of diffusion directions. After self-supervised masked autoencoding pretraining, tests on several downstream tasks show that the learned representations and the pretrained model can provide competitive or superior performance compared to several baselines in these downstream tasks (even compared to a fully trained baseline); the finetuned features from our pretrained encoder resulted in a 6% higher accuracy in classifying mild cognitive impairment and a 0.05 increase in the correlation coefficient when predicting cognitive scores. Code is available at: github.com/gustavochau/D-RoPE.