[논문 리뷰] MambaCapsule: Towards Transparent Cardiac Disease Diagnosis with Electrocardiography Using Mamba Capsule Network

Yinlong Xu, Xiaoqiang Liu외 7명|arXiv (Cornell University)|2024. 07. 30.

ECG Monitoring and AnalysisMedicine인용 수 3

한 줄 요약

MambaCapsule은 장기 시퀀스 특징 추출을 위한 Mamba와 계층적 특징 표현 및 재구성에 적합한 캡슐 네트워크를 융합하여 심 electrocardiogram (ECG) 기반 심장 부정맥 분류를 위한 새로운 해석 가능한 딥러닝 프레임워크를 제안한다. MIT-BIH 및 PTB 데이터셋에서 각각 99.54%와 99.59%의 정확도를 달성하면서도, 학습된 특징과 그 관계를 드러내는 신호 재구성 방식을 통해 투명한 의사결정을 가능하게 한다.

ABSTRACT

Cardiac arrhythmia, a condition characterized by irregular heartbeats, often serves as an early indication of various heart ailments. With the advent of deep learning, numerous innovative models have been introduced for diagnosing arrhythmias using Electrocardiogram (ECG) signals. However, recent studies solely focus on the performance of models, neglecting the interpretation of their results. This leads to a considerable lack of transparency, posing a significant risk in the actual diagnostic process. To solve this problem, this paper introduces MambaCapsule, a deep neural networks for ECG arrhythmias classification, which increases the explainability of the model while enhancing the accuracy.Our model utilizes Mamba for feature extraction and Capsule networks for prediction, providing not only a confidence score but also signal features. Akin to the processing mechanism of human brain, the model learns signal features and their relationship between them by reconstructing ECG signals in the predicted selection. The model evaluation was conducted on MIT-BIH and PTB dataset, following the AAMI standard. MambaCapsule has achieved a total accuracy of 99.54% and 99.59% on the test sets respectively. These results demonstrate the promising performance of under the standard test protocol.

연구 동기 및 목표

• 딥러닝 기반 ECG 부정맥 진단의 해석 가능성 부족 문제를 해결하여 임상적 신뢰도와 신뢰성을 향상시키기 위해.
• 높은 정확도로 부정맥을 분류할 뿐 아니라, 신호 재구성을 통해 예측의 근거를 설명할 수 있는 모델을 개발하기 위해.
• 캡슐 네트워크의 뇌 유사한 불변성과 동적 라우팅 기능을 활용하여 ECG 신호 내 특징 표현 및 관계 모델링을 향상시키기 위해.
• Mamba 아키텍처를 통합하여 ECG 데이터에서 장기적인 시간 패턴을 효율적으로 모델링하고 장기적 시퀀스에 대한 특징 추출 능력을 향상시키기 위해.

제안 방법

• 모델은 원시 ECG 신호에서 장기적이고 계층적인 표현을 추출하기 위해 Mamba 기반 인코더를 사용하며, 상태 공간 모델을 활용해 계산 효율성과 장기 시퀀스 모델링을 확보한다.
• 인코딩된 특징을 처리하기 위해 캡슐 네트워크 디코더를 활용하며, 클래스 확률과 신호 구성 요소 간의 공간적 관계를 포함한 캡슐 벡터를 출력한다.
• 입력 ECG 신호를 예측된 캡슐 상태에서 재구성하는 재구성 헤드를 포함하여, 모델이 주로 집중하는 특징을 시각화함으로써 설명 가능성 기능을 제공한다.
• 재구성 과정을 통해 입력의 변형(예: 펌퍼런스)에 대한 P파 및 R파 형태의 변화를 분석함으로써 특징 학습의 타당성과 모델 의사결정의 해석 가능성을 검증한다.
• 인코더-디코더가 종합적으로 끝내기로 훈련되는 공동 훈련 방식을 채택하며, 재구성 손실이 특징 학습을 이끄는 데 기여한다.
• 캡슐 네트워크의 라우팅 메커니즘은 관련 특징에 동적으로 주목함으로써 예측의 강건성과 해석 가능성 향상에 기여한다.

실험 결과

연구 질문

• RQ1 Mamba 기반 인코더가 CNN, Transformer, ResNet 아키텍처 대비 장기 시퀀스 ECG 표현 학습에서 성능 향상을 이룰 수 있는가?
• RQ2 캡슐 네트워크의 동적 라우팅 메커니즘이 ECG 신호 내 특징 계층과 관계 모델링에 얼마나 기여하는가?
• RQ3 캡슐 상태에서의 신호 재구성은 모델의 내부 지식과 의사결정 이론을 얼마나 효과적으로 드러내는가?
• RQ4 Mamba와 캡슐 네트워크의 통합이 분류 정확도를 향상시키면서도 모델의 투명성을 유지하거나 향상시키는가?
• RQ5 입력의 변형(예: 이동)이 재구성된 ECG 신호에 미치는 영향은 무엇이며, 이는 모델의 강건성과 특징 이해 수준을 어떻게 반영하는가?

주요 결과

• MambaCapsule은 표준 AAMI 평가 프로토콜을 기준으로 MIT-BIH 부정맥 데이터셋에서 테스트 정확도 99.54%를, PTB 데이터셋에서 99.59%를 기록하여 최신 기술 수준의 성능을 입증하였다.
• 입력의 변형에 대해 캡슐 출력에서 재구성된 신호가 높은 일관성을 유지하여, 모델이 안정적이고 의미 있는 특징 및 그 상호의존성을 학습하고 있음을 시사한다.
• 정상(N) 및 심방성 premature beat(S) 레이블에 대한 재구성 신호는 S 클래스 신호에서 P파가 존재하지 않는 것을 명확히 보여주어 모델의 분류 근거를 시각적으로 설명한다.
• Mamba 인코더는 F1 점수와 재구성 정밀도 측면에서 CNN, ResNet, Transformer 기반 인코더보다 뛰어난 성능을 보이며 정확도 및 특징 표현 능력에서 슈퍼리어한 성능을 보였다.
• 캡슐 기반 디코더는 더 긴 훈련 시간이 걸리지만, 특징 간 관계를 유지할 수 있는 능력 덕분에 표준 MLP 디코더보다 더 해석 가능하고 강건한 예측을 제공하였다.
• 모델의 재구성 능력은 단순히 개별 파형 형태(P, R)뿐 아니라 그 상대적 구성까지도 포괄적으로 포착하고 있음을 드러내어, 모델의 해석 가능성 주장에 대한 지원을 제공한다.