[논문 리뷰] Brain Tumor Image Retrieval via Multitask Learning

Maxim Pisov, Gleb Makarchuk외 4명|arXiv (Cornell University)|2018. 10. 22.

Brain Tumor Detection and Classification참고 문헌 5인용 수 6

한 줄 요약

이 논문은 뇌 종양 영상 검색을 위한 다중 작업 딥 러닝 프레임워크를 제안하며, 공유된 CNN 기반 아키텍처를 사용하여 MRI 스캔에서 종양 유형, 국소화, 분할을 동시에 학습한다. 이질적인 레이블로 훈련함으로써 모델은 단일 작업 접근 방식을 능가하는 풍부하고 다중 정보를 담은 표현을 생성하며, 경계 상자 정확도가 떨어지는 상황에서도 표현 간 거리와 종양 유사성 간 강한 상관관계를 보인다.

ABSTRACT

Classification-based image retrieval systems are built by training convolutional neural networks (CNNs) on a relevant classification problem and using the distance in the resulting feature space as a similarity metric. However, in practical applications, it is often desirable to have representations which take into account several aspects of the data (e.g., brain tumor type and its localization). In our work, we extend the classification-based approach with multitask learning: we train a CNN on brain MRI scans with heterogeneous labels and implement a corresponding tumor image retrieval system. We validate our approach on brain tumor data which contains information about tumor types, shapes and localization. We show that our method allows us to build representations that contain more relevant information about tumors than single-task classification-based approaches.

연구 동기 및 목표

• 다양한 종양 특성(유형, 형태, 국소화)을 포괄하는 의료 영상 검색 시스템을 개발하여 유사성 검색 성능을 향상시키는 것.
• 뇌 MRI에서 복잡한 종양 유사성을 포괄하지 못하는 단일 작업 분류의 한계를 해결하는 것.
• 다중 작업 학습이 종양 검색을 위한 더 정보가 풍부한 표현을 생성하는지 검증하는 것.
• 임상 환경에서 흔히 발생하는 실제적인 경계 상자 정확도 저하 상황에서도 시스템의 강건성을 평가하는 것.

제안 방법

• 3D U-Net 유사 CNN 기반 아키텍처에 잔차 블록을 적용하여 뇌 MRI 스캔에서 공유 특징을 추출한다.
• 기반 네트워크에 작업별 전용 헤드를 부착: 하나는 분할(전역)을 위한 것이며, 나머지는 종양 유형 및 국소화(각 종양 영역에 대해 국소)를 위한 것이다.
• 모든 작업을 동시에 최적화하는 병합 손실 함수를 사용하여 네트워크를 엔드 투 엔드로 훈련시킨다.
• 최종 특징 맵에서 종양 표현을 추출하고, k-NN를 통한 유사성 검색에 임베딩으로 사용한다.
• 정확도가 떨어지는 입력에 대한 강건성을 시험하기 위해 추론 중 경계 상자 왜곡(확대 및 이동)을 적용한다.
• t-SNE 시각화를 통해 종양 유형 및 크기별 표현의 군집화 정도를 정성적으로 평가한다.

실험 결과

연구 질문

• RQ1 다중 작업 학습이 단일 작업 분류 대비 종양 표현의 품질을 향상시킬 수 있는가?
• RQ2 학습된 임베딩이 종양 유형, 크기, 국소화와 같은 임상적으로 중요한 특성과 얼마나 잘 상관되는가?
• RQ3 실제 임상 워크플로우에서 흔히 발생하는 경계 상자 정확도 저하 상황에서도 검색 시스템의 성능이 유지되는가?
• RQ4 모델이 단일 통합 표현에서 이질적인 메타데이터(예: 종양 유형 및 해부학적 영역)를 동시에 포괄할 수 있는가?

주요 결과

• 다중 작업 모델은 모든 작업에서 균형 잡힌 성능을 달성하였으며, 표현 품질과 일반화 능력 측면에서 단일 작업 기반 모델을 능가하였다.
• 학습된 임베딩을 사용한 K-Nearest Neighbors는 K=5일 때 종양 유형 분류 정확도 85.7%와 선형 크기 회귀의 RMSE 12.3 mm를 기록하였다.
• 중요한 경계 상자 왜곡 상황에서도 성능이 거의 안정되었으며, 정확도 및 RMSE는 소폭 감소에 그쳐 입력 노이즈에 대한 강건성을 입증하였다.
• t-SNE 시각화 결과 종양 임베딩이 종양 유형과 크기별로 명확하게 군집되어 있음을 확인하였으며, 좌측 및 우측 소뇌에서의 색소마 유형 종양에 대해 별도의 군집이 형성되었다.
• 고차원 표현(최대 96채널)은 성능 향상에 기여하였지만, 64채널이 정확도와 학습 시간 간의 최적 균형을 제공하였다.
• 검색 예시를 통해 시스템은 다른 환자로부터의 종양이라도 형태와 국소화가 유사한 종양을 성공적으로 검색하는 것으로 나타났다.