[논문 리뷰] Benchmarking CNN on 3D Anatomical Brain MRI: Architectures, Data Augmentation and Deep Ensemble Learning

Benoît Dufumier, Pietro Gori외 4명|arXiv (Cornell University)|2021. 06. 02.

AI in cancer detection인용 수 10

한 줄 요약

이 논문은 대규모 다중 기관 3D 뇌 MRI 데이터셋(N=10,000)을 바탕으로 최신 3D CNN 아키텍처를 연령 예측, 성별 분류, 정신분열증 진단 등의 과제에서 벤치마킹한다. VBM 사전처리가 원시 데이터보다 성능을 크게 향상시키며, 선형 모델이 VBM 데이터에서 최신 SOTA CNN과 유사한 성능을 보이며, 딥 앙상블 학습이 정확도를 희생시키지 않고도 깊이 있는 모델의 캘리브레이션을 효과적으로 향상시킨다—반면 데이터 증강은 일관된 이점을 제공하지 못한다.

ABSTRACT

Deep Learning (DL) and specifically CNN models have become a de facto method for a wide range of vision tasks, outperforming traditional machine learning (ML) methods. Consequently, they drew a lot of attention in the neuroimaging field in particular for phenotype prediction or computer-aided diagnosis. However, most of the current studies often deal with small single-site cohorts, along with a specific pre-processing pipeline and custom CNN architectures, which make them difficult to compare to. We propose an extensive benchmark of recent state-of-the-art (SOTA) 3D CNN, evaluating also the benefits of data augmentation and deep ensemble learning, on both Voxel-Based Morphometry (VBM) pre-processing and quasi-raw images. Experiments were conducted on a large multi-site 3D brain anatomical MRI data-set comprising N=10k scans on 3 challenging tasks: age prediction, sex classification, and schizophrenia diagnosis. We found that all models provide significantly better predictions with VBM images than quasi-raw data. This finding evolved as the training set approaches 10k samples where quasi-raw data almost reach the performance of VBM. Moreover, we showed that linear models perform comparably with SOTA CNN on VBM data. We also demonstrated that DenseNet and tiny-DenseNet, a lighter version that we proposed, provide a good compromise in terms of performance in all data regime. Therefore, we suggest to employ them as the architectures by default. Critically, we also showed that current CNN are still very biased towards the acquisition site, even when trained with N=10k multi-site images. In this context, VBM pre-processing provides an efficient way to limit this site effect. Surprisingly, we did not find any clear benefit from data augmentation techniques. Finally, we proved that deep ensemble learning is well suited to re-calibrate big CNN models without sacrificing performance.

연구 동기 및 목표

• 대규모 다중 기관 3D 뇌 MRI 데이터셋에서 최신 3D CNN 아키텍처의 성능을 평가하기 위해.
• VBM 대비 준-원시 이미지 사전처리의 모델 성능에 미치는 영향을 평가하기 위해.
• 소규모 데이터 환경에서 임상 신경영상 분야의 특성에 맞는 데이터 증강 및 딥 앙상블 학습의 효과를 조사하기 위해.
• 임상 뇌 MRI 과제를 위한 딥 러닝 모델의 캘리브레이션과 불확실성 정도를 정량화하기 위해.
• 코드와 BHB-10K 데이터셋을 공개함으로써 재현 가능성을 촉진하기 위해.

제안 방법

• 연구는 공공 자료에서 확보한 N=10,000건의 스캔으로 구성된 대규모 이질적인 다중 기관 3D T1-강도 MRI 데이터셋을 사용한다.
• 다양한 최신 3D CNN 아키텍처(예: DenseNet 및 제안된 경량 버전인 tiny-DenseNet)를 연령 예측, 성별 분류, 정신분열증 진단 등의 세 가지 신경영상 과제에서 평가한다.
• 비선형 사전처리의 역할을 평가하기 위해 VBM 처리된 이미지와 최소한의 사전처리(준-원시)된 이미지에서의 성능을 비교한다.
• 데이터 증강 전략으로는 표준 변환(애핀, 블러)과 합성 MRI 아티팩트(그림자, 스팁)를 사용하며, 훈련 중 적용한다.
• 딥 앙상블 학습은 복수의 독립된 모델을 훈련하고 예측을 평균화하여 불확실성 캘리브레이션을 향상시킨다.
• 모델 캘리브레이션은 기대 캘리브레이션 오차(Expected Calibration Error, ECE)를 사용해 평가하며, 불확실성 추정을 위해 MC-Dropout와 비교한다.

실험 결과

연구 질문

• RQ1 최신 3D CNN 아키텍처는 대규모 다중 기관 3D 뇌 MRI 데이터셋에서 단순 선형 모델보다 어떻게 성능을 내는가?
• RQ2 VBM 사전처리가 준-원시 이미지 대비 다양한 과제와 데이터 환경에서 일관되게 성능 향상을 이끌어내는가?
• RQ3 표준 변환과 합성 MRI 아티팩트를 포함한 데이터 증강 전략이 소규모 데이터 환경에서 모델 성능에 어떤 영향을 미치는가?
• RQ4 딥 앙상블 학습이 임상 신경영상 과제에서 정확도를 저하시키지 않고 모델 캘리브레이션을 효과적으로 향상시킬 수 있는가?
• RQ5 대규모 다중 기관 데이터셋을 기반으로 훈련하더라도, 3D CNN은 여전히 데이터 수집 기관에 편향되어 있는가?

주요 결과

• 모든 CNN 모델이 VBM 처리된 이미지에서 준-원시 이미지보다 유의미하게 높은 성능을 보였으며, 훈련 데이터셋 크기가 10,000개에 수렴할수록 성능 격차가 점차 줄어들었다.
• 선형 모델이 VBM 데이터에서 최신 SOTA 3D CNN과 유사한 성능을 보이며, 깊이 있는 신경망이 데이터의 의미 있는 비선형성을 포착하지 못하고 있음을 시사한다.
• DenseNet과 제안된 tiny-DenseNet 아키텍처는 모든 데이터 환경과 과제에서 뛰어난 성능 대 복잡도 균형을 확보했다.
• 합성 MRI 아티팩트를 포함한 광범위한 데이터 증강에도 불구하고 일관된 성능 향상은 관찰되지 않았으며, 일부 증강 전략은 오히려 성능을 악화시켰다.
• 딥 앙상블 학습은 모든 모델과 과제에서 캘리브레이션 오차를 효과적으로 감소시켰으며, 정확도를 희생시키지 않고도 불확실성 추정을 향상시켰다—반면 MC-Dropout는 유사한 이점을 제공하지 못했다.
• N=10,000건의 다중 기관 스캔을 사용하더라도 3D CNN은 여전히 수집 기관에 편향되어 있었으며, VBM 사전처리가 이 편향을 완화하는 데 도움이 되었다.