[Paper Review] Medical Transformer: Universal Brain Encoder for 3D MRI Analysis

Eunji Jun, Seungwoo Jeong+2 more|arXiv (Cornell University)|Apr 28, 2021

Advanced Neural Network Applications24 references32 citations

TL;DR

A self-supervised, multi-view Transformer-based backbone (Medical Transformer) pre-trained on large-scale 3D brain MRI data to enable parameter-efficient transfer learning for brain disease diagnosis, brain age prediction, and brain tumor segmentation.

ABSTRACT

Transfer learning has gained attention in medical image analysis due to limited annotated 3D medical datasets for training data-driven deep learning models in the real world. Existing 3D-based methods have transferred the pre-trained models to downstream tasks, which achieved promising results with only a small number of training samples. However, they demand a massive amount of parameters to train the model for 3D medical imaging. In this work, we propose a novel transfer learning framework, called Medical Transformer, that effectively models 3D volumetric images in the form of a sequence of 2D image slices. To make a high-level representation in 3D-form empowering spatial relations better, we take a multi-view approach that leverages plenty of information from the three planes of 3D volume, while providing parameter-efficient training. For building a source model generally applicable to various tasks, we pre-train the model in a self-supervised learning manner for masked encoding vector prediction as a proxy task, using a large-scale normal, healthy brain magnetic resonance imaging (MRI) dataset. Our pre-trained model is evaluated on three downstream tasks: (i) brain disease diagnosis, (ii) brain age prediction, and (iii) brain tumor segmentation, which are actively studied in brain MRI research. The experimental results show that our Medical Transformer outperforms the state-of-the-art transfer learning methods, efficiently reducing the number of parameters up to about 92% for classification and

Motivation & Objective

• Motivate transfer learning in 3D medical imaging due to limited annotated data.
• Propose a universal, parameter-efficient backbone that models 3D MRI as multi-view 2D slices (sagittal, coronal, axial).
• Pre-train via self-supervised masked encoding vector prediction on large healthy brain MRI data.
• Demonstrate transfer learning to brain disease diagnosis, brain age prediction, and brain tumor segmentation.
• Show reduced parameter count while achieving or surpassing state-of-the-art performance.

Proposed method

• Multi-view 3D to 2D slice representation from sagittal, coronal, and axial planes.
• Backbone with separate convolutional encoders per plane feeding a transformer to model inter-slice dependencies.
• Self-supervised pre-training with exemplar-based loss and masked encoding vector prediction (BERT-like) across slice tokens.
• Fine-tuning with single-scale (classification/regression) or multi-scale (segmentation) predictions.
• Comparison against state-of-the-art 3D transfer-learning and self-supervised methods on ADNI BraTS and related datasets.

Experimental results

Research questions

• RQ1 Can a self-supervised, multi-view backbone pre-trained on healthy brain MRI serve as a universal encoder for downstream 3D MRI tasks?
• RQ2 Does the Medical Transformer reduce model parameters while maintaining or improving performance compared with existing 3D transfer-learning methods?
• RQ3 How do brain disease diagnosis, brain age prediction, and brain tumor segmentation benefit from a shared universal encoder?
• RQ4 Is a multi-scale fine-tuning strategy advantageous for segmentation versus classification/regression tasks?

Key findings

• Outperforms state-of-the-art 3D transfer-learning methods across the three tasks.
• Achieves up to ~92% parameter reduction for classification and regression and ~97% for segmentation compared to competitors.
• Demonstrates strong performance with partial training data, preserving accuracy with fewer labeled samples.
• Brain disease diagnosis and brain age prediction show notable gains using the universal backbone.
• Brain tumor segmentation achieves competitive Dice scores, with strongest improvement on the challenging ET region.