[論文レビュー] Bootstrapping Your Own Positive Sample: Contrastive Learning With Electronic Health Record Data

Tingyi Wanyan, Jing Zhang+4 more|arXiv (Cornell University)|Apr 7, 2021

Machine Learning in Healthcare参考文献 37被引用数 8

ひとこと要約

本論文は、COVID-19患者の死亡予測を向上させるために、電子歴史記録（EHR）データに対する特徴ベースおよび属性ベースの新しい陽性サンプリング戦略を備えた対照的学習フレームワークを提案する。特徴空間における患者の類似性と事前計算された類似度メトリクスを活用することで、クラス間分離を強化し、クラス内多様性を保持する。結果として、AUROCが0.959に達し、交差エントロピー（0.873）およびフォーカル損失（0.931）を上回る性能を達成した。

ABSTRACT

Electronic Health Record (EHR) data has been of tremendous utility in Artificial Intelligence (AI) for healthcare such as predicting future clinical events. These tasks, however, often come with many challenges when using classical machine learning models due to a myriad of factors including class imbalance and data heterogeneity (i.e., the complex intra-class variances). To address some of these research gaps, this paper leverages the exciting contrastive learning framework and proposes a novel contrastive regularized clinical classification model. The contrastive loss is found to substantially augment EHR-based prediction: it effectively characterizes the similar/dissimilar patterns (by its "push-and-pull" form), meanwhile mitigating the highly skewed class distribution by learning more balanced feature spaces (as also echoed by recent findings). In particular, when naively exporting the contrastive learning to the EHR data, one hurdle is in generating positive samples, since EHR data is not as amendable to data augmentation as image data. To this end, we have introduced two unique positive sampling strategies specifically tailored for EHR data: a feature-based positive sampling that exploits the feature space neighborhood structure to reinforce the feature learning; and an attribute-based positive sampling that incorporates pre-generated patient similarity metrics to define the sample proximity. Both sampling approaches are designed with an awareness of unique high intra-class variance in EHR data. Our overall framework yields highly competitive experimental results in predicting the mortality risk on real-world COVID-19 EHR data with a total of 5,712 patients admitted to a large, urban health system. Specifically, our method reaches a high AUROC prediction score of 0.959, which outperforms other baselines and alternatives: cross-entropy(0.873) and focal loss(0.931).

研究の動機と目的

• 臨床予測タスクにおけるEHRデータのクラス不均衡と高いクラス内分散の課題に対処すること。
• 視覚タスクとは異なり、データ拡張が比較的容易でないEHR分野における従来の対照的学習の限界を克服すること。
• 意味的類似性を捉えつつ、同じ結果クラス内での患者多様性を保持するように適合されたEHRデータ向けの陽性サンプリング戦略を開発すること。
• 最小限の教師信号を伴う対照的正則化を用いて、不均衡なEHRデータにおける臨床モデルの予測性能を向上させること。
• 対照的学習に情報に基づいた陽性サンプリングを組み合わせることで、現実世界のEHR応用においてより一般化可能で頑健な表現が得られることを示すこと。

提案手法

• クラス不均衡の処理と特徴学習の向上を目的に、フォーカル損失と対照的損失を組み合わせた対照的正則化分類モデルを導入する。
• 学習済み特徴空間内の近傍を選び、陽性ペアを形成する特徴ベースの陽性サンプリング戦略を提案する。これにより、局所的近傍学習が向上する。
• 臨床的特徴に基づくなど、事前計算された患者類似度メトリクス（例：臨床的特徴に基づく）を用いて、意味的に意味のある陽性ペアを定義する属性ベースの陽性サンプリング戦略を開発する。
• 「引き寄せ・引き離し」メカニズムを持つ対照的損失関数を適用する：類似した患者（陽性サンプル）を埋め込み空間で近づけ、類似しない患者（陰性サンプル）を遠ざける。
• UMAP可視化と埋め込み分離スコア（ESS）を用いて、学習済み表現空間におけるクラス間およびクラス内クラスタリングの質を評価する。
• 線形分類器を対照的埋め込み上で訓練し、下流の予測性能を評価するとともに、学習済み表現の質を検証する。

実験結果

リサーチクエスチョン

• RQ1 標準的な教師あり損失と比較して、カスタムの陽性サンプリング戦略を用いた対照的学習は、不均衡なEHRデータにおける予測性能を向上させることができるか？
• RQ2 特徴ベースおよび属性ベースの陽性サンプリング戦略は、ランダムサンプリングと比較してEHRデータにおけるクラス内多様性をどれほど効果的に捉えられるか？
• RQ3 提案されたサンプリング戦略は、学習済み埋め込み空間におけるクラス間分離をどの程度向上させるか？
• RQ4 対照的正則化は、EHRにおける高いデータ多様性とクラス不均衡の下でも、より頑健で一般化可能な表現を生成するか？
• RQ5 提案手法は、COVID-19患者の入院内死亡予測といった現実世界の臨床予測タスクにおいて、優れた性能を達成できるか？

主な発見

• 提案された対照的正則化モデルは、現実のCOVID-19 EHRデータ上でAUROC 0.959を達成し、交差エントロピー（0.873）およびフォーカル損失（0.931）を顕著に上回った。
• 属性ベースの陽性サンプリング戦略が最も高い性能を示し、ランダムサンプリングと比較して線形分類器の正確度が4.5%向上した。
• 特徴および属性ベースのサンプリング戦略は、ランダムサンプリングよりも高い埋め込み分離スコア（ESS）を達成し、クラス間分離の向上を示した。
• 両方の新規サンプリング戦略は、ランダムサンプリングと比較して、より高いクラス内分散（クラスタ中心からの距離の標準偏差）を維持しており、患者の多様性のより良いモデリングを示した。
• UMAP可視化により、提案手法が死亡患者と非死亡患者のクラスタをランダムサンプリングよりも明確に分離したことが確認された。
• 対照的正則化は、ランダムな陽性サンプリングでよく見られる埋め込み空間の収縮を緩和し、各クラス内での多様なサブグループ表現を維持した。