[논문 리뷰] AMI-Net+: A Novel Multi-Instance Neural Network for Medical Diagnosis from Incomplete and Imbalanced Data

Zeyuan Wang, Josiah Poon외 1명|arXiv (Cornell University)|2019. 07. 03.

Machine Learning in Healthcare참고 문헌 17인용 수 7

한 줄 요약

AMI-Net+는 포커스 손실, 게이트형 어텐션, 그리고 자기적응형 다중 인스턴스 풀링 메커니즘을 통합함으로써 부족하고 불균형한 임상 데이터로부터 의료 진단을 향상시키는 새로운 다중 인스턴스 신경망이다. 두 개의 실제 임상 데이터셋에서 기존 모델들을 능가하며, 부정확하고 편향된 증상-질병 관계를 다루는 데 뛰어난 성능을 보여준다.

ABSTRACT

In medical real-world study (RWS), how to fully utilize the fragmentary and scarce information in model training to generate the solid diagnosis results is a challenging task. In this work, we introduce a novel multi-instance neural network, AMI-Net+, to train and predict from the incomplete and extremely imbalanced data. It is more effective than the state-of-art method, AMI-Net. First, we also implement embedding, multi-head attention and gated attention-based multi-instance pooling to capture the relations of symptoms themselves and with the given disease. Besides, we propose var-ious improvements to AMI-Net, that the cross-entropy loss is replaced by focal loss and we propose a novel self-adaptive multi-instance pooling method on instance-level to obtain the bag representation. We validate the performance of AMI-Net+ on two real-world datasets, from two different medical domains. Results show that our approach outperforms other base-line models by a considerable margin.

연구 동기 및 목표

• 실제 연구에서 부족하고 극도로 불균형한 의료 데이터로부터 강건한 진단 모델을 훈련시키는 데 도전하는 것.
• 고도화된 어텐션 메커니즘과 손실 함수를 도입함으로써 AMI-Net 프레임워크를 개선하는 것.
• 의료 질환에 대한 관련성을 기반으로 동적으로 인스턴스 가중치를 학습하는 자기적응형 다중 인스턴스 풀링 전략을 개발하는 것.
• 실제 임상 데이터셋을 활용해 다양한 의료 분야에서 모델의 효과성을 검증하는 것.

제안 방법

• 훈련 중 클래스 불균형 문제를 완화하기 위해 표준 교차 엔트로피 손실을 포커스 손실로 대체한다.
• 내부 증상 간 관계와 증상-질환 관계를 모델링하기 위해 다중 헤드 및 게이트형 어텐션 메커니즘을 활용한다.
• 질환에 대한 관련성에 기반해 동적으로 인스턴스 가중치를 학습하는 자기적응형 다중 인스턴스 풀링 레이어를 도입한다.
• 어텐션 및 풀링 연산 이전에 증상 특징을 인코딩하기 위해 임bedding 레이어를 사용한다.
• 백 수준의 증상 데이터로부터 질환 분류를 최적화하기 위해 엔드 투 엔드 훈련을 적용한다.
• 각 환자는 증상(인스턴스)의 '백'으로 간주되며, 질환 레이블은 백에 할당되는 다중 인스턴스 학습(MIL) 프레임워크를 활용한다.

실험 결과

연구 질문

• RQ1 실제 환경에서 부족하고 극도로 불균형한 의료 데이터로부터 효과적으로 학습할 수 있는 다중 인스턴스 신경망이 존재하는가?
• RQ2 희귀 질환 진단의 맥락에서 표준 교차 엔트로피 손실 대비 포커스 손실이 모델 성능을 어떻게 향상시키는가?
• RQ3 MIL에서 고정된 풀링 전략 대비 자기적응형 풀링 전략이 표현 학습을 얼마나 향상시키는가?
• RQ4 어텐션 메커니즘의 통합이 진단을 위한 핵심 증상을 식별하는 데 모델의 능력을 향상시키는가?
• RQ5 AMI-Net+는 다양한 데이터 특성을 지닌 다양한 의료 분야로 일반화 가능한가?

주요 결과

• AMI-Net+는 서로 다른 분야의 두 개의 실제 임상 데이터셋에서 기준 모델, 특히 AMI-Net를 크게 능가한다.
• 포커스 손실의 사용은 소수 질환 클래스에서 성능 향상을 이끌어내어 클래스 불균형의 부정적 영향을 감소시킨다.
• 자기적응형 다중 인스턴스 풀링 메커니즘은 가장 관련성이 높은 증상에 집중함으로써 모델의 해석 가능성과 정확도를 향상시킨다.
• 특히 게이트형 어텐션은 복잡한 증상-질환 의존성을 포착하는 데 모델의 능력을 향상시킨다.
• AUC 및 F1 스코어 측면에서 기존의 MIL 기반 접근법 대비 측정 가능한 성능 향상을 달성하며 최신 기술 수준의 성능을 기록한다.