Towards end-to-end LLM-based censoring-aware survival analysis

Resumen Técnico: LLMSurvival

Enunciado del Problema
El análisis de supervivencia es una piedra angular de la predicción médica, que tiene como objetivo modelar tanto la ocurrencia de un evento como el tiempo hasta su ocurrencia. Sin embargo, la aplicación de Modelos de Lenguaje Grandes (LLM) como modelos de supervivencia de extremo a extremo se ha visto obstaculizada por el desafío de la censores. En los datos censurados, el tiempo exacto de un evento es desconocido porque un sujeto puede retirarse o el período de observación termina antes de que ocurra el evento. Esta falta de etiquetas de verdad fundamental explícitas impide el uso directo de los marcos estándar de ajuste fino supervisado (SFT) y aprendizaje por refuerzo, que dependen de objetivos definitivos. En consecuencia, la investigación previa ha relegado en gran medida a los LLM al papel de extractores de características o ha requerido modificaciones arquitectónicas sustanciales (por ejemplo, añadir redes neuronales sobre los LLM) en lugar de utilizar el LLM como el motor predictivo principal.

Metodología: LLMSurvival
Los autores proponen LLMSurvival, un marco que permite el análisis de supervivencia consciente de la censura utilizando LLMs sin modificar que operan directamente sobre datos clínicos tabulares. La innovación central es la reformulación de la predicción del tiempo hasta el evento como un problema de clasificación por pares, evitando la necesidad de etiquetas de tiempo absolutas.

1. Conversión de Datos: Las variables tabulares estructuradas (por ejemplo, edad, signos vitales, resultados de laboratorio) se convierten en declaraciones en lenguaje natural utilizando una plantilla (por ejemplo, "La edad es 79; El género es FEMENINO").
2. Formulación por Pares: En lugar de predecir un tiempo de supervivencia específico, el modelo se entrena para determinar cuál de dos sujetos comparables experimentará el evento primero.
   • Pares de Entrenamiento: Se construye un par de entrenamiento válido $(i, j)$ donde el sujeto $i$ experimentó el evento ($e_i=1$) y lo hizo antes del tiempo de seguimiento del sujeto $j$ ($t_i < t_j$). Esto incluye sujetos censurados siempre que hayan sido observados durante más tiempo que el tiempo del evento del otro sujeto.
   • Objetivo: El modelo se ajusta fino mediante aprendizaje supervisado para predecir el ordenamiento relativo. La función de pérdida es un objetivo de entropía cruzada binaria sobre pares comparables válidos, optimizando la probabilidad de que el modelo asigne una puntuación de riesgo más alta al sujeto que experimenta el evento antes.
   • Muestreo: Para gestionar el costo computacional de generar todos los pares posibles, se emplea una estrategia de muestreo de casos y controles anidados, seleccionando un número fijo de controles (sujetos con seguimiento más largo) para cada caso de evento.
3. Inferencia mediante Comparación con Anclaje: Durante las pruebas, la puntuación de riesgo de un sujeto no se predice directamente. En su lugar, el sujeto de prueba se compara contra un conjunto fijo de $K$ individuos "ancla" muestreados de la cohorte de entrenamiento.
   • El modelo predice si el sujeto de prueba tiene un riesgo mayor que cada ancla.
   • La puntuación de riesgo final es la proporción de anclas contra las cuales el sujeto de prueba se clasifica como de mayor riesgo.
4. Modelos: El marco se evaluó utilizando dos LLMs ajustados para instrucciones y desplegable localmente: Llama 3.1-8B-Instruct y Qwen 2.5-7B-Instruct.

Contribuciones Clave

• Modelado de Supervivencia de LLM de Extremo a Extremo: El artículo presenta una prueba de concepto para entrenar LLMs directamente como modelos de supervivencia sin modificaciones arquitectónicas, superando su papel tradicional como extractores de características.
• Resolución de Censura: Al reformular el problema como una tarea de clasificación, el método incorpora naturalmente las observaciones censuradas en el proceso de entrenamiento, abordando el obstáculo fundamental de los tiempos de evento faltantes.
• Portabilidad: Se demuestra que el marco es portable a diversos contextos clínicos (mortalidad aguda en UCI y riesgo de fractura crónica) sin cambios arquitectónicos específicos de la tarea.
• Despliegue Local: El uso de modelos base compactos y públicamente disponibles (7B/8B parámetros) permite el despliegue local, mitigando las preocupaciones de privacidad asociadas con las llamadas a APIs externas.

Resultados
El marco se evaluó en dos tareas clínicas distintas:

1. Predicción de Mortalidad en UCI (MIMIC-IV):
   • LLMSurvival logró un índice C global de 0.776 (Llama) y 0.784 (Qwen).
   • Esto representa una mejora del 3.1% sobre el modelo de riesgos proporcionales de Cox (0.753) y una mejora del 3.0% sobre la puntuación clínica establecida SAPS-II (0.746).
   • Superó a tres líneas base de aprendizaje profundo (Cox-nnet, DeepHit, DeepSurv) en promedio.
2. Predicción de Fracturas por Fragilidad (Cohorte NYP/WCM):
   • LLMSurvival logró un índice C global de 0.742 (Llama) y 0.741 (Qwen).
   • Esto mostró una mejora del 0.5% sobre el modelo de Cox y una mejora significativa del 8.2% sobre la puntuación clínica FRAX (0.660).
   • El rendimiento fue competitivo o superior a las líneas base de aprendizaje profundo en varios horizontes temporales (2, 3 y 5 años).

En ambas tareas, los análisis de Kaplan–Meier confirmaron que LLMSurvival estratificó efectivamente a los pacientes en grupos de alto y bajo riesgo con trayectorias de supervivencia distintas (Razones de Peligro de 4.80 para UCI y 4.02 para fracturas). El análisis cualitativo de las salidas del modelo indicó que los LLMs generaron razonamientos clínicamente coherentes para sus comparaciones por pares, citando factores de riesgo relevantes como la gravedad de la admisión, los niveles de oxigenación y la carga de comorbilidades.

Significado y Afirmaciones
Los autores afirman que LLMSurvival demuestra que el modelado de supervivencia con censura puede hacerse compatible con el ajuste fino de LLM mediante una reformulación basada en comparaciones. El estudio sirve como una prueba de concepto de que los LLMs pueden expandir sus capacidades para resolver problemas estadísticos complejos de forma independiente, en lugar de depender de arquitecturas híbridas. La capacidad del marco para superar las puntuaciones clínicas curadas por expertos (SAPS-II y FRAX) sugiere que los LLMs pueden aprender una estratificación de riesgos significativa directamente a partir de los datos. Además, el trabajo ilustra un principio más amplio: reformular problemas estadísticamente desafiantes en tareas alineadas con las capacidades de los LLM (como el razonamiento comparativo) puede extender la aplicabilidad de los modelos de lenguaje en el aprendizaje automático biomédico.

Limitaciones
Los autores reconocen varias limitaciones:

• El entrenamiento mostró cierta variabilidad entre las configuraciones de muestreo, lo que sugiere margen para mejoras en la estabilidad de la optimización.
• La selección de pares dependió de un muestreo aleatorio de controles; estrategias estructuradas podrían mejorar la eficiencia.
• La inferencia requiere múltiples comparaciones por pares contra anclas, introduciendo una sobrecarga computacional.
• La implementación actual utiliza únicamente entradas tabulares estructuradas; se necesita trabajo futuro para integrar señales multimodales como notas clínicas e imágenes médicas.