[论文解读] Residual Vision Transformer (ResViT) Based Self-Supervised Learning Model for Brain Tumor Classification

Meryem Altın Karagöz, Özkan Ufuk Nalbantoğlu+1 more|arXiv (Cornell University)|Nov 19, 2024

Brain Tumor Detection and ClassificationNeuroscience被引用 3

一句话总结

该论文提出了一种基于残差视觉Transformer（ResViT）的自监督学习模型，用于脑肿瘤分类。该模型以MRI生成作为先验任务进行预训练，并在分类任务上进行微调。通过结合CNN与ViT以提取局部和全局特征，利用合成MRI数据进行数据增强，并在MRI数据而非ImageNet上进行预训练，该模型在BraTS数据集上达到90.56%的准确率，在Figshare数据集上达到98.53%，在Kaggle数据集上达到98.47%，优于当前最先进模型。

ABSTRACT

Deep learning has proven very promising for interpreting MRI in brain tumor diagnosis. However, deep learning models suffer from a scarcity of brain MRI datasets for effective training. Self-supervised learning (SSL) models provide data-efficient and remarkable solutions to limited dataset problems. Therefore, this paper introduces a generative SSL model for brain tumor classification in two stages. The first stage is designed to pre-train a Residual Vision Transformer (ResViT) model for MRI synthesis as a pretext task. The second stage includes fine-tuning a ResViT-based classifier model as a downstream task. Accordingly, we aim to leverage local features via CNN and global features via ViT, employing a hybrid CNN-transformer architecture for ResViT in pretext and downstream tasks. Moreover, synthetic MRI images are utilized to balance the training set. The proposed model performs on public BraTs 2023, Figshare, and Kaggle datasets. Furthermore, we compare the proposed model with various deep learning models, including A-UNet, ResNet-9, pix2pix, pGAN for MRI synthesis, and ConvNeXtTiny, ResNet101, DenseNet12, Residual CNN, ViT for classification. According to the results, the proposed model pretraining on the MRI dataset is superior compared to the pretraining on the ImageNet dataset. Overall, the proposed model attains the highest accuracy, achieving 90.56% on the BraTs dataset with T1 sequence, 98.53% on the Figshare, and 98.47% on the Kaggle brain tumor datasets. As a result, the proposed model demonstrates a robust, effective, and successful approach to handling insufficient dataset challenges in MRI analysis by incorporating SSL, fine-tuning, data augmentation, and combining CNN and ViT.

研究动机与目标

• 解决深度学习方法在脑肿瘤分类中因标注脑MRI数据集有限带来的挑战。
• 克服因医学影像中小样本、类别不平衡数据集导致的过拟合和泛化能力差的问题。
• 通过自监督学习开发一种数据高效、鲁棒且可泛化的脑肿瘤分类模型。
• 通过结合CNN与ViT架构，提升MRI扫描中局部与全局特征的捕捉能力。
• 证明在下游分类任务中，使用MRI数据预训练优于ImageNet预训练。

提出的方法

• 使用生成式自监督学习先验任务（MRI重建）在无标签MRI数据上预训练ResViT模型。
• 利用ResViT架构，通过残差CNN模块联合提取局部特征，通过视觉Transformer层捕捉全局上下文信息。
• 在预训练过程中生成合成MRI图像，以增强训练集并平衡类别分布。
• 将预训练的ResViT模型微调为下游脑肿瘤分类任务的分类器。
• 以T1加权MRI序列作为主要输入模态，因其具有较高的诊断价值。
• 对比在ImageNet上预训练与在MRI数据上预训练的性能，以评估领域特定迁移学习的优势。

Figure 1: The proposed Residual Vision Transformer (ResViT) based Self-supervised learning model.

实验结果

研究问题

• RQ1 与ImageNet预训练相比，基于MRI数据的自监督预训练策略是否能提升下游脑肿瘤分类性能？
• RQ2 在混合式ResViT架构中结合CNN与ViT是否能增强脑肿瘤分类的特征学习能力？
• RQ3 通过先验任务生成的合成MRI数据在多大程度上能提升小样本、类别不平衡数据集上的模型泛化能力与准确率？
• RQ4 在使用所提模型时，哪种MRI序列（T1、T2、FLAIR等）能获得最高的分类准确率？
• RQ5 该模型能否在少量微调下有效迁移到其他脑肿瘤数据集？

主要发现

• 所提出的基于ResViT的模型在BraTS 2023数据集上使用T1加权序列实现了90.56%的准确率，优于所有对比模型。
• 在Figshare数据集上，模型准确率达到98.53%，超过此前最先进方法如LCDEiT（98.11%）和ResNet-50（97.48%）。
• 在Kaggle脑肿瘤数据集上，模型准确率达到98.47%，表明其在不同数据源间具有强大的泛化能力。
• 在MRI数据上进行预训练显著优于ImageNet预训练，凸显了领域特定自监督学习的优势。
• 模型在T1加权序列上的表现最为稳健，表明该序列对肿瘤类型分类具有高度信息量。
• 通过先验任务整合合成MRI数据，提升了模型的鲁棒性，并缓解了小样本数据集上的过拟合问题。

Figure 2: The several axial slices of MRI include labels and information of tumor region with coverage that all belong to the same case. The enhancing tumor (ET), the tumor core (TC), and the whole tumor (WT) are shown yellow, blue, and green, respectively.