[논문 리뷰] Residual Vision Transformer (ResViT) Based Self-Supervised Learning Model for Brain Tumor Classification

Meryem Altın Karagöz, Özkan Ufuk Nalbantoğlu외 1명|arXiv (Cornell University)|2024. 11. 19.

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 스캔에서 국소적 및 전역적 특징을 포괄적으로 캡처함으로써 성능을 향상시킨다.
• 하류 분류 작업에서 ImageNet 사전학습보다 MRI 데이터로의 사전학습이 더 우수한 성능을 내는지를 입증한다.

제안 방법

• MRI 재구성이라는 생성적 자기지도 학습 사전과업을 통해 레이블이 없는 MRI 데이터로 ResViT 모델을 사전학습시킨다.
• ResViT 아키텍처를 활용해 잔차 CNN 블록을 통해 국소적 특징을 추출하고, 비전 트랜스포머 레이어를 통해 전역적 맥락을 추출한다.
• 사전학습 기간 동안 합성 MRI 영상을 생성하여 훈련 세트를 증강하고 클래스 분포를 균형 잡는다.
• 사전학습된 ResViT 모델을 하류 뇌 종양 분류 작업에서 분류기로 미세조정한다.
• T1 가중 MRI 시퀀스를 주된 입력 모odal로 사용한다. 이는 진단적 유용성이 높기 때문이다.
• 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.