cs.CVMar 26, 2026
SaveDiffusion MRI Transformer with a Diffusion Space Rotary Positional Embedding (D-RoPE)
Organizations: Stanford University
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.
Figures & tables
Figure 1: Overview of dMRI data structure and model input. The purple sphere represents diffusion sampling directions at a specific b-value. 3D dMRI volumes acquired at different directions are divided into 3D patches. Spatial coordinates define the image domain and spherical coordinates describe the diffusion space. Patches are linearized with absolute positional embeddings and processed by a transformer encoder incorporating relative positional encoding (D-RoPE; see Section 3.1 and Fig. 3 ).
Figure 2: Summary of the model architecture. The different diffusion volumes are patchified via a projection layer and the absolute positional information of the image space and the diffusion space is added to the the patch embeddings. The tokens are passed through modified attention blocks that alternate attention in both image and diffusion spaces. Additionally, relative information of different diffusion directions and b-values is encoded via D-RoPE. The obtained latent representations from the encoder are used for downstream task evaluation. To close the pretraining loop, some tokens are masked and fed into transformer blocks as the ones used in the encoder, followed by 3D convolutional blocks to reconstruct the original volumes.
Figure 3: Attention calculation with D-RoPE: The relative rotation matrix depends on the distance between diffusion volumes, which is defined as a weighted combination of the distance between b-values and the angle between b-vectors.
Figure 4: Qualitative evaluation of the reconstructions for a test subject for a b-value of 2000 . The reconstruction obtained with different masking strategies (Spatial, Diffusion, and Alternating) and with and without D-RoPE are shown in the three left-most columns. The reference ground-truth image and the spatial mask that was used are shown in the two right-most columns (white=masked). In general, a good reconstruction of the brain structures are observed. The MAE without D-RoPE case shows an erroneous contrast when only diffusion space masking is performed, possibly because of the lack of spatial information for a given diffusion direction.
| Method | Masking | PSNR ( ) | SSIM ( ) | FID ( ) | |
|---|---|---|---|---|---|
| b-value = 1000 | MAE | Spatial | 14.54 ± 0.16 | 0.971 ± 0.003 | 13.37 |
| MAE | Diffusion | 10.13 ± 0.26 | 0.875 ± 0.005 | 39.76 | |
| MAE | Alternating | 14.70 ± 0.17 | 0.971 ± 0.004 | 11.13 | |
| MAE w/D-RoPE | Spatial | 14.84 ± 0.17 | 0.973 ± 0.003 | 10.28 | |
| MAE w/D-RoPE | Diffusion | 14.86 ± 0.17 | 0.974 ± 0.003 | 10.73 | |
| MAE w/D-RoPE | Alternating | 14.94 ± 0.17 | 0.974 ± 0.003 | 8.36 |
Table 1: Quantitative evaluation of the reconstructions obtained from different masking strategies (Spatial, Diffusion, and Alternating) and with and without D-RoPE in the test set of HCP-YA. The mean standard deviation is shown for Peak Signal-to-Noise Ratio (PSNR) and Structural Similarity Index Measure (SSIM). The value of the Fréchet Inception distance (FID) is shown in the last column. : the higher, the better; : the lower, the better. The best results per b-value are highlighted in bold.
| Method | Masking | FA error ( ) | MD error x ( ) |
|---|---|---|---|
| MAE | Spatial | 0.074 ± 0.013 | 0.098 ± 0.009 |
| MAE | Diffusion | 0.185 ± 0.012 | 0.233 ± 0.015 |
| MAE | Alternating | 0.080 ± 0.013 | 0.109 ± 0.008 |
| MAE w/D-RoPE | Spatial | 0.061 ± 0.007 | 0.104 ± 0.008 |
| MAE w/D-RoPE | Diffusion | 0.058 ± 0.007 | 0.099 ± 0.008 |
| MAE w/D-RoPE | Alternating | 0.055 ± 0.007 | 0.103 ± 0.008 |
Table 2: Evaluation of fractional anisotropy (FA) and mean diffusivity (MD) obtained from the reconstructions with and without D-RoPE in HCP-YA. For MD, the x1000 error is shown. The mean standard deviation of the absolute error is shown. : the lower, the better. The best results are highlighted in bold.
| Age prediction | Sex classification | |||||||
|---|---|---|---|---|---|---|---|---|
| Method | Features | Prediction Head | Supervision | ( ) | MSE ( ) | ACC ( ) | AUROC ( ) | |
| HC Features | Tract FA, MD | Linear | Full | 0.90 ± 0.01 | 0.20 ± 0.01 | 70.90 ± 2.60 % | 0.77 ± 0.02 | |
| HC Features | Tract FA, MD | MLP | Full | 0.90 ± 0.01 | 0.19 ± 0.02 | 56.90 ± 2.50 % | 0.59 ± 0.03 | |
| HC Features | Tract based RISH | MLP | Full | 0.61 ± 0.21 | 1.15 ± 1.24 | 61.70 ± 2.00 % | 0.67 ± 0.03 | |
| 3D ResNet | FA and MD maps | MLP | Full | 0.96 ± 0.01 | 0.09 ± 0.02 | 76.00 ± 1.20 % | 0.83 ± 0.04 | |
| ViT w/RoPE | FA and MD maps | MLP | Full | 0.93 ± 0.01 | 0 .15 ± 0.02 | 65.40 ± 3.70 % | 0.72 ± 0.04 | |
Table 3: Evaluation of different methods (handcrafted features, 3D ResNet, ViT, and our partially supervised methods) on the age prediction and sex classification tasks in HCP-A and HCP-D. For age prediction, the correlation coefficient ( ) and mean square error (MSE) are reported. For sex classification, accuracy (ACC) and area under the receiver operating curve (AUROC) are reported. : the higher, the better, : the lower, the better. The best result per metric is highlighted in bold, and the second best is underlined. Our proposed methods are highlighted in yellow. *HC Features means Hand-Crafted Features.
| CN vs MCI classification | ADAS prediction | ||||||||
|---|---|---|---|---|---|---|---|---|---|
| Method | Features | Prediction Head | Supervision | BA ( ) | AUPRC ( ) | AUROC ( ) | ( ) | MSE ( ) | |
| HC Features | Tract FA, MD | Linear | Full | 6 1.90 ± 7.30 % | 0.53 ± 0.06 | 0.69 ± 0.06 | 0.06 ± 0.15 | 10.92 ± 4.08 | |
| HC Features | Tract FA, MD | MLP | Full | 57.40 ± 5.40 % | 0.44 ± 0.09 | 0.62 ± 0.04 | 0 .33 ± 0.09 | 0 .81 ± 0.13 | |
| HC Features | Tract based RISH | MLP | Full | 53.90 ± 9.60 % | 0.31 ± 0.07 | 0.52 ± 0.08 | 0.28 ± 0.14 | 0.93 ± 0.18 | |
| 3D ResNet | FA and MD maps | MLP | Full | 58.40 ± 8.30 % | 0 .50 ± 0.10 | 0 .68 ± 0.08 | 0 .33 ± 0.08 | 0.86 ± 0.14 | |
| ViT w/RoPE | FA and MD maps | MLP | Full | 58.10 ± 6.60 % | 0.47 ± 0.08 | 0.65 ± 0.06 | 0.25 ± 0.12 | 0.93 ± 0.15 | |
Table 4: Evaluation of different methods (handcrafted features, 3D ResNet, ViT, and our partially supervised methods) on the CN vs MCI classification and ADAS prediction tasks in ADNI. For CN vs MCI, balance accuracy (BA), area under the precision-recall curve (AUPRC), and area under the receiver operating curve (AUROC) are reported. For ADAS prediction, the correlation coefficient ( ) and mean square error (MSE) are reported. : the lower, the better. : the lower, the better. The best result per metric is highlighted in bold, and the second best is underlined. Our proposed methods are highlighted in yellow. *HC Features means Hand-Crafted Features.
| Hyperparameter | Value |
|---|---|
| Number of transformer encoding Layers | 10 |
| Number of transformer decoding Layers | 3 |
| Number of convolutional decoding Layers | 3 |
| Latent dimension | 384 |
| Number of heads | 3 |
| Convolutional kernel size | 3 |
Table 5: Parameters using for the pre-training stage of the MAE.
| Task | ||||
| hyperparameter | Age prediction | Sex Classification | MCI classification | ADAS prediction |
| optimizer | AdamW | AdamW | AdamW | AdamW |
| optimizer momentum | , | , | , | , |
| learning rate schedule | Constant | cosine | Constant | cosine |
| start learning rate | 5.00E-05 | 1.00E-05 | 5.00E-05 | 5.00E-05 |
| final learning rate | 5.00E-05 | 1.00E-06 | 5.00E-05 | 1.00E-06 |
Table 6: Parameters using for the partial finetuning of the MAE in different downstream tasks.
Figure 5: Qualitative evaluation of the reconstructions for a test subject for a b-value of 1000 . The reconstruction obtained from different masking strategies (Spatial, Diffusion and Alternating) and with and without D-RoPE are shown in the three left-most columns. The reference ground-truth image and the spatial mask that was used are shown in the two right-most columns.
Figure 6: Qualitative evaluation of the reconstructions for a test subject for a b-value of 3000 . The reconstruction obtained from different masking strategies (Spatial, Diffusion and Alternating) and with and without D-RoPE are shown in the three left-most columns. The reference ground-truth image and the spatial mask that was used are shown in the two right-most columns.
| Age prediction | Sex classification | |||||
|---|---|---|---|---|---|---|
| Trained on | Tested on | ( ) | MSE ( ) | ACC ( ) | AUROC ( ) | |
| Protocol 1 | Protocol 1 | 0.895 ± 0.010 | 0.22 ± 0.02 | 75.0 ± 1.2 % | 0.827 ± 0.015 | |
| Protocol 1 | Protocol 2 | 0.897 ± 0.007 | 0.21 ± 0.02 | 75.4 ± 2.0 % | 0.828 ± 0.017 | |
Table 7: Comparison of linear probing on top of latent representations trained on protocol 1 and tested on this and a second different protocol. Minimal changes in performance are observed.
Figure 7: Linear probing performance of the frozen latents with different patch sizes (Top) and at different stages of pretraining (Bottom). A clear performance improvement in all downstream tasks is observed as one goes to smaller patch sizes. The optimal number of pretraining epochs seems to vary depending on the downstream task.
Figure 8: Two dimensional UMAP projections for the HCP-A samples colored by sex on the left and combined HCP-A and HCP-D samples colore by age on the right. We observed a structure of the latent space in which female subjects appear to be concentrated towards the lower part of the manifold, and in which there is a gradient from younger to older subjects going upwards in the manifold.