[논문 리뷰] Using Multitask Learning to Improve 12-Lead Electrocardiogram Classification

J. Weston Hughes, Taylor Sittler외 4명|arXiv (Cornell University)|2018. 12. 02.

ECG Monitoring and Analysis인용 수 4

한 줄 요약

이 논문은 공유 표현 학습을 통해 희귀 부정맥의 성능을 크게 향상시키는 다중 작업 컨볼루션 신경망(CNN)을 제안한다. 12 lead ECG 분류에서 다수의 심장 질환을 동시에 학습함으로써, 특히 희귀 부정맥에 대해 성능을 향상시킨다. 24개의 이진 레이블을 가진 41,522개의 ECG 샘플 데이터셋을 사용하여, 다중 헤드 모델은 희귀 클래스인 2차도 AV 블록(Mobitz I)에서 F1 점수를 0.58에서 0.73으로 향상시키며, 더 높은 F1 점수를 달성한다. 이는 관련된 일반적인 클래스로부터의 감독이 희귀 클래스의 일반화 능력을 향상시킨다는 것을 보여준다.

ABSTRACT

We develop a multi-task convolutional neural network (CNN) to classify multiple diagnoses from 12-lead electrocardiograms (ECGs) using a dataset comprised of over 40,000 ECGs, with labels derived from cardiologist clinical interpretations. Since many clinically important classes can occur in low frequencies, approaches are needed to improve performance on rare classes. We compare the performance of several single-class classifiers on rare classes to a multi-headed classifier across all available classes. We demonstrate that the addition of common classes can significantly improve CNN performance on rarer classes when compared to a model trained on the rarer class in isolation. Using this method, we develop a model with high performance as measured by F1 score on multiple clinically relevant classes compared against the gold-standard cardiologist interpretation.

연구 동기 및 목표

• 희귀 심장 부정맥의 12-lead ECG 분류 성능이 부족한 이유로 인한 훈련 데이터 부족 문제를 해결하기 위해.
• 다중 작업 학습이 관련된 더 흔한 클래스를 활용하여 희귀 ECG 진단의 분류 정확도를 향상시킬 수 있는지 조사하기 위해.
• 다양한 임상적으로 관련성이 있는 진단에 대해 일반화되는 고성능의 다중 헤드 딥 러닝 모델을 개발하기 위해.
• 희귀 클래스에 대해 성능 향상을 가져올 수 있는 단일 관련 일반 클래스를 추가하는 것이, 희귀 상태에 대한 대규모 데이터 수집의 필요성을 줄일 수 있는지 평가하기 위해.

제안 방법

• 12-lead ECG를 250Hz로 샘플링하여 10초 동안의 입력 텐서(12×2500)를 생성하고, 34층의 컨볼루션 신경망(CNN)을 훈련한다.
• 모든 24개의 이진 분류 작업을 동시에 학습할 수 있도록, 각 진단 클래스별로 하나의 출력 헤드를 가진 다중 헤드 아키텍처를 사용한다.
• 각 헤드는 이진 교차 엔트로피 손실을 사용하여 특정 심장 질환의 유무를 예측하며, 공유된 컨볼루션 레이어를 통해 역전파된 기울기를 공유한다.
• 클래스 불균형 문제를 완화하기 위해 데이터셋을 재샘플링하여 클래스 빈도를 균형 잡았으며, 소수 클래스는 4,000개의 예시로 업샘플링하였다.
• 모델 성능 평가에는 모든 클래스에서의 F1 점수를 사용하였으며, 단일 헤드 모델과 다중 헤드 모델 간의 성능을 다양한 클래스 조합에서 훈련한 결과로 비교하였다.
• 예를 들어 AV 블록, 빈혈성 빈맥 등과 같은 계층적이고 상관관계가 있는 클래스 간의 관계를 활용하여, 공유된 특징 레이어에서의 표현 학습을 향상시켰다.

실험 결과

연구 질문

• RQ1 희귀 12-lead ECG 진단에 대해, 더 흔한 클래스와의 표현 공유를 통해 다중 작업 학습이 딥 러닝 모델의 F1 점수를 향상시킬 수 있는가?
• RQ2 희귀 진단인 2차도 AV 블록(Mobitz I)에 대해, 단일 관련 일반 클래스(예: 1차도 AV 블록)를 포함시키는 것이, 다수의 관련 없는 클래스를 포함시키는 것과 비슷한 성능 향상을 가져오는가?
• RQ3 희귀 클래스에 대해 단일 헤드 모델과 다중 헤드 모델의 성능를 비교했을 때, 특히 훈련 예제가 1,000개 미만인 희귀 클래스에 대해 어떤 차이가 있는가?
• RQ4 다중 작업 학습이 이질적이고 저주파수의 진단 레이블을 가진 ECG에 적용되었을 때, CNN의 깊은 레이어에서 특징 학습이 얼마나 향상되는가?

주요 결과

• 다중 헤드 모델은 모든 클래스에서 단일 헤드 모델을 능가했으며, 특히 희귀 클래스에서 가장 뚜렷한 F1 점수 향상을 보였다. 예를 들어 완전한 심장 블록(F1: 0.62 vs. 0.09)과 심방조직재입 빈맥(F1: 0.49 vs. 0.39)에서 두드러진 개선이 있었다.
• 626개의 예시만 있는 희귀 클래스인 2차도 AV 블록(Mobitz I)에 대해, 다중 헤드 모델은 F1 점수 0.73을 기록하여 단일 헤드 모델의 0.58보다 뚜렷한 향상을 보였다.
• 2차도 AV 블록(Mobitz I)의 훈련에 단일 관련 일반 클래스(1차도 AV 블록)를 추가하는 것만으로도, 모든 다른 클래스를 추가한 것과 동일한 성능 향상을 얻었으며, 이는 관련 감독의 효과성이 매우 높다는 것을 시사한다.
• 일반 클래스인 심방 리듬(F1: 0.91)과 fibrillation(F1: 0.85)은 단일 및 다중 헤드 모델 모두에서 높은 성능를 기록했지만, 다중 헤드 모델 역시 약간의 성능 향상을 보였다.
• 공유된 표현 학습이 여러 작업 간에 이루어지면서, 특히 깊은 레이어에서의 특징 추출 능력이 향상되어, 자원이 제한된 클래스에 대한 일반화 능력이 향상됨을 입증하였다.
• 결과적으로 다중 작업 학습은 희귀 ECG 진단에 대해 데이터 효율성 장벽을 줄일 수 있으며, 희귀 클래스당 1,000개 미만의 예시로도 고성능 모델을 개발할 수 있음을 시사한다.