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

Resumen Técnico: CleanSurvival

Enunciado del Problema
El preprocesamiento de datos es un componente crítico, aunque a menudo pasado por alto, en los flujos de trabajo de aprendizaje automático, impactando significativamente el rendimiento de los modelos. Si bien los marcos de Aprendizaje Automático Automatizado (AutoML) han comenzado a integrar el preprocesamiento automatizado para tareas de clasificación y regresión, existe una notable falta de soluciones automatizadas y especializadas para el análisis de supervivencia (modelado de tiempo hasta el evento). El análisis de supervivencia presenta desafíos únicos debido a los datos censurados, donde los tiempos de evento se observan solo parcialmente. Las herramientas de AutoML existentes o bien no manejan los datos censurados de forma nativa, dependen de un preprocesamiento manual, o convierten las tareas de supervivencia en problemas de clasificación (por ejemplo, mediante discretización del tiempo), lo que conduce a una pérdida de información. Además, las herramientas actuales de preprocesamiento automatizado a menudo carecen de soporte para las estructuras de recompensa y los tipos de modelos específicos requeridos para el análisis de supervivencia, como el modelo de riesgos proporcionales de Cox o los Bosques Aleatorios de Supervivencia.

Metodología
El artículo introduce CleanSurvival, un marco de preprocesamiento de datos automatizado diseñado específicamente para el análisis de supervivencia, construido sobre un enfoque de Aprendizaje por Refuerzo (RL) utilizando Q-learning.

• Arquitectura: El marco trata la selección de pasos de preprocesamiento como un problema de toma de decisiones secuencial. Un agente navega por un espacio de estados definido por las características del conjunto de datos para seleccionar una secuencia de acciones (técnicas de preprocesamiento) que maximice una función de recompensa.
• Acciones de Preprocesamiento: El marco soporta una variedad de técnicas para:
  • Imputación de Valores Faltantes: Análisis de Casos Completos (CCA), imputación por media/mediana, Imputación Múltiple por Ecuaciones Encadenadas (MICE) y k-Vecinos Más Cercanos (KNN).
  • Detección de Valores Atípicos: Elipse de Envolvente (basada en la distancia de Mahalanobis) y residuos de Martingala.
  • Selección/Extracción de Características: Selección univariada de Riesgos Proporcionales de Cox, LASSO, Eliminación Recursiva de Características (RFE) y Ganancia de Información.
• Modelos de Supervivencia: El sistema es compatible con diversos modelos de tiempo hasta el evento, incluyendo el modelo de Riesgos Proporcionales de Cox, Bosques Aleatorios de Supervivencia (RSF), DeepHit (un enfoque de aprendizaje profundo) y modelos proporcionados por el usuario.
• Estructura de Recompensa: El agente de Q-learning se guía por una señal de recompensa basada en el Índice de Concordancia (C-index) evaluado en un conjunto de retención. Esto asegura que el flujo de preprocesamiento seleccionado esté optimizado específicamente para el rendimiento predictivo del modelo de supervivencia aguas abajo.
• Modos de Funcionamiento: La herramienta ofrece cuatro modos:
  1. Algoritmo Principal: Optimización mediante Q-learning.
  2. Limpieza Aleatoria: Búsqueda aleatoria para comparación de línea base.
  3. Flujo Personalizado: Flujos fijos definidos por el usuario mediante archivos de configuración.
  4. Sin Preparación: Entrada de datos crudos para rendimiento de línea base.

Contribuciones Clave

1. Primero Preprocesamiento Basado en RL para Supervivencia: CleanSurvival extiende el marco Learn2Clean (originalmente para regresión/clasificación) para manejar las restricciones específicas de los datos de tiempo hasta el evento censurados.
2. Optimización Específica del Modelo: A diferencia del preprocesamiento estático, el marco selecciona dinámicamente flujos que son óptimos para clases de modelos específicas (por ejemplo, un flujo óptimo para Cox puede diferir de uno para RSF).
3. Implementación de Código Abierto: Los autores proporcionan un paquete de Python disponible en GitHub, integrado con bibliotecas estándar como scikit-survival.
4. Manejo de Censura: El marco incorpora explícitamente la ponderación por probabilidad inversa de censura (IPCW) en métricas como la Puntuación Brier Integrada (IBS) para tener en cuenta las observaciones censuradas durante la evaluación.

Resultados Experimentales
El marco fue evaluado en tres conjuntos de datos del mundo real: Rotterdam (cáncer de mama), Flchain (cadena ligera libre en suero) y GBSG (Grupo de Estudio de Cáncer de Mama Alemán). Se introdujo ausencia sintética (MCAR, MAR, MNAR) para probar la robustez.

• Rendimiento: CleanSurvival superó consistentemente a las líneas base simples (Análisis de Casos Completos e Imputación por Media) y a menudo igualó o superó la búsqueda aleatoria no dirigida y Optuna (optimización bayesiana) en términos de C-index.
  • En el conjunto de datos Flchain, CleanSurvival logró un C-index de 0.67, en comparación con 0.39 para CCA y 0.62 para Imputación por Media.
  • En el conjunto de datos Rotterdam con un 50% de ausencia, CleanSurvival logró C-index entre 0.82 y 0.83, superando significativamente a CCA (0.77–0.79) e Imputación por Media (0.78–0.79).
• Selección de Características: La política de Q-learning seleccionó frecuentemente operadores de selección de características, apareciendo LASSO en aproximadamente el 68% de los flujos optimizados y RFE en aproximadamente el 42%.
• Calibración: La Puntuación Brier Integrada (IBS) indicó que la calibración fue generalmente estable a través de los mecanismos de ausencia, con la excepción de las condiciones MCAR del 40% que mostraron mayor varianza.
• Tiempo de Ejecución: El comportamiento del tiempo de ejecución dependió de las condiciones. Aunque el marco identificó con éxito flujos de alto rendimiento, el tiempo para alcanzar el 95% del mejor C-index observado varió según el conjunto de datos y el mecanismo de ausencia.

Significado y Afirmaciones
El artículo afirma que CleanSurvival aborda una brecha crítica en el ecosistema de AutoML al generalizar el preprocesamiento automatizado al análisis de supervivencia. Su importancia radica en:

• Robustez: Demostrar que el preprocesamiento automatizado guiado por RL puede producir modelos de supervivencia más robustos a través de niveles variables de ausencia y ruido en comparación con enfoques manuales o estáticos.
• Accesibilidad: Hacer que los estudios de análisis de supervivencia sean "más fáciles y rápidos de realizar" para profesionales que pueden carecer de la experiencia de dominio para ajustar manualmente flujos de preprocesamiento complejos para datos censurados.
• Integración: Proporcionar un puente fluido entre la limpieza de datos y el modelado de supervivencia aguas abajo, asegurando que las decisiones de limpieza estén guiadas directamente por su impacto en las predicciones de supervivencia en lugar de métricas genéricas de calidad de datos.

Los autores mantienen un tono modesto, reconociendo que, aunque los resultados son prometedores, la generalización a conjuntos de datos clínicos de alta dimensión y multimodales requiere mayor verificación, y la sobrecarga computacional del bucle de Q-learning (especialmente con modelos KNN y de aprendizaje profundo) sigue siendo una limitación. Se propone trabajo futuro para extender el marco a tareas adicionales de preprocesamiento y mejorar la escalabilidad.