CleanSurvival: Automated data preprocessing for time-to-event models using reinforcement learning

Technical Summary: CleanSurvival

Problem Statement
Data preprocessing is a critical yet often overlooked component of machine learning pipelines, significantly impacting model performance. While Automated Machine Learning (AutoML) frameworks have begun integrating automated preprocessing for classification and regression tasks, there is a distinct lack of tailored, automated solutions for survival analysis (time-to-event modeling). Survival analysis presents unique challenges due to censored data, where event times are only partially observed. Existing AutoML tools either fail to handle censored data natively, rely on manual preprocessing, or convert survival tasks into classification problems (e.g., via time discretization), leading to information loss. Furthermore, current automated preprocessing tools often lack support for the specific reward structures and model types required for survival analysis, such as the Cox proportional hazards model or Random Survival Forests.

Methodology
The paper introduces CleanSurvival, an automated data preprocessing framework specifically designed for survival analysis, built upon a Reinforcement Learning (RL) approach using Q-learning.

• Architecture: The framework treats the selection of preprocessing steps as a sequential decision-making problem. An agent navigates a state space defined by the dataset's characteristics to select a sequence of actions (preprocessing techniques) that maximizes a reward function.
• Preprocessing Actions: The framework supports a variety of techniques for:
  • Missing Value Imputation: Complete Case Analysis (CCA), Mean/Median imputation, Multiple Imputation by Chained Equations (MICE), and k-Nearest Neighbors (KNN).
  • Outlier Detection: Elliptic Envelope (based on Mahalanobis distance) and Martingale residuals.
  • Feature Selection/Extraction: Univariate Cox Proportional Hazards selection, LASSO, Recursive Feature Elimination (RFE), and Information Gain.
• Survival Models: The system is compatible with various time-to-event models, including the Cox Proportional Hazards model, Random Survival Forests (RSF), DeepHit (a deep learning approach), and user-supplied models.
• Reward Structure: The Q-learning agent is guided by a reward signal based on the Concordance Index (C-index) evaluated on a hold-out set. This ensures that the selected preprocessing pipeline is optimized specifically for the predictive performance of the downstream survival model.
• Working Modes: The tool offers four modes:
  1. Main Algorithm: Q-learning optimization.
  2. Random Cleaning: Random search for baseline comparison.
  3. Custom Pipeline: User-defined fixed pipelines via configuration files.
  4. No Preparation: Raw data input for baseline performance.

Key Contributions

1. First RL-based Preprocessing for Survival: CleanSurvival extends the Learn2Clean framework (originally for regression/classification) to handle the specific constraints of censored time-to-event data.
2. Model-Specific Optimization: Unlike static preprocessing, the framework dynamically selects pipelines that are optimal for specific model classes (e.g., a pipeline optimal for Cox may differ from one for RSF).
3. Open-Source Implementation: The authors provide a Python package available on GitHub, integrating with standard libraries like scikit-survival.
4. Handling of Censoring: The framework explicitly incorporates inverse probability of censoring weighting (IPCW) in metrics like the Integrated Brier Score (IBS) to account for censored observations during evaluation.

Experimental Results
The framework was evaluated on three real-world datasets: Rotterdam (breast cancer), Flchain (serum free light chain), and GBSG (German Breast Cancer Study Group). Synthetic missingness (MCAR, MAR, MNAR) was introduced to test robustness.

• Performance: CleanSurvival consistently outperformed simple baselines (Complete Case Analysis and Mean Imputation) and often matched or exceeded undirected random search and Optuna (Bayesian optimization) in terms of C-index.
  • On the Flchain dataset, CleanSurvival achieved a C-index of 0.67, compared to 0.39 for CCA and 0.62 for Mean Imputation.
  • On the Rotterdam dataset with 50% missingness, CleanSurvival achieved C-indices between 0.82 and 0.83, significantly outperforming CCA (0.77–0.79) and Mean Imputation (0.78–0.79).
• Feature Selection: The Q-learning policy frequently selected feature selection operators, with LASSO appearing in ~68% of optimized pipelines and RFE in ~42%.
• Calibration: The Integrated Brier Score (IBS) indicated that calibration was generally stable across missingness mechanisms, with the exception of 40% MCAR conditions which showed higher variance.
• Runtime: Runtime behavior was condition-dependent. While the framework successfully identified high-performing pipelines, the time to reach 95% of the best-observed C-index varied based on the dataset and missingness mechanism.

Significance and Claims
The paper claims that CleanSurvival addresses a critical gap in the AutoML ecosystem by generalizing automated preprocessing to survival analysis. Its significance lies in:

• Robustness: Demonstrating that automated, RL-guided preprocessing can yield more robust survival models across varying levels of missingness and noise compared to manual or static approaches.
• Accessibility: Making survival analysis studies "easier and quicker to perform" for practitioners who may lack the domain expertise to manually tune complex preprocessing pipelines for censored data.
• Integration: Providing a seamless bridge between data cleaning and downstream survival modeling, ensuring that cleaning decisions are directly guided by their impact on survival predictions rather than generic data quality metrics.

The authors maintain a modest tone, acknowledging that while the results are promising, generalization to high-dimensional, multi-modal clinical datasets requires further verification, and the computational overhead of the Q-learning loop (especially with KNN and deep learning models) remains a limitation. Future work is proposed to extend the framework to additional preprocessing tasks and improve scalability.