[論文レビュー] SODEN: A Scalable Continuous-Time Survival Model through Ordinary Differential Equation Networks

Weijing Tang, Jiaqi Ma+2 more|arXiv (Cornell University)|Aug 19, 2020

Statistical Methods and Inference参考文献 38被引用数 10

ひとこと要約

SODENは、打ち切りデータにおける尤度積分の非可解性という課題を克服するため、通常微分方程式（ODE）ネットワークを用いたスケーラブルな連続時間生存モデルを提案する。尤度推定を微分方程式制約付き最適化問題として定式化することで、ODEソルバーやアドジョイント感度解析を活用し、尤度と勾配を効率的に計算可能とし、ミニバッチ確率的勾配降下法（SGD）を用いたエンドツーエンド学習を可能にした。実世界の生存データセットにおいて最先端の性能を達成した。

ABSTRACT

In this paper, we propose a flexible model for survival analysis using neural networks along with scalable optimization algorithms. One key technical challenge for directly applying maximum likelihood estimation (MLE) to censored data is that evaluating the objective function and its gradients with respect to model parameters requires the calculation of integrals. To address this challenge, we recognize that the MLE for censored data can be viewed as a differential-equation constrained optimization problem, a novel perspective. Following this connection, we model the distribution of event time through an ordinary differential equation and utilize efficient ODE solvers and adjoint sensitivity analysis to numerically evaluate the likelihood and the gradients. Using this approach, we are able to 1) provide a broad family of continuous-time survival distributions without strong structural assumptions, 2) obtain powerful feature representations using neural networks, and 3) allow efficient estimation of the model in large-scale applications using stochastic gradient descent. Through both simulation studies and real-world data examples, we demonstrate the effectiveness of the proposed method in comparison to existing state-of-the-art deep learning survival analysis models. The implementation of the proposed SODEN approach has been made publicly available at https://github.com/jiaqima/SODEN.

研究の動機と目的

• 打ち切りデータにおける尤度評価における非可解な積分という、ディープ生存分析の根本的課題に取り組む。
• 比例ハザード仮定などの制限的仮定や時間の離散化に依存する既存手法の限界を克服する。
• ミニバッチ確率的勾配降下法（SGD）を用いた、柔軟なニューラルネットワークベースの生存モデルのスケーラブルな学習を可能にする。
• 強い構造的仮定を必要としない、広範な連続時間生存分布の族をサポートする一般化されたフレームワークを開発する。
• ヒューリスティック的または部分的に正当化された最適化手法に代わる、理論的に整合的で計算的に効率的な代替手法を提供する。

提案手法

• 累積ハザード関数をニューラルネットワークでパrameter化されたODEの解としてモデル化する。
• 打ち切りデータに対する最尤推定を微分方程式制約付き最適化問題として定式化する。
• 効率的な数値的ODEソルバー（例：適応的ステップサイズ法）を用いて生存関数と尤度を計算する。
• アドジョイント感度解析を活用し、ODE解を介して勾配を逆伝播することで、エンドツーエンド学習を可能にする。
• ODEベースの尤度計算をミニバッチSGDと互換性のあるディープラーニングフレームワークに統合する。
• 比例ハザード仮定やパラメトリックな形を仮定せず、ハザード関数の柔軟なモデリングを可能にする。

実験結果

リサーチクエスチョン

• RQ1 ODEソルバーを用いることで、数値的に扱える尤度評価と勾配計算が可能な連続時間生存モデルを定式化できるか？
• RQ2 ニューラルネットワークベースのODEモデルは、実世界の打ち切りデータにおいて、既存のディープラーニング生存モデルを上回る予測性能を示すか？
• RQ3 比例ハザード仮定を緩和することで、モデルのフィットと予測精度はどの程度向上するか？
• RQ4 ミニバッチ確率的勾配降下法を用いて、大規模データセットへのスケーラブルな適用が可能か？
• RQ5 リスクスコアリングタスクにおいて、 concordance index（ concordance index）とログランク検定で測定した際、本手法の性能はいかがなものか？

主な発見

• SODENは、MIMIC、SUPPORT、METABRICなどのすべての実世界データセットで最良の負の対数尤度（NLL）を達成し、ベースラインモデル比で15–20%の改善を示した。
• MIMICデータセットでは、SODENは時間依存型 concordance index（C^td）を 0.746 ± 0.003 で達成し、Cox（0.734 ± 0.003*）とDeepSurv（0.706 ± 0.003*）を顕著に上回った。
• アブレーションスタディの結果、PH制限付きのSODEN-PHおよびCox制限付きのSODEN-Coxモデルは、完全なSODENに比べて性能が劣ることが確認され、これらのデータセットではPH仮定が成り立たないことが示された。
• SODEN-PHおよびSODEN-CoxよりSODENのNLLが顕著に優れており、対比較t検定のp値が0.05未満であった。これは非線形的かつ非PHモデリングの重要性を示している。
• カプラン・マイヤー曲線では、MIMICデータセットにおいてSODENが特定した高リスク群と低リスク群の間で、統計的に有意な分離（p < 0.001）が確認された。
• 本手法はミニバッチSGDによるスケーラブルな学習を可能とし、複雑で高次元の特徴を有する大規模な生存分析アプリケーションに適している。