[논문 리뷰] Masked Transformer for Electrocardiogram Classification

Zhou Ya, Xiaolin Diao외 4명|arXiv (Cornell University)|2023. 08. 31.

ECG Monitoring and AnalysisMedicine인용 수 3

한 줄 요약

이 논문은 220,251건의 전문가가 검토한 대규모 ECG 데이터셋에서 마스킹된 자동에코드(pre-training)를 활용한 전기심장도(ECG) 분류를 위한 마스킹된 트랜스포머 모델인 MTECG를 제안한다. ECG 신호를 겹치지 않는 세그먼트로 나누고 경량 디코더를 사용해 변동하는 복원 목표를 적용함으로써, MTECG-T는 사설 및 공개 ECG 데이터셋에서 최신 기술(SOTA) 성능을 달성하며, 경량 5.7M 파라미터 아키텍처를 유지하면서도 최근 SOTA 방법들을 능가한다.

ABSTRACT

Electrocardiogram (ECG) is one of the most important diagnostic tools in clinical applications. With the advent of advanced algorithms, various deep learning models have been adopted for ECG tasks. However, the potential of Transformer for ECG data has not been fully realized, despite their widespread success in computer vision and natural language processing. In this work, we present Masked Transformer for ECG classification (MTECG), a simple yet effective method which significantly outperforms recent state-of-the-art algorithms in ECG classification. Our approach adapts the image-based masked autoencoders to self-supervised representation learning from ECG time series. We utilize a lightweight Transformer for the encoder and a 1-layer Transformer for the decoder. The ECG signal is split into a sequence of non-overlapping segments along the time dimension, and learnable positional embeddings are added to preserve the sequential information. We construct the Fuwai dataset comprising 220,251 ECG recordings with a broad range of diagnoses, annotated by medical experts, to explore the potential of Transformer. A strong pre-training and fine-tuning recipe is proposed from the empirical study. The experiments demonstrate that the proposed method increases the macro F1 scores by 3.4%-27.5% on the Fuwai dataset, 9.9%-32.0% on the PTB-XL dataset, and 9.4%-39.1% on a multicenter dataset, compared to the alternative methods. We hope that this study could direct future research on the application of Transformer to more ECG tasks.

연구 동기 및 목표

• 비록 시각 및 자연어 처리(NLP) 분야에서 성공을 거둔 트랜스포머임에도 불구하고 ECG 분류에서의 성능 한계를 해결하기 위해.
• 작은 레이블이 부여된 ECG 데이터셋 문제를 해결하기 위해 220,251건의 ECG 기록으로 구성된 대규모 전문가가 검토한 데이터셋을 구축하기 위해.
• 시각 및 NLP 분야에서 입증된 마스킹된 자동에코드(MAE) 기법을 시간 시리즈인 ECG 데이터에 적응시켜 자기지도 학습(pre-training)을 위해 활용하기 위해.
• 경량 트랜스포머 모델의 ECG 작업에서 효과적인 학습 레시피, 즉 복원 목표 설계, 학습률 스케줄링, 정규화 기법 등을 조사하기 위해.
• 적절한 사전학습 및 미세조정 전략을 결합했을 때, 일반적인 비전 트랜스포머(Vision Transformer) 아키텍처가 전문화된 CNN보다 뛰어난 성능을 낼 수 있는지 평가하기 위해.

제안 방법

• 시퀀스 길이를 줄이고 구조적 정보를 유지하기 위해 ECG 신호를 겹치지 않는 고정 길이 세그먼트로 나눈다.
• 순서 정보를 유지하고 세그먼트 위치를 구분하기 위해 학습 가능한 위치 임베딩을 사용한다.
• 비라벨 ECG 데이터에 마스킹된 자동에코드(MAE) 사전학습 전략을 적용하여, 무작위로 마스킹된 세그먼트를 경량 디코더를 사용해 복원하도록 모델을 훈련시킨다.
• 모델이 단순한 패턴이 아닌 의미 있는 파형 특징을 학습하도록 하기 위해 변동하는 복원 목표를 도입한다.
• 미세조정 단계에서 계층별 학습률 감소와 DropPath 정규화를 적용하여 과적합을 줄이고 일반화 성능을 향상시킨다.
• 최종 모델인 MTECG-T는 이 완전한 사전학습 및 미세조정 레시피를 통해 훈련된 경량 비전 트랜스포머 기반 아키텍처이다.

실험 결과

연구 질문

• RQ1 마스킹된 자동에코드 기법이 ECG 시간 시리즈 데이터에 효과적으로 적용되어 자기지도 표현 학습을 향상시킬 수 있는가?
• RQ2 복원 목표 설계, 훈련 스케줄 길이, 정규화 기법과 같은 훈련 구성 요소들이 ECG 분류 성능에 미치는 영향은 어떠한가?
• RQ3 적절한 사전학습 및 미세조정 전략을 결합했을 때, 경량의 일반적인 비전 트랜스포머 아키텍처가 ECG 분류에서 최신 기술(SOTA) 성능을 달성할 수 있는가?
• RQ4 제안된 마스킹 사전학습 방법은 레이블이 부족한 임상 환경에서 다른 모odalities(예: 에코카디오그래피)의 레이블이 부족한 상황에서도 비라벨 ECG 데이터를 효과적으로 활용할 수 있도록 하는가?
• RQ5 트랜스포머 기반 모델이 CNN보다 성능이 뛰어난 이유가 아키텍처적 우월성 때문인지, 사전학습 레시피의 효과성 때문인가?

주요 결과

• MTECG-T는 사설 및 공개 ECG 데이터셋에서 최신 기술(SOTA) 성능을 달성하며, 최근 SOTA 방법들보다 분류 정확도에서 뛰어나다.
• 570만 개의 파라미터를 가진 모델은 마스킹 비율(5%에서 75%)이 넓은 범위에서 안정적으로 성능을 유지하여 사전학습 초모수에 대한 강건성을 보여준다.
• 제거 실험(ablation studies) 결과, 변동하는 복원 목표, 더 긴 훈련 스케줄, 계층별 학습률 감소, 최적의 DropPath 비율이 성능 향상에 핵심적임을 확인하였다.
• 모델 크기를 일정 수준 이상으로 늘리면 성능 저하가 발생함을 확인하여, 데이터셋 크기가 더 큰 아키텍처의 이점을 제한함을 시사한다.
• 트랜스포머가 경량 ECG 분류에 부적합하다는 가정을 뒤집으며, 적절한 훈련 레시피를 적용한 일반적인 ViT가 전문화된 CNN을 능가할 수 있음을 보여준다.
• MTECG의 자기지도 학습 성격 덕분에 비라벨 ECG 데이터를 효과적으로 활용할 수 있어, 레이블 데이터가 제한된 새로운 임상 환경에 적합하다.