Stage-specific cancer survival prediction enriched by explainable machine learning

TL;DR

This work tackles the problem that survival predictions trained on mixed cancer stages may misrepresent stage-specific outcomes. It deploys stage-stratified ML on SEER data for colorectal, stomach, and liver cancers, using four models (Logistic Regression, AdaBoost, Random Forest, CatBoost) and SHAP for global explanations plus LIME for local, case-level interpretations. CatBoost consistently delivers the strongest $AUC-ROC$ across stages and cancers, while SHAP and LIME reveal that age, tumor extension, lymph node involvement, surgical factors, and histologic subtypes drive stage-dependent survival patterns in clinically coherent ways. The findings support more transparent, stage-tailored prognostic tools that can inform personalized treatment planning, though external validation in diverse cohorts is needed to confirm generalizability and guide deployment in practice.

Abstract

Despite the fact that cancer survivability rates vary greatly between stages, traditional survival prediction models have frequently been trained and assessed using examples from all combined phases of the disease. This method may result in an overestimation of performance and ignore the stage-specific variations. Using the SEER dataset, we created and verified explainable machine learning (ML) models to predict stage-specific cancer survivability in colorectal, stomach, and liver cancers. ML-based cancer survival analysis has been a long-standing topic in the literature; however, studies involving the explainability and transparency of ML survivability models are limited. Our use of explainability techniques, including SHapley Additive exPlanations (SHAP) and Local Interpretable Model-agnostic Explanations (LIME), enabled us to illustrate significant feature-cancer stage interactions that would have remained hidden in traditional black-box models. We identified how certain demographic and clinical variables influenced survival differently across cancer stages and types. These insights provide not only transparency but also clinical relevance, supporting personalized treatment planning. By focusing on stage-specific models, this study provides new insights into the most important factors at each stage of cancer, offering transparency and potential clinical relevance to support personalized treatment planning.

Figures (8)

• Figure 1: ML workflow, including data collection, preprocessing, model training, evaluation, and interpretation using SHAP lundberg2017unified and LIME ribeiro2016should.
• Figure 2: Stage distribution of colorectal, stomach, and liver cancers. The plots illustrate the distribution of cases across stages, emphasizing the differing patterns of disease presentation among cancer types.
• Figure 3: Rule-based criteria for defining five-year cancer survival status pour2018stage.
• Figure 4: Correlation heatmaps between numerical features for colorectal, stomach, and liver cancers.
• Figure 5: The receiver operating characteristic (ROC) curves of the four models (Logistic Regression, Random Forest, AdaBoost, CatBoost) for colorectal cancer across different stages: (a) Localized, (b) Regional, and (c) Distant.