[논문 리뷰] Multi-task Learning via Adaptation to Similar Tasks for Mortality Prediction of Diverse Rare Diseases

Luchen Liu, Zequn Liu외 5명|PubMed|2020. 04. 11.

Machine Learning in Healthcare참고 문헌 30인용 수 10

한 줄 요약

이 논문은 다양한 희귀질환의 사망 예측을 위한 공유 파라미터 초기화의 빠른 적응을 가능하게 하기 위해 모델 공간 내에서 동적 태스크 유사도 측정을 사용하는 새로운 다중태스크 학습 방법인 Ada-SiT를 제안한다. 메타학습 중 태스크 간 유사성을 활용함으로써 Ada-SiT는 데이터가 적은 희귀질환 예측 성능을 크게 향상시켜, 실제 EHR 데이터셋에서 전역 단일태스크 모델 및 기존의 다중태스크 학습 기준 모델을 모두 능가한다.

ABSTRACT

The mortality prediction of diverse rare diseases using electronic health record (EHR) data is a crucial task for intelligent healthcare. However, data insufficiency and the clinical diversity of rare diseases make it hard for deep learning models to be trained. Mortality prediction for these patients with different diseases can be viewed as a multi-task learning problem with insufficient data but a large number of tasks. On the other hand, insufficient training data makes it difficult to train task-specific modules in multi-task learning models. To address the challenges of data insufficiency and task diversity, we propose an initialization-sharing multi-task learning method (Ada-SiT). Ada-Sit can learn the parameter initialization and dynamically measure the tasks' similarities, used for fast adaptation. We use Ada-SiT to train long short-term memory networks (LSTM) based prediction models on longitudinal EHR data. The experimental results demonstrate that the proposed model is effective for mortality prediction of diverse rare diseases.

연구 동기 및 목표

• 희귀질환 사망 예측에서 데이터 부족성과 임상적 다양성을 EHR 데이터를 통해 해결하기 위해.
• 각 질환 당 학습 데이터가 적어 태스크별 모듈을 효과적으로 학습하지 못하는 기존 다중태스크 학습 모델의 한계를 극복하기 위해.
• 초기화 공유에 동적 태스크 유사도를 통합함으로써 메타학습에서의 빠른 적응을 향상시키기 위해.
• 다양하고 자원이 부족한 희귀질환 태스크 간 일반화 성능을 향상시키는 모델에 종속되지 않는 방법을 개발하기 위해.

제안 방법

• Ada-SiT는 메타학습 프레임워크를 사용하여 장기 단기 기억망(LSTM) 네트워크의 공유 초기화를 학습함으로써 새로운 태스크에 대한 빠른 적응을 가능하게 한다.
• 이 방법은 정적 임상 특징에 의존하는 대신, 태스크별 모델 파라미터 간 코사인 유사도를 사용하여 모델 공간 내에서 동적 태스크 유사도 측정을 도입한다.
• 유사한 태스크는 모델 공간 내에서 학습된 거리 임계값 내에서 이웃으로 식별되며, 각 태스크 $\text{task}_i$에 대해 이웃 집합 $N_{\text{cos}}(\text{task}_i)$ 가 형성된다.
• 모델은 두 단계 과정을 통해 학습된다: 메타학습을 통한 초기화 학습 및 유사 태스크에 대한 적응, 이후 개별 태스크에 대한 피니튜닝.
• 이 방법은 모델에 종속되지 않으며, LSTM 또는 TCN와 같은 어떤 딥러닝 아키텍처와도 결합 가능하여 EHR 기반 예측에 활용할 수 있다.
• 태스크 유사도는 훈련된 모델의 $L_2$-정규화된 파라미터 벡터를 사용하여 계산되며, 임상적으로 유사한 질환 간 공유 표현을 포착할 수 있다.

실험 결과

연구 질문

• RQ1 모델 공간 내 동적 태스크 유사도가 희귀질환 사망 예측을 위한 다중태스크 학습에서의 빠른 적응에 기여하는가?
• RQ2 개별 태스크의 학습 데이터가 매우 적을 경우, 유사 태스크로부터 학습하는 것이 성능에 어떤 영향을 미치는가?
• RQ3 메타학습에서 태스크 유사도 측정에 있어 모델 공간 내 유사도가 정적 임상 특징(예: 사망률)보다 우월한가?
• RQ4 Ada-SiT가 데이터가 적은 희귀질환 태스크에서 전역 단일태스크 모델 및 표준 다중태스크 학습 기준 모델을 얼마나 뛰어나게 향상시키는가?

주요 결과

• MiniMIMIC 데이터셋에서 Ada-SiT는 LSTM에 비해 AUC를 6.9% 향상시키고 AP를 19.1% 향상시켜, 데이터가 적은 희귀질환 예측에서 뛰어난 성능을 입증한다.
• MiniMIMIC에서 Ada-SiT는 강력한 다중태스크 기준 모델인 Multi-SAnD에 비해 AUC를 8.6% 향상시키고 AP를 65.0% 향상시켜 정보 공유 능력 향상을 입증한다.
• 유사도에 대해 k-최근접 이웃(KNN)을 사용하는 Ada-SiT (KNN)는 전체 Ada-SiT 모델보다 성능이 열 劣함을 보여, 거리 기반 이웃 선택이 k-NN보다 더 효과적임을 시사한다.
• 제거 실험 결과, 모델 공간 내 태스크 유사도 측정이 정적 특징 기반 유사도(예: 사망률)에 비해 AUC를 3.0% 향상시키고 AP를 9.7% 향상시키는 것으로 확인되었다.
• t-SNE 시각화 결과, 모델 공간 내 태스크 클러스터는 사망률과 강하게 상관되며, 저사망률 및 고사망률 질환은 명확히 분리된 군집을 이룬다. 이는 유사도 측정의 임상적 타당성을 검증한다.
• Ada-SiT는 MAML(Ada-SiT -)에 비해 AUC 1.5% 향상 및 AP 6.0% 향상으로, 태스크 유사도를 통합함으로써 표준 메타학습을 뛰어넘는 적응 능력을 향상시킨다.