[論文レビュー] Selective prediction-set models with coverage guarantees

Jean Feng, Arjun Sondhi+2 more|arXiv (Cornell University)|Jun 13, 2019

Machine Learning in Healthcare参考文献 25被引用数 12

ひとこと要約

本論文は、高不確実性の予測を避けることで信頼性を向上させる、選択的予測集合（SPS）モデルを学習するためのペナルティ付き損失最小化フレームワークを提案する。K-フォールド交差検証を用いてモデルをアンサンブル化することで、名目水準に近い被覆率と narrower な信頼区間を達成し、MNIST画像分類およびICU患者のリスク予測タスクの両方で、既存手法を上回る性能を発揮する。

ABSTRACT

Though black-box predictors are state-of-the-art for many complex tasks, they often fail to properly quantify predictive uncertainty and may provide inappropriate predictions for unfamiliar data. Instead, we can learn more reliable models by letting them either output a prediction set or abstain when the uncertainty is high. We propose training these selective prediction-set models using an uncertainty-aware loss minimization framework, which unifies ideas from decision theory and robust maximum likelihood. Moreover, since black-box methods are not guaranteed to output well-calibrated prediction sets, we show how to calculate point estimates and confidence intervals for the true coverage of any selective prediction-set model, as well as a uniform mixture of K set models obtained from K-fold sample-splitting. When applied to predicting in-hospital mortality and length-of-stay for ICU patients, our model outperforms existing approaches on both in-sample and out-of-sample age groups, and our recalibration method provides accurate inference for prediction set coverage.

研究の動機と目的

• 不確実性が高い場合に予測を控えることができる汎用的な機械学習モデルの学習フレームワークを構築すること。
• 特に高リスクの医療現場において、既存の機械学習ベースの臨床予測モデルに理論的被覆保証が欠如している問題に対処すること。
• 点推定ではなく予測集合を出力できるようにすることで、臨床意思決定における不確実性の定量化を向上させること。
• K-フォールド交差検証を用いたモデルに依存しない統計的推論手順を開発し、マージナル被覆率の推定精度を向上させること。
• 実世界の医療データセットを用いて、提案手法の実証的妥当性を検証し、精度と被覆安定性の向上を示すこと。

提案手法

• 予測精度と予測控除行動を同時に最適化するペナルティ付き経験的損失関数を提案し、分布外または予測が難しいケースでモデルが控除するよう促進する。
• 異なるデータサブグループ全体にわたって一貫した条件付き被覆を促進するために、一様な受容ペナルティを導入する。
• K-フォールド交差検証を用いて複数のSPSモデルを学習し、それらをアンサンブル化することでマージナル被覆率推定の精度を向上させる。
• 有限標本被覆保証を持つ有効な予測集合を構築するために、統計的補正の原則を適用する。
• K-フォールドモデルの集約予測を用いて、アンサンブルのマージナル被覆率に関する統計的推論を実施する。
• MNISTおよびMIMIC-IIIデータセット上で、pSPS目的関数を用いて深層ニューラルネットワークを学習し、性能と被覆率を評価する。

実験結果

リサーチクエスチョン

• RQ1 統一されたペナルティ付き損失フレームワークは、機械学習モデルが選択的に予測集合を出力するか控除するのを可能にし、医療応用における信頼性を向上させることができるか？
• RQ2 単一分割検証と比較して、SPSモデルのK-フォールド交差検証に基づくアンサンブル化は、マージナル被覆率推定の精度と正確性にどのように影響するか？
• RQ3 pSPSモデルは、分布外または複雑なケースで控除することで、特に臨床現場において予測精度をどの程度向上させることができるか？
• RQ4 提案手法は、特に実世界のICU患者予測タスクにおいて、既存手法と比較して名目水準に近い被覆率を達成するか？
• RQ5 ブラックボックス手法を用いて学習されたSPSモデルのマージナル被覆率を、モデルに依存しない推論手順が信頼性を持って推定できるか？

主な発見

• 入院死亡リスク予測タスクにおいて、pSPSモデルはテスト損失0.3768 (0.006) を達成し、ベースライン手法を上回った。
• 入院死亡リスク予測タスクでは、pSPSモデルの受容率は51.44% (2.44)、マージナル被覆率推定値は0.768 (0.047) であり、名目水準0.75に近く、良好な結果を示した。
• 入院滞在期間予測タスクでは、マージナル被覆率推定値が0.759 (0.007) で、95%信頼区間 (0.752, 0.766) を示し、テスト被覆率0.756とよく一致した。
• アンサンブルモデルのマージナル被覆率の信頼区間幅は、個々のモデルの約半分であり、推定の精度が向上していることを示した。
• pSPSモデルは、両タスクにおいて若年層の患者に対してより頻繁に控除するよう学習し、受容されたケースのテスト損失と精度が低下した。
• MIMIC-IIIデータセットでは、最終的なLOS用pSPSモデルが、若年層の患者に対して有意に低い受容確率を割り当てており、これは変動性の高さを反映し、より良いキャリブレーションが達成されたことを示している。