Scalable Random Survival Forests via Risk Set Sampling

Resumen Técnico: Bosques de Supervivencia Aleatorios Escalables mediante el Muestreo del Conjunto de Riesgo

Planteamiento del Problema
Los Bosques de Supervivencia Aleatorios (RSF, por sus siglas en inglés) son una potente extensión de los bosques aleatorios para datos de tiempo hasta un evento con censura. Sin embargo, su aplicación a conjuntos de datos grandes o de alta dimensionalidad suele ser computacionalmente prohibitiva. El principal cuello de botella surge de la evaluación repetida de los "conjuntos de riesgo" necesarios para calcular la regla de división de log-rank durante la construcción del árbol. Para cada tiempo de evento, el algoritmo debe evaluar a todos los individuos que aún están en riesgo, lo que provoca que los costes computacionales escalen rápidamente con el tamaño de la muestra y el número de tiempos de evento distintos. Si bien el muestreo del conjunto de riesgo se ha integrado con éxito en las redes neuronales basadas en Cox (por ejemplo, Cox-Time) y es un diseño estándar en la epidemiología clásica (casos y controles anidados), aún no se ha adaptado para los RSF.

Metodología
Los autores proponen una modificación del algoritmo RSF que integra el muestreo del conjunto de riesgo directamente en el proceso de construcción de árboles.

• Mecanismo Central: En lugar de utilizar el conjunto de riesgo completo ($R_k$) en cada tiempo de evento $t_k$ para calcular el estadístico de división de log-rank, el método toma una muestra de una fracción fija de individuos del conjunto de riesgo ($\tilde{R}_k$).
• Ajuste Estadístico: El estadístico de log-rank depende del número esperado de eventos y su varianza dentro de un nodo. Los autores modifican estos cálculos para utilizar el conjunto de riesgo muestreado $\tilde{R}_k$ en lugar del conjunto completo. Específicamente, el número esperado de eventos $E(D_{\alpha,k})$ y la varianza $Var(D_{\alpha,k})$ para un nodo hija $\alpha$ se recalculan utilizando la proporción de individuos en riesgo muestreados ($\tilde{R}_{\alpha,k}$) respecto al conjunto de riesgo muestreado total ($\tilde{R}_k$).
• Implementación: Este enfoque se implementó en el paquete de R ranger, una implementación de alto rendimiento de bosques aleatorios. Los autores probaron fracciones de muestreo que van desde el 10% hasta el 100% (donde el 100% representa el conjunto de riesgo completo estándar).
• Diseño Experimental:
  • Simulaciones: Se generaron datos de supervivencia heterogéneos utilizando distribuciones mixtas de Exponencial y Weibull con tamaños de muestra de $n=500, 1000$ y $2000$y$p=15$ covariables.
  • Datos del Mundo Real: Se utilizaron dos conjuntos de datos: METABRIC (1904 observaciones de cáncer de mama) y el conjunto de datos del Grupo de Estudio de Cáncer de Mama Alemán (GBSG, 686 observaciones).
  • Métricas: El rendimiento se evaluó mediante el Índice de Concordancia (C-index) vía validación cruzada de 5 pliegues (5-fold) y estimaciones de fuera de la muestra (out-of-bag, OOB), junto con el tiempo de ejecución computacional.

Resultos Clave

• Rendimiento Predictivo: En todos los conjuntos de datos simulados y del mundo real, los modelos RSF entrenados con conjuntos de riesgo muestreados (específicamente utilizando del 25% al 40% del conjunto de riesgo) alcanzaron valores de C-index virtualmente indistinguibles de los modelos entrenados con conjuntos de riesgo completos.
• Eficiencia Computacional: Reducir el conjunto de riesgo al rango del 25–40% resultó en una reducción del tiempo de entrenamiento de más del 50%. El tiempo de ejecución aumentó monótonamente con la fracción del conjunto de riesgo, observándose las ganancias más significativas en los conjuntos de datos más grandes.
• Estabilidad y Sesgo: El estudio encontró que el muestreo del conjunto de riesgo afecta principalmente a la varianza de las estimaciones en lugar de introducir un sesgo sistemático. Los intervalos de confianza para el C-index mostraron un solapamiento sustancial en todas las fracciones de muestreo, lo que indica que no hay una diferencia estadísticamente significativa en la discriminación predictiva.
• Efecto de Regularización: En varios escenarios, los modelos entrenados con fracciones de muestreo moderadas (25–40%) exhibieron valores de C-index marginalmente superiores a los modelos de riesgo completo, lo que sugiere un posible efecto de regularización beneficioso que reduce el ruido.

Significación y Reivindicaciones
El artículo sostiene que el muestreo del conjunto de riesgo es una modificación simple pero efectiva que mejora sustancialmente la escalabilidad de los Bosques de Supervivencia Aleatorios sin comprometer el rendimiento predictivo o la estabilidad del modelo. Los autores argumentan que:

1. Puente entre Disciplinas: El enfoque proporciona un puente metodológico entre el diseño de muestreo epidemiológico clásico (casos y controles anidados) y el aprendizaje automático moderno para el análisis de supervivencia.
2. Utilidad Práctica: El cálculo del riesgo completo a menudo ofrece poca ventaja práctica para lograr un buen rendimiento predictivo en el modelado RSF, mientras que incurre en costes computacionales significativamente mayores.
3. Accesibilidad: Al implementar esta metodología en el ampliamente utilizado paquete ranger, los autores facilitan la aplicación de RSF escalables en la práctica, particularmente para datos de supervivencia de gran escala y alta dimensionalidad.

Los autores concluyen que, si bien el método es robusto, trabajos futuros podrían explorar su aplicación en marcos más complejos (riesgos competitivos, modelos de múltiples estados) e investigar las interacciones con otros hiperparámetros de RSF.