[論文レビュー] Masked Transformer for Electrocardiogram Classification

Zhou Ya, Xiaolin Diao+4 more|arXiv (Cornell University)|Aug 31, 2023

ECG Monitoring and AnalysisMedicine被引用数 3

ひとこと要約

本稿では、220,251件のエキスパートアノテート済みECGデータセットを用いたマスク化自己符号化事前学習を活用する、心電図（ECG）分類向けのマスク化トランスフォーマーモデルMTECGを提案する。ECG信号を重複のないセグメントに分割し、軽量なデコーダーを用いて変動する再構成ターゲットを適用することで、MTECG-TはプライベートおよびパブリックなECGデータセットの両方で最先端の性能を達成し、最近のSOTA手法を上回る結果を得るとともに、5.7Mパラメータの軽量アーキテクチャを維持している。

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に適応し、自己教師付き事前学習を実現する。
• 軽量トランスフォーマーモデルにおけるECGタスクの性能向上に寄与する、再構成ターゲット設計、学習率スケジューリング、正則化手法などの有効なトレーニングレシピを調査する。
• 適切な事前学習および微調整戦略を組み合わせた際、ベーシックなビジョントランスフォーマー（ViT）アーキテクチャが、特化したCNNを上回る性能を示すかどうかを検証する。

提案手法

• ECG信号を重複のない固定長のセグメントに分割することで、シーケンス長を短縮し、構造的情報を保持する。
• 順序の維持とセグメント位置の区別を図るため、学習可能な位置埋め込みを用いる。
• ラベルなしECGデータに対してマスク化自己符号化（MAE）事前学習戦略を適用し、ランダムにマスクされたセグメントを軽量デコーダーを用いて再構成する。
• モデルが単純なパターンではなく、意味のある波形特徴を学習するよう促すために、変動する再構成ターゲットを導入する。
• 微調整段階では、レイヤーごとの学習率減衰とDropPath正則化を適用し、過学習を低減し一般化性能を向上させる。
• 最終的なモデル、MTECG-Tは、この包括的な事前学習および微調整レシピで訓練された、軽量なビジョントランスフォーマー基盤アーキテクチャである。

実験結果

リサーチクエスチョン

• RQ1 マスク化自己符号化は、ECG時系列データに効果的に適応可能であり、自己教師付き表現学習の向上に寄与するか？
• RQ2 再構成ターゲット設計、トレーニングスケジュールの長さ、正則化手法といったトレーニング要因が、ECG分類におけるモデル性能に与える影響は何か？
• RQ3 適切な事前学習および微調整戦略を組み合わせた場合、軽量でベーシックなビジョントランスフォーマー（ViT）アーキテクチャがECG分類で最先端の性能を達成できるか？
• RQ4 提案されたマスク化事前学習手法は、他のモodal（例：エコーカーディオグラフィ）からのラベルが限られる臨床的状況において、ラベルなしECGデータの有効な活用を可能にするか？
• RQ5 トランスフォーマー基盤モデルがCNNを上回る性能を示す要因は、アーキテクチャの優位性にあるのか、それとも事前学習レシピの有効性にあるのか？

主な発見

• MTECG-Tは、プライベートおよびパブリックなECGデータセットの両方で最先端の性能を達成し、分類精度において最近のSOTA手法を上回っている。
• 570万パラメータのモデルは、マスキング比（5%〜75%）の広い範囲で安定した性能を示しており、事前学習ハイパーパrameterに対する頑健性を示している。
• アブレーションスタディにより、変動する再構成ターゲット、より長いトレーニングスケジュール、レイヤーごとの学習率減衰、最適なDropPath率が性能向上に不可欠であることが確認された。
• モデルを特定のサイズを超えて拡大すると性能が低下する傾向を示しており、これはデータセットサイズがより大きなアーキテクチャの恩恵を制限していることを示唆している。
• トランスフォーマーが軽量ECG分類に不適切であるという仮定に反し、適切なトレーニングレシピを備えたベーシックなViTが、特化したCNNを上回ることを示した。
• MTECGの自己教師付き性質により、ラベルなしECGデータの有効活用が可能となり、ラベルデータが限られる新しい臨床的状況に適している。