[論文レビュー] 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

ひとこと要約

本稿では、MRI再構成を事前学習タスクとして用いる自己教師あり学習モデルであるResidual Vision Transformer（ResViT）を用いた脳腫瘍分類手法を提案する。CNNとViTを組み合わせて局所的およびグローバルな特徴抽出を実現し、合成MRIデータを用いた増幅処理と、ImageNetではなくMRIデータを用いた事前学習を行うことで、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データを用いた事前学習の利点を実証すること。

提案手法

• 生成的自己教師あり学習の事前学習タスクとしてのMRI再構成を用い、ラベルなしMRIデータ上でResViTモデルを事前学習する。
• ResViTアーキテクチャを活用し、残差CNNブロックで局所的特徴を抽出するとともに、ビジョントランスフォーマーレイヤーでグローバルなコンテキストを抽出する。
• 事前学習中に合成MRI画像を生成し、トレーニングセットの拡張とクラス分布のバランスを図る。
• 下流の脳腫瘍分類タスクにおいて、事前学習済みResViTモデルを微調整して分類器として利用する。
• 診断的有用性が高いため、主にT1強調MRIシーケンスを入力モodalitiyとして用いる。
• ImageNetでの事前学習とMRIデータでの事前学習を比較し、ドメイン特化型の転移学習の利点を評価する。

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

実験結果

リサーチクエスチョン

• RQ1 MRIデータを用いた自己教師あり事前学習戦略は、ImageNetでの事前学習と比較して、下流の脳腫瘍分類性能を向上させることができるか？
• RQ2 ハイブリッドResViTアーキテクチャにおけるCNNとViTの統合は、脳腫瘍分類のための特徴学習を強化するか？
• RQ3 合成MRI画像生成は、小規模で不均衡なデータセットにおいて、モデルの一般化性能と精度をどの程度向上させるか？
• RQ4 T1、T2、FLAIRなど、どのMRIシーケンスが、提案モデルを用いた際の分類精度を最も高めるか？
• RQ5 提案モデルは、最小限の微調整で他の脳腫瘍データセットへ効果的に転送可能か？

主な発見

• 提案されたResViTベースのモデルは、T1強調シーケンスを用いてBraTS 2023データセットで90.56%の精度を達成し、比較対象の全モデルを上回った。
• Figshareデータセットでは98.53%の精度に達し、LCDEiT（98.11%）やResNet-50（97.48%）といった先行研究を上回った。
• Kaggle脳腫瘍データセットでは98.47%の精度を達成し、多様なデータソースにわたる強力な一般化性能を示した。
• ImageNetでの事前学習よりもMRIデータでの事前学習が顕著に優れた性能を示し、ドメイン特化型自己教師あり学習の利点を裏付けた。
• モデルの性能は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.