[論文レビュー] Deep Cox Mixtures for Survival Regression

Chirag Nagpal, Steve Yadlowsky+2 more|arXiv (Cornell University)|Jan 16, 2021

Statistical Methods and Inference参考文献 9被引用数 5

ひとこと要約

本論文は、患者を潜在的サブグループに分け、それぞれが独自の比例ハザード構造をとるモデルとして、深層ニューラルネットワークと非パrametricベースラインハザードを用いて推定する、柔軟な生存回帰モデル「Deep Cox Mixtures」を提案する。この手法は、特にマイノリティ層のデモグラフィック要因において、古典的および深層学習ベースラインを上回る大幅なキャリブレーションの向上を達成する。

ABSTRACT

Survival analysis is a challenging variation of regression modeling because of the presence of censoring, where the outcome measurement is only partially known, due to, for example, loss to follow up. Such problems come up frequently in medical applications, making survival analysis a key endeavor in biostatistics and machine learning for healthcare, with Cox regression models being amongst the most commonly employed models. We describe a new approach for survival analysis regression models, based on learning mixtures of Cox regressions to model individual survival distributions. We propose an approximation to the Expectation Maximization algorithm for this model that does hard assignments to mixture groups to make optimization efficient. In each group assignment, we fit the hazard ratios within each group using deep neural networks, and the baseline hazard for each mixture component non-parametrically. We perform experiments on multiple real world datasets, and look at the mortality rates of patients across ethnicity and gender. We emphasize the importance of calibration in healthcare settings and demonstrate that our approach outperforms classical and modern survival analysis baselines, both in terms of discriminative performance and calibration, with large gains in performance on the minority demographics.

研究の動機と目的

• 生存モデルのキャリブレーションが不十分である医療現場の課題に取り組むこと、特に比例ハザード仮定が成り立たない場合に焦点を当てる。
• 特に人種的・性的マイノリティ層のような代表が不足しているデモグラフィックグループにおける、予測性能（特にキャリブレーション）の向上を図ること。
• 潜在的患者サブグループを学習することで、比例ハザード仮定を越えた柔軟な生存モデルの構築を実現すること。
• 区別性能（ concordance index で測定）に悪影響を及げることなく、キャリブレーションが向上することを実証すること。
• 実世界の臨床データセットに適用可能な実用的でオープンソースの深層学習フレームワークを提供すること。

提案手法

• 各成分が潜在的患者サブグループを表し、それぞれに独自のハザード関数を持つコックス回帰成分の混合モデルを用いて生存結果をモデル化する。
• 患者特徴に基づくグループ固有のハザード比を深層ニューラルネットワークで推定し、非線形な関係を柔軟に表現する。
• スプラインに基づくスムージングを用いて、各グループ内でベースライン生存関数を非パラメトリックに推定し、パラメトリックな分布仮定を回避する。
• 最適化効率を向上させるために、混合成分へのハードアサインメントを用いた近似モンテカルロ期待最大化（EM）アルゴリズムを実装する。
• Eステップをカプラン＝マイヤー推定値で近似し、Mステップをグループ固有パラメータの確率的勾配降下法で実行する。
• 特に人種および性別ごとのサブグループレベルでのパフォーマンスに注目し、不確実性とキャリブレーションメトリクスをモデル評価に統合する。

実験結果

リサーチクエスチョン

• RQ1 深層ニューラルネットワークを用いたハザード比推定を伴うコックスモデルの混合モデルは、標準的なコックス回帰および深層学習ベースラインと比較して、生存予測のキャリブレーションを改善できるか？
• RQ2 提案手法は、ブラックおよび女性患者のような代表が不足しているデモグラフィックグループにおける生存予測の性能格差を低減できるか？
• RQ3 潜在的サブグループによる比例ハザード仮定の緩和が、区別性能とキャリブレーションの両方をどの程度向上できるか？
• RQ4 SUPPORTのような実世界のデータセットにおいて、推定されたサブグループ固有の生存曲線は、実証的なカプラン＝マイヤー推定値とどの程度一致するか？
• RQ5 モデルの解釈可能性およびサブグループの同定が、臨床的意思決定支援にどのように役立つか？

主な発見

• Deep Cox Mixturesは、すべてのデモグラフィックグループにおいてキャリブレーションを顕著に向上させ、特にブラックおよび女性患者のようなマイノリティ層で最も大きなパフォーマンス向上を達成した。
• 最先端のベースラインと同等またはそれを上回る区別性能（C-index）を維持しており、区別性能とキャリブレーションのトレードオフがないことが示された。
• サブグループ固有の生存曲線が交差する傾向を示しており、非比例ハザードを示しているが、モデルは潜在的クラスタリングによりこれを効果的に捉えている。
• サブグループのカプラン＝マイヤー推定値とモデル予測生存曲線がよく一致しており、モデルが意味のある患者サブグループを回復できていることが妥当性を裏付けている。
• 既存のツール（2013年ACC/AHAプールドコhort方程式）に見られるブラック集団における心血管リスクの過大評価をモデルが低減した。
• GitHub上のオープンソース実装により、将来的な医療分野および研究応用の再現性と拡張性が確保された。