[논문 리뷰] MSW-Transformer: Multi-Scale Shifted Windows Transformer Networks for 12-Lead ECG Classification

Renjie Cheng, Zhemin Zhuang외 3명|arXiv (Cornell University)|2023. 06. 21.

ECG Monitoring and Analysis인용 수 7

한 줄 요약

이 논문은 다중 척도 이동 창과 학습 가능한 특징 융합을 사용하는 경량 단일층 트랜스포머 아키텍처인 MSW-Transformer를 제안한다. 이는 12-lead ECG 분류 성능을 향상시키기 위해 고안되었으며, PTBXL-2020 데이터셋에서 최신 기술 수준의 성능을 달성한다. 초급분류군에 대해 최대 77.85%의 매크로-F1 스코어를 기록하고, 샘플에 대해선 81.26%의 스코어를 기록한다. 또한 주의 메커니즘의 시각화를 통해 해석 가능성도 향상시켰다.

ABSTRACT

Automatic classification of electrocardiogram (ECG) signals plays a crucial role in the early prevention and diagnosis of cardiovascular diseases. While ECG signals can be used for the diagnosis of various diseases, their pathological characteristics exhibit minimal variations, posing a challenge to automatic classification models. Existing methods primarily utilize convolutional neural networks to extract ECG signal features for classification, which may not fully capture the pathological feature differences of different diseases. Transformer networks have advantages in feature extraction for sequence data, but the complete network is complex and relies on large-scale datasets. To address these challenges, we propose a single-layer Transformer network called Multi-Scale Shifted Windows Transformer Networks (MSW-Transformer), which uses a multi-window sliding attention mechanism at different scales to capture features in different dimensions. The self-attention is restricted to non-overlapping local windows via shifted windows, and different window scales have different receptive fields. A learnable feature fusion method is then proposed to integrate features from different windows to further enhance model performance. Furthermore, we visualize the attention mechanism of the multi-window shifted mechanism to achieve better clinical interpretation in the ECG classification task. The proposed model achieves state-of-the-art performance on five classification tasks of the PTBXL-2020 12-lead ECG dataset, which includes 5 diagnostic superclasses, 23 diagnostic subclasses, 12 rhythm classes, 17 morphology classes, and 44 diagnosis classes, with average macro-F1 scores of 77.85%, 47.57%, 66.13%, 34.60%, and 34.29%, and average sample-F1 scores of 81.26%, 68.27%, 91.32%, 50.07%, and 63.19%, respectively.

연구 동기 및 목표

• 자동 분류를 방해하는 ECG 신호 내에서 낮은 차이를 보이는 병리학적 특징 문제를 해결한다.
• CNN과 표준 트랜스포머가 미세한 ECG 변형을 포착하거나 장거리 시퀀스를 처리하는 데 한계를 보이는 문제를 극복한다.
• 복잡도와 계산 비용을 줄인 경량이며 효율적인 트랜스포머 모델을 개발하여 임상 현장에 적용 가능하게 한다.
• 임상 진단 기준과 연계된 주의 메커니즘의 시각화를 통해 모델의 해석 가능성을 향상시킨다.
• 최소한의 파라미터로 다중 클래스 12-lead ECG 분류 작업에서 최신 기술 수준의 성능를 달성한다.

제안 방법

• 다중 척도 이동 창 기반 메커니즘을 제안하여, 자기주의를 서로 겹치지 않는 국소 창에서 다양한 척도로 제한함으로써 계층적 특징을 포착한다.
• 이동 창을 활용해 창 간 정보 흐름을 가능하게 하면서도 국소 주의의 효율성과 계산 비용을 유지한다.
• 다양한 창 척도에서의 주의 출력을 융합하는 학습 가능한 특징 융합 모듈을 도입하여 표현 학습을 향상시킨다.
• 12-lead ECG 시퀀스의 장거리 의존성을 모델링하기 위해 위치 인코딩을 사용한 단일층 트랜스포머 블록을 활용한다.
• 다양한 창 척도에서의 주의 스코어를 시각화하고, 임상 진단 기준과 연계하여 모델의 해석 가능성을 확보한다.

실험 결과

연구 질문

• RQ1 다중 척도 이동 창을 갖춘 단일층 트랜스포머가 12-lead ECG 분류에서 더 깊거나 더 복잡한 모델보다 우월한 성능을 낼 수 있는가?
• RQ2 다중 척도 창 기반 메커니즘이 미세한 병리학적 차이를 보이는 ECG 신호의 특징 추출에 어떻게 기여하는가?
• RQ3 고정된 연결 또는 평균화 방식과 비교해 학습 가능한 특징 융합이 모델 성능 향상에 얼마나 기여하는가?
• RQ4 MSW-Transformer의 주의 메커니즘이 임상 진단 기준과 의미 있는 방식으로 해석될 수 있는가?
• RQ5 모델은 특히 희귀하거나 임상적으로 유사한 arrhythmia에 대해 불균형한 ECG 분류 작업에서 어떻게 성능을 내는가?

주요 결과

• MSW-Transformer는 PTBXL-2020 데이터셋에서 초급분류군 5개에 대해 최신 기술 수준의 매크로-F1 스코어 77.85%를 기록했으며, 23개의 하위분류군에선 47.57%, 12개의 리듬 클래스에선 66.13%, 17개의 형태학적 클래스에선 34.60%, 44개의 진단 클래스에선 34.29%를 기록했다.
• 샘플 매크로-F1 스코어는 각각 81.26%, 68.27%, 91.32%, 50.07%, 63.19%를 기록하여 균형 잡힌 클래스뿐 아니라 불균형 클래스에서도 뛰어난 성능를 보였다.
• 혼동 행렬을 분석한 결과, Norm(89.13% 매크로-F1), CLBBB(90.00%), SR(95.96%), STACH(92.62%)와 같은 공통 클래스에 대해 높은 정확도를 기록하여 오분류율이 낮음을 확인했다.
• 데이터 불균형에도 불구하고 공통적인 ECG 패tern을 효과적으로 구분하는 데에 강건함을 보였지만, ISCA, NST, LAFB/LPFB와 같은 희귀 레이블에 대해서는 유사성으로 인해 성능 저하가 발생했다.
• 주의 시각화 결과, 모델이 임상적으로 중요한 ECG 세그먼트—예를 들어 ST 세그먼트와 QRS 복합파—에 주목하고 있음을 확인했으며, 이는 전문가의 진단 기준과 일치하여 해석 가능성을 향상시켰다.
• 표준 트랜스포머보다 파라미터 수가 적고 추론 속도가 더 빨라 실생활 임상 적용에 적합하다.