[論文レビュー] X-CAL: Explicit Calibration for Survival Analysis

Mark Goldstein, Xintian Han+3 more|PubMed|Jan 13, 2021

Machine Learning in Healthcare参考文献 17被引用数 13

ひとこと要約

本稿では、トレーニング中に直接分布型キャリブレーション（d-キャリブレーション）を最適化する、微分可能で明示的なキャリブレーション目的関数であるX-CALを導入する。最大尤度推定と統合することで、シミュレートされたデータ、MNISTベースの生存予測タスク、MIMIC-III、TCGAを含む多様なデータセットにおいて、 concordance や尤度を著しく損なわずにモデルのキャリブレーションが向上し、医療分野のリスク予測の信頼性が向上する。

ABSTRACT

Survival analysis models the distribution of time until an event of interest, such as discharge from the hospital or admission to the ICU. When a model's predicted number of events within any time interval is similar to the observed number, it is called <i>well-calibrated</i>. A survival model's calibration can be measured using, for instance, distributional calibration (D-CALIBRATION) [Haider et al., 2020] which computes the squared difference between the observed and predicted number of events within different time intervals. Classically, calibration is addressed in post-training analysis. We develop explicit calibration (X-CAL), which turns D-CALIBRATION into a differentiable objective that can be used in survival modeling alongside maximum likelihood estimation and other objectives. X-CAL allows practitioners to directly optimize calibration and strike a desired balance between predictive power and calibration. In our experiments, we fit a variety of shallow and deep models on simulated data, a survival dataset based on MNIST, on length-of-stay prediction using MIMIC-III data, and on brain cancer data from The Cancer Genome Atlas. We show that the models we study can be miscalibrated. We give experimental evidence on these datasets that X-CAL improves D-CALIBRATION without a large decrease in concordance or likelihood.

研究の動機と目的

• 生存モデルにおいてトレーニング中におけるキャリブレーション最適化の欠如に取り組むこと。これは、強力な予測性能を示すものの、しばしばキャリブレーションが不十分であるためである。
• トレーニング中に分布型キャリブレーション（d-キャリブレーション）を明示的に最適化する手法を開発すること。後処理によるキャリブレーションに依存するのではなく、トレーニング段階で直接最適化する。
• 予測精度とキャリブレーションのバランスを調整可能なハイパーパrameterを用いて、実務家が調整可能にすることで、解釈可能性と臨床的信頼性を向上させること。
• X-CALの有効性を、浅いモデルと深いモデル、実世界データと合成データを含む多様な生存モデリングシナリオで評価すること。
• トレーニング中に明示的なキャリブレーションを実施することで、尤度や concordance に著しい損失を伴わせることなく、よりキャリブレートされた生存予測が得られることを示すこと。

提案手法

• X-CALは、d-キャリブレーションを微分可能な目的関数として定式化し、標準的な生存モデリングフレームワーク内でのエンドツーエンド最適化を可能にする。
• 経験的分布の微分可能な緩和を用いて、時間間隔における予測済みイベント数と観測済みイベント数の二乗差を近似する。
• X-CALは最大尤度推定と併せて正則化項として統合され、予測尤度とキャリブレーションの両方を同時に最適化可能になる。
• 計算のスケーラビリティを確保するため、データのサブサンプリングを用いてキャリブレーション目的関数の計算を効率化する。
• パラメトリック（例：Weibull、Log-Normal）およびノンパラメトリック（例：Cat、mtlr）の両方の生存モデルに対応可能であり、広範な適用性を有する。
• ハイパーパrameter λ により、尤度とキャリブレーションのトレードオフを制御可能であり、実務家が望むキャリブレーション-予測バランスに調整可能である。

実験結果

リサーチクエスチョン

• RQ1 トレーニング中に明示的なキャリブレーションを実施することで、予測性能を損なわせることなく、d-キャリブレーションが著しく改善されるか？
• RQ2 異なるモデルアーキテクチャとデータセットにおいて、X-CALは標準的な最大尤度トレーニングと比較して、キャリブレーションと concordance の両面で優れているか？
• RQ3 MIMIC-III やがんゲノムアトラス（The Cancer Genome Atlas）のような実世界の生存データセットにおいて、X-CALはどれほどmiscalibrationを低減できるか？
• RQ4 X-CALの統合により、特に滞在期間予測や生存予測タスクにおいて、臨床意思決定に役立つより信頼性の高いリスク推定が得られるか？
• RQ5 X-CALは、浅いモデルおよび深い生存モデル、特に複雑なニューラルネットワーク構造を有するモデルに対しても、効果的に適用可能か？

主な発見

• X-CALは、すべてのデータセットでd-キャリブレーション誤差を著しく低減した。MNISTベースの生存タスクでは、d-キャリブレーションがλ=1000のとき、ベースラインの0.018から0.002に低下したが、concordanceは0.80以上を維持した。
• MIMIC-IIIの滞在期間予測では、X-CALを適用したCat-iモデルがλ=1000でd-キャリブレーション0.001を達成した。X-CALなしでは0.002であった。concordanceはわずかに低下（λ=0のとき0.758から0.748に）したが、ほとんど差がない。
• TCGAの膠芽腫（glioma）データでは、X-CALによりd-キャリブレーションがWeibullモデルの0.035から、Cat-iモデル（λ=1000）で0.003に低下した。ほとんどのモデルでconcordanceは0.80以上を維持した。
• TCGAデータにおいて、X-CALを適用したLog-Normalモデルでは、d-キャリブレーションが70％減少（0.059から0.004に）したが、concordanceは0.80前後で安定した。
• X-CALは、パラメトリック、ノンパラメトリック、およびディープなモデルのすべてでキャリブレーションを向上させ、その強靭性と一般化能力を示した。
• この方法により、制御可能なトレードオフが可能である。λを大きくするとキャリブレーションが向上するが、尤度はわずかに低下する。しかし、concordanceはほとんど変化しなかった。