3D Mamba MAE

Abstract

Learning meaningful and interpretable representations from high-dimensional volumetric magnetic resonance (MR) images is essential for advancing personalized medicine. While Vision Transformers (ViTs) have shown promise in handling image data, their application to 3D multi-contrast MR images faces challenges due to computational complexity and interpretability. To address this, we propose a novel state-space-model (SSM)-based masked autoencoder which scales ViT-like models to handle high-resolution data effectively while also enhancing the interpretability of learned representations. We propose a latent-to-spatial mapping technique that enables direct visualization of how latent features correspond to specific regions in the input volumes in the context of SSM. We validate our method on two key neuro-oncology tasks: identification of isocitrate dehydrogenase mutation status and 1p/19q co-deletion classification, achieving state-of-the-art accuracy. Our results highlight the potential of SSM-based self-supervised learning to transform radiomics analysis by combining efficiency and interpretability.

Results under different patch sizes

Installing the dependencies

After you clone the project from github, and then please install the dependencies using the following commands:

  conda create -n vim python=3.10
  conda activate vim
  conda install pytorch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 pytorch-cuda=12.1 -c pytorch -c nvidia
  conda install cuda -c nvidia/label/cuda-12.1.0
  pip install -r vim_requirements.txt
  pip install -e mamba2
  pip install tensorboard torchio rotary-embedding-torch openpyxl

Dataset Description

• In this work, we use Brats2022 (Brats2021) dataset for pre-training and EGD dataset for fine-tuning.
• A label description file is provided for reference.
• After you have downloaded the data, please remember to process them into 160 x 160 x 160 volumes and make sure the data value is integer and the range is [0, 255].
• Create a file named config.ini, with the following content:

  [brats_160]
  BASE_PATH = "your preprocessed brats2022 path"
  TEST_BASE_PATH = "your preprocessed egd path"

Pre-training

Please note that both pre-training and fine-tuning stages require at least 24GB GPU memory for models with path size 4.

To pre-training the model, please use the following script:

python -m torch.distributed.launch --nproc_per_node=4 --master_port=27590 --use_env main_pretrain.py --model mae_3d_vim_small_patch4 --batch_size 4  --norm_pix_loss --mask_ratio 0.75 --epochs 1000 --warmup_epochs 40 --blr 2.5e-4 --weight_decay 0.05 --use_3d --name_dataset brats-160  --input_size 160 --seed 42 --output_dir vim_small_brats_160_patch4 --log_dir vim_small_brats_160_patch4  --num_workers 8 --use_z_score --pin_mem

NOTE: To train models with other patch sizes, you just need to change the model name. If you don't have enough RAM, you can use --no_pin_men flag to reduce the RAM computation

After you download the pre-trained weight, please unzip the file and put folder in your main directory.

Finetuning script with 5-fold cross-validtion

To evaluate the model performance on 1p/19q co-deletion classification, the launch script are as follows:

python -m torch.distributed.launch --nproc_per_node=4 --master_port=27590 --use_env main_cr_finetune.py --model vim_3D_small_patch4_stride4_224_bimambav2_final_pool_mean_abs_pos_embed_div2  --batch_size 2  --epochs 100   --name_dataset brats-160  --lr 1e-4 --weight_decay 0.05 --input_size 160 --seed 42 --finetune vim_small_brats_160_patch4/checkpoint-999.pth --nb_classes 2 --output_dir vim_small_brats_160_patch4_cr --log_dir vim_small_brats_160_patch4_cr --use_scheduler --use_3d --use_z_score --label_type 1p19q --k_fold 5

NOTE: If you want to evaluate on IDH classification task, please remove --label_type 1p19q flag. You may want to change the number of gpus used in the scripts and the patch sizes in the scripts as well.

Visualization

To visualize the pre-trained model, you just need to run the following command:

cd vis/
python visualize.py

To visualize your own fine-tuned model, you just need to run the following command:

cd vis/
python visualize_finetuned.py

Citation

@misc{hu2024learningbraintumorrepresentation,
      title={Learning Brain Tumor Representation in 3D High-Resolution MR Images via Interpretable State Space Models}, 
      author={Qingqiao Hu and Daoan Zhang and Jiebo Luo and Zhenyu Gong and Benedikt Wiestler and Jianguo Zhang and Hongwei Bran Li},
      year={2024},
      eprint={2409.07746},
      archivePrefix={arXiv}
}

Acknowledgement

This training framework and the MoCoV3 pre-training framework is built upon MAE.

The Mamba models are built upon Mamba and Vision Mamba.

The evaluation code and MoCoV3 training code is from vit_ae_plus_plus.

Some part of the MoCoV3 finetuning code is from DeiT.

Thanks for their wonderful open-source works.