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

Riepilogo Tecnico: CleanSurvival

Enunciato del Problema
La pre-elaborazione dei dati è un componente critico, ma spesso trascurato, delle pipeline di machine learning, con un impatto significativo sulle prestazioni del modello. Sebbene i framework di Automated Machine Learning (AutoML) abbiano iniziato a integrare la pre-elaborazione automatizzata per compiti di classificazione e regressione, esiste una marcata carenza di soluzioni automatizzate e su misura per l'analisi di sopravvivenza (modellazione tempo-evento). L'analisi di sopravvivenza presenta sfide uniche a causa dei dati censurati, dove i tempi di evento sono osservati solo parzialmente. Gli strumenti AutoML esistenti o non gestiscono nativamente i dati censurati, o si affidano a una pre-elaborazione manuale, o convertono i compiti di sopravvivenza in problemi di classificazione (ad esempio, tramite discretizzazione del tempo), portando a una perdita di informazioni. Inoltre, gli attuali strumenti di pre-elaborazione automatizzata spesso mancano del supporto per le strutture di ricompensa e i tipi di modello specifici richiesti dall'analisi di sopravvivenza, come il modello dei rischi proporzionali di Cox o le Random Survival Forests.

Metodologia
Il documento introduce CleanSurvival, un framework di pre-elaborazione automatizzata dei dati specificamente progettato per l'analisi di sopravvivenza, basato su un approccio di Reinforcement Learning (RL) che utilizza il Q-learning.

• Architettura: Il framework tratta la selezione delle fasi di pre-elaborazione come un problema di decisione sequenziale. Un agente naviga in uno spazio degli stati definito dalle caratteristiche del dataset per selezionare una sequenza di azioni (tecniche di pre-elaborazione) che massimizza una funzione di ricompensa.
• Azioni di Pre-elaborazione: Il framework supporta una varietà di tecniche per:
  • Imputazione dei Valori Mancanti: Analisi dei Casi Completamente Osservati (CCA), imputazione per media/mediana, Imputazione Multipla tramite Equazioni Concatenate (MICE) e k-Nearest Neighbors (KNN).
  • Rilevamento degli Outlier: Ellittico (basato sulla distanza di Mahalanobis) e residui Martingala.
  • Selezione/Estrazione delle Caratteristiche: Selezione univariata dei Rischi Proporzionali di Cox, LASSO, Eliminazione Ricorsiva delle Caratteristiche (RFE) e Information Gain.
• Modelli di Sopravvivenza: Il sistema è compatibile con vari modelli tempo-evento, inclusi il modello dei Rischi Proporzionali di Cox, le Random Survival Forests (RSF), DeepHit (un approccio di deep learning) e modelli forniti dall'utente.
• Struttura di Ricompensa: L'agente Q-learning è guidato da un segnale di ricompensa basato sull'Indice di Concordanza (C-index) valutato su un set di hold-out. Ciò garantisce che la pipeline di pre-elaborazione selezionata sia ottimizzata specificamente per le prestazioni predittive del modello di sopravvivenza a valle.
• Modalità di Funzionamento: Lo strumento offre quattro modalità:
  1. Algoritmo Principale: Ottimizzazione Q-learning.
  2. Pulizia Casuale: Ricerca casuale per il confronto con la baseline.
  3. Pipeline Personalizzata: Pipeline fisse definite dall'utente tramite file di configurazione.
  4. Nessuna Preparazione: Input di dati grezzi per le prestazioni di base.

Contributi Chiave

1. Prima Pre-elaborazione basata su RL per la Sopravvivenza: CleanSurvival estende il framework Learn2Clean (originariamente per regressione/classificazione) per gestire i vincoli specifici dei dati tempo-evento censurati.
2. Ottimizzazione Specifica per Modello: A differenza della pre-elaborazione statica, il framework seleziona dinamicamente pipeline ottimali per specifiche classi di modelli (ad esempio, una pipeline ottimale per Cox può differire da una per RSF).
3. Implementazione Open-Source: Gli autori forniscono un pacchetto Python disponibile su GitHub, integrato con librerie standard come scikit-survival.
4. Gestione della Censura: Il framework incorpora esplicitamente la ponderazione inversa della probabilità di censura (IPCW) in metriche come il Brier Score Integrato (IBS) per tenere conto delle osservazioni censurate durante la valutazione.

Risultati Sperimentali
Il framework è stato valutato su tre dataset reali: Rotterdam (cancro al seno), Flchain (catene leggere sieriche libere) e GBSG (German Breast Cancer Study Group). È stata introdotta una mancanza sintetica (MCAR, MAR, MNAR) per testare la robustezza.

• Prestazioni: CleanSurvival ha costantemente superato le semplici baseline (Analisi dei Casi Completamente Osservati e Imputazione per Media) e spesso ha eguagliato o superato la ricerca casuale non guidata e Optuna (ottimizzazione bayesiana) in termini di C-index.
  • Sul dataset Flchain, CleanSurvival ha raggiunto un C-index di 0,67, rispetto a 0,39 per la CCA e 0,62 per l'Imputazione per Media.
  • Sul dataset Rotterdam con il 50% di dati mancanti, CleanSurvival ha raggiunto C-index compresi tra 0,82 e 0,83, superando significativamente la CCA (0,77–0,79) e l'Imputazione per Media (0,78–0,79).
• Selezione delle Caratteristiche: La politica Q-learning ha selezionato frequentemente operatori di selezione delle caratteristiche, con LASSO presente in circa il 68% delle pipeline ottimizzate e RFE in circa il 42%.
• Calibrazione: Il Brier Score Integrato (IBS) ha indicato che la calibrazione era generalmente stabile tra i meccanismi di mancanza, fatta eccezione per le condizioni MCAR al 40% che hanno mostrato una varianza più elevata.
• Tempo di Esecuzione: Il comportamento del tempo di esecuzione era dipendente dalle condizioni. Sebbene il framework abbia identificato con successo pipeline ad alte prestazioni, il tempo necessario per raggiungere il 95% del miglior C-index osservato variava in base al dataset e al meccanismo di mancanza.

Significato e Affermazioni
Il documento afferma che CleanSurvival colma una lacuna critica nell'ecosistema AutoML generalizzando la pre-elaborazione automatizzata all'analisi di sopravvivenza. Il suo significato risiede in:

• Robustezza: Dimostrando che la pre-elaborazione automatizzata e guidata da RL può produrre modelli di sopravvivenza più robusti attraverso vari livelli di mancanza e rumore rispetto agli approcci manuali o statici.
• Accessibilità: Rendendo gli studi di analisi di sopravvivenza "più semplici e veloci da eseguire" per i praticanti che potrebbero non possedere l'esperienza di dominio necessaria per sintonizzare manualmente pipeline di pre-elaborazione complesse per dati censurati.
• Integrazione: Fornendo un ponte fluido tra la pulizia dei dati e la modellazione di sopravvivenza a valle, garantendo che le decisioni di pulizia siano guidate direttamente dal loro impatto sulle previsioni di sopravvivenza piuttosto che da metriche generiche di qualità dei dati.

Gli autori mantengono un tono modesto, riconoscendo che, sebbene i risultati siano promettenti, la generalizzazione a dataset clinici ad alta dimensionalità e multimodali richiede ulteriore verifica, e il sovraccarico computazionale del ciclo Q-learning (specialmente con modelli KNN e deep learning) rimane una limitazione. È proposto un lavoro futuro per estendere il framework ad attività aggiuntive di pre-elaborazione e migliorare la scalabilità.