[논문 리뷰] Multi-task UNet: Jointly Boosting Saliency Prediction and Disease Classification on Chest X-ray Images

Hongzhi Zhu, Robert Rohling외 1명|arXiv (Cornell University)|2022. 02. 15.

Visual Attention and Saliency Detection인용 수 6

한 줄 요약

이 논문은 깊이 신경망을 사용하여 흉부 X선(CXR) 영상에서 시각적 주목점 예측과 질병 분류를 동시에 수행하는 다중 작업 UNet 아키텍처를 제안한다. 공유된 특징 학습과 과적합 방지를 위한 최적화된 다중 작업 학습 기반으로, 이 모델은 두 작업 모두에서 최신 기술 수준의 성능을 달성하며, 별도로 전문화된 주목점 예측 모델이나 분류 모델보다 뛰어난 성능을 보인다.

ABSTRACT

Human visual attention has recently shown its distinct capability in boosting machine learning models. However, studies that aim to facilitate medical tasks with human visual attention are still scarce. To support the use of visual attention, this paper describes a novel deep learning model for visual saliency prediction on chest X-ray (CXR) images. To cope with data deficiency, we exploit the multi-task learning method and tackles disease classification on CXR simultaneously. For a more robust training process, we propose a further optimized multi-task learning scheme to better handle model overfitting. Experiments show our proposed deep learning model with our new learning scheme can outperform existing methods dedicated either for saliency prediction or image classification. The code used in this paper is available at https://github.com/hz-zhu/MT-UNet.

연구 동기 및 목표

• 다중 작업 학습을 활용하여 표현 학습을 향상시켜 의료 영상에서의 데이터 부족 문제를 해결한다.
• CXR 영상에서 인간의 시각적 주목점과 질병을 동시에 예측하는 통합된 딥 러닝 프레임워크를 개발한다.
• 과적합을 줄이기 위해 다중 작업 학습 기반을 최적화하여 모델의 강인성과 일반화 능력을 향상시킨다.
• 기존의 주목점 예측 또는 질병 분류에 전용으로 개발된 최신 기술 수준의 모델들을 초월하여 CXR 데이터셋에서 성능을 높인다.
• 인간의 시각적 주목 주의를 딥 러닝 파이프라인에 통합하여 의료 영상 작업에 대해 이식 가능하고 효율적인 학습을 가능하게 한다.

제안 방법

• 주목점 예측과 질병 분류 헤드 간에 인코더 특징을 공유하는 다중 작업 UNet 아키텍처를 설계한다.
• CXR 영상에서 계층적 특징을 추출하기 위해 U-Net 아키텍처를 기반으로 한 공유 인코더를 활용한다.
• 두 개의 병렬 디코더 헤드를 구현한다: 하나는 2차원 주목점 히트맵을 예측하고, 다른 하나는 다중 클래스 질병 분류를 수행한다.
• 기울기 균형 조정과 과적합 감소를 위해 동적 손실 가중치를 적용한 최적화된 다중 작업 학습 기반을 도입한다.
• 일반화 능력을 향상시키기 위해 데이터 증강 및 정규화 기법을 적용하여 제한된 의료 데이터셋에서의 성능을 개선한다.
• 주목점 예측에 이진 교차 엔트로피와 분류에 다중 클래스 교차 엔트로피를 조합한 복합 손실 함수를 사용해 모델을 종합적으로 훈련한다.

실험 결과

연구 질문

• RQ1 주목점 예측과 질병 분류를 동시에 학습시키는 것이 단일 작업 모델 대비 두 작업의 성능을 향상시키는가?
• RQ2 공유된 특징을 활용한 다중 작업 학습이 낮은 데이터 의료 영상 환경에서 표현 학습을 어떻게 향상시키는가?
• RQ3 제안된 최적화된 다중 작업 학습 기반은 작은 CXR 데이터셋에서 과적합을 어느 정도 감소시키는가?
• RQ4 주목점 예측을 통해 인간의 시각적 주목 주의를 통합하면 질병 분류 모델의 해석 가능성과 성능이 향상되는가?
• RQ5 통합 아키텍처가 CXR 영상에서 주목점 예측과 질병 분류 양쪽에서 전문화된 모델보다 뛰어난 성능을 낼 수 있는가?

주요 결과

• 제안된 다중 작업 UNet은 주목점 예측 및 질병 분류 양쪽 작업에서 최신 기술 수준의 모델을 초월한다.
• 최적화된 다중 작업 학습 기반은 특히 작은 의료 데이터셋에서 과적합을 크게 감소시켜 일반화 능력을 향상시킨다.
• 통합 학습은 더 나은 특징 학습을 이끌어내며, 단일 작업 기반 모델 대비 주목점 예측 및 분류 지표에서 성능 향상을 입증한다.
• 질병 분류에서 더 높은 AUC 스코어를 기록하고, 기존 방법들보다 인간이 애너테이션한 주목점 맵과 더 높은 상관관계를 보였다.
• 제거 실험을 통해 공유된 특징 학습과 동적 손실 가중치가 성능 향상에 핵심적인 역할을 한다는 것이 확인되었다.
• 코드와 훈련된 모델은 GitHub에 공개되어 있어 재현성과 의료 분야의 다중 작업 학습 연구를 위한 후속 연구를 가능하게 한다.