[논문 리뷰] Multi-Task Attention-Based Semi-Supervised Learning for Medical Image Segmentation

Shuai Chen, Gerda Bortsova외 3명|arXiv (Cornell University)|2019. 07. 29.

Advanced Neural Network Applications참고 문헌 13인용 수 13

한 줄 요약

이 논문은 다중 작업 어텐션 기반 준지도 학습(MASSL)을 제안하며, 전경과 배경 재구성에 대해 어텐션 메커니즘을 사용하여 분리하는 방식으로 분할 및 재구성 작업을 동시에 최적화한다. 제한된 레이블이 부여된 데이터로 훈련되었을 때, 이 방법은 완전히 지도 학습된 모델과 사전 훈련된 CNN보다 우수한 성능을 보이며, 교차 훈련과 어텐션 메커니즘이 특징의 분류 가능성을 높이고 뇌 종양 및 백질 비정상 볼록부위 데이터셋에서 분할 정확도를 향상시킨다.

ABSTRACT

We propose a novel semi-supervised image segmentation method that simultaneously optimizes a supervised segmentation and an unsupervised reconstruction objectives. The reconstruction objective uses an attention mechanism that separates the reconstruction of image areas corresponding to different classes. The proposed approach was evaluated on two applications: brain tumor and white matter hyperintensities segmentation. Our method, trained on unlabeled and a small number of labeled images, outperformed supervised CNNs trained with the same number of images and CNNs pre-trained on unlabeled data. In ablation experiments, we observed that the proposed attention mechanism substantially improves segmentation performance. We explore two multi-task training strategies: joint training and alternating training. Alternating training requires fewer hyperparameters and achieves a better, more stable performance than joint training. Finally, we analyze the features learned by different methods and find that the attention mechanism helps to learn more discriminative features in the deeper layers of encoders.

연구 동기 및 목표

• 제한된 정규화된 의료 영상 문제를 해결하기 위해, 준지도 학습을 통해 대량의 비정상 데이터를 활용한다.
• 스킵 커넥션으로 인해 클래스별 정보를 유지하지 못하는 기존 U-Net 아키텍처의 표준 오토에인코더 기반 재구성 방법의 한계를 극복한다.
• 어텐션 메커니즘을 통한 분할 인식 재구성 통합을 통해 딥 인코더의 특징 학습을 향상시킨다.
• 두 가지 다중 작업 훈련 전략—일괄 훈련과 교차 훈련—이 분할 성능에 미치는 영향을 평가한다.
• 어텐션 가이드된 재구성이 표준 재구성 방식보다 더 분류 가능한 특징을 깊은 네트워크 레이어에서 유도한다는 것을 입증한다.

제안 방법

• 공유 인코더가 분할 헤드(스킵 커넥션 포함)와 재구성 헤드(스킵 커넥션 없음)에 연결되는 다중 작업 프레임워크를 제안하며, U-Net 유사 아키텍처를 형성한다.
• 부드러운 분할 예측에서 전경과 배경 마스크를 생성하는 어텐션 메커니즘을 도입하여, 클래스 인식 재구성을 가능하게 한다.
• 어텐션 기반 마스크를 사용해 클래스별 영역을 분리하고, 가중치가 부여된 MSE를 독립적으로 적용하는 재구성 손실(L2)을 정의한다.
• 두 가지 훈련 전략을 구현한다: 일괄 훈련(양측 작업을 동시에 최적화)과 교차 훈련(작업별 학습률을 사용해 순차적으로 최적화).
• 양측 작업에 대해 공유 인코더를 사용하여, 재구성 목적이 인코더의 더 분류 가능한 특징 학습을 이끌도록 한다. 특히 깊은 레이어에서 효과가 뚜렷하다.
• 소수의 레이블이 부여된 이미지와 대량의 비정상 데이터를 사용하여, 두 가지 의료 영상 작업—뇌 종양(BraTS18) 및 백질 비정상 볼록부위(WMH17) 분할—에 적용한다.

실험 결과

연구 질문

• RQ1 표준 오토에인코더 기반 재구성 방식과 비교해, 어텐션 기반 재구성 목적이 준지도 학습 의료 영상 분할 성능을 향상시키는가?
• RQ2 어텐션을 통한 분할 및 재구성 통합이 표준 다중 작업 학습보다 인코더에서 더 분류 가능한 특징을 생성하는가?
• RQ3 일괄 훈련과 교차 훈련 전략이 다중 작업 준지도 학습 분할에서 성능과 안정성에 미치는 영향은 어떠한가?
• RQ4 소수의 레이블이 부여된 이미지만 존재할 경우, 제안된 방법이 완전히 지도 학습된 CNN 및 사전 훈련 기반 기준 모델보다 얼마나 뛰어나게 성능을 발휘하는가?
• RQ5 어텐션 메커니즘이 깊은 인코더 레이어에서 학습된 특징의 품질에 어떤 영향을 미치는가? 선형 분류 성능 측정 기준으로 평가한다.

주요 결과

• 교차 훈련을 사용한 MASSL은 BraTS18(50개 레이블, 70개 비레이블)에서 0.7111 ± 0.02의 Dice 스코어를 기록하여, 완전히 지도 학습된 CNN 기준 모델(0.6670 ± 0.03)과 MSSL 변종(0.6646 ± 0.03)보다 유의미하게 뛰어나다.
• WMH17에서 MASSL은 교차 훈련을 통해 0.7204 ± 0.02의 Dice 스코어를 기록하여, 완전히 지도 학습된 CNN 기준 모델(0.6806 ± 0.03)과 MSSL 기준 모델(0.6880 ± 0.02)을 초월했다.
• 어텐션 메커니즘이 성능 향상에 크게 기여했다: 교차 훈련을 사용한 MASSL은 BraTS18에서 0.7111 ± 0.02의 성능을 기록했고, 어텐션 없이 사용한 MSSL 기준 모델은 0.6670 ± 0.03이었다.
• 교차 훈련은 안정성과 성능 면에서 일괄 훈련을 능가했으며, 손실 가중치 계수 γ의 초모수 조정 없이도 모든 데이터셋에서 더 높은 Dice 스코어를 기록했다.
• 특징 분석 결과 어텐션 메커니즘이 깊은 인코더 레이어의 분류 가능성을 향상시켰다: BraTS18에서 MASSL은 레이어 5(가장 깊은 레이어)에서 R² = 0.535 ± 0.03을 기록했고, 표준 CNN은 0.486 ± 0.04를 기록했다.
• 제거 분석 결과 어텐션 메커니즘이 성능 향상의 핵심 요인임을 확인했으며, 동일한 훈련 전략과 손실 가중치를 사용하더라도 MSSL(어 attention 없음)은 MASSL보다 성능이 열 劣하다.