Resumen Técnico: Modelos de Lenguaje de Gran Escala Aumentados por Grafos de Conocimiento para la Predicción de Enfermedades

Planteamiento del Problema

Las Historias Clínicas Electrónicas (HCE) contienen datos longitudinales adecuados para la predicción clínica; sin embargo, los modelos existentes a menudo no logran proporcionar un razonamiento a nivel de paciente que sea accionable para los clínicos. Los modelos clásicos (p. ej., regresión logística, ensambles de árboles) ofrecen robustez e interpretabilidad global, pero típicamente solo proporcionan explicaciones post-hoc toscas (p. ej., puntuaciones de importancia de características) que carecen de la narrativa paso a paso requerida para decisiones clínicas específicas. Por el contrario, aunque los Modelos de Lenguaje de Gran Escala (LLM) pueden generar narrativas de Cadena de Pensamiento (CoT), su aplicación no restringida en medicina se ve obstaculizada por la escasez de conjuntos de datos de razonamiento de alta calidad y validados clínicamente, así como por el riesgo de pasos intermedios alucinados o médicamente incorrectos. El desafío central es permitir que los LLM realicen una predicción de enfermedades precisa a nivel de visita, generando simultáneamente explicaciones fundamentadas fácticamente en el conocimiento biomédico y temporalmente consistentes con la historia del paciente.

Metodología

Los autores proponen un marco de CoT guiado por Grafos de Conocimiento (KG) diseñado para estructurar el razonamiento del LLM utilizando la estructura de PrimeKG, un grafo de conocimiento biomédico a gran escala. El flujo de trabajo consta de cuatro etapas primarias:

1. Mapeo de Entidades del KG:

   • Los códigos ICD-9 de MIMIC-III se alinean con los nodos de PrimeKG mediante un procedimiento de tres etapas:
     • Coincidencia Exacta: Emparejamiento directo de cadenas entre las descripciones de ICD-9 y las etiquetas de los nodos del KG.
     • Coincidencia por Similitud: Recuperación basada en embeddings (similitud de coseno > 0.85) para coincidencias no exactas.
     • Filtrado por LLM: GPT-4o valida o revisa los mapeos provisionales para asegurar la corrección clínica.
   • Este proceso genera 1,513 nodos de características mapeados y 10 nodos de enfermedades mapeados.

2. Minería de Relevancia y Rutas del KG:

   • Para cada enfermedad objetivo, el sistema identifica un conjunto compacto de nodos de características relevantes ($R_d$) utilizando GPT-4o.
   • Rutas de Razonamiento: Se extraen rutas cortas (limitadas a $L=5$ saltos) entre los nodos de características relevantes y el nodo de la enfermedad objetivo dentro de PrimeKG.
   • Selección de Rutas: GPT-4o filtra estas rutas para retener solo aquellas que son clínicamente relevantes y mecánicamente sólidas, podando cadenas espurias.

3. Generación y Filtrado de CoT:

   • Se le solicita al LLM que genere explicaciones de CoT ($C_{t,d}$) para una visita específica $t$, incorporando las características de ICD-9 observadas, el conjunto de relevancia específico de la enfermedad y las rutas de KG minadas.
   • Filtrado de Consistencia de Etiquetas: Los trazos de CoT generados se retienen solo si la conclusión binaria explícita ($\hat{y}$) derivada del razonamiento coincide con el resultado real ($y$). Esto crea un corpus de supervisión de alta calidad y fundamentado en el KG.

4. Ajuste Fino del Modelo (Fine-Tuning):

   • Modelos ligeros de pesos abiertos (LLaMA-3.1-Instruct-8B y Gemma-7B) se ajustan finamente sobre este corpus curado.
   • En la inferencia, estos modelos toman las características de ICD-9 y la evidencia del KG específica de la enfermedad como entrada para producir tanto una predicción de enfermedad como un trazo de razonamiento al estilo clínico.

Contribuciones Clave

• Marco de CoT Guiado por KG: Un método novedoso que integra la estructura del KG biomédico directamente en el proceso de generación de CoT, yendo más allá de las explicaciones libres no restringidas hacia un razonamiento clínicamente fundamentado.
• Supervisión Eficiente en Datos: La construcción de un corpus de supervisión compacto y de alta calidad mediante el filtrado de consistencia de etiquetas permite que los LLM pequeños (8B de parámetros) aprendan razonamiento pronóstico con datos limitados (400–1,000 muestras).
• Transferencia Zero-Shot: El marco demuestra la capacidad de transferir los patrones de razonamiento aprendidos a una cohorte distinta (CRADLE) sin reentrenamiento, manteniendo una alta precisión en el pronóstico de eventos cardiovasculares tras el diagnóstico de diabetes.
• Evaluación Centrada en el Clínico: Un estudio humano ciego valida que las explicaciones guiadas por KG son preferidas sobre las de los modelos base en términos de claridad, relevancia y corrección clínica.

Resultados Experimentales

El marco fue evaluado en la cohorte MIMIC-III (10 enfermedades mapeadas a PrimeKG) y la cohorte CRADLE (Diabetes Tipo 2 a eventos cardiovasculares).

• Desempeño en MIMIC-III:

  • En entornos de bajos datos (400 y 1,000 visitas de índice), el ajuste fino de CoT guiado por KG mejoró el desempeño de LLaMA-3.1-8B y Gemma-7B, logrando resultados competitivos o superiores con bases clásicas fuertes (XGBoost, Random Forest, Regresión Logística).
  • LLaMA-3.1 (1,000 muestras): Alcanzó un AUROC de 0.6995 y un macro-AUPR de 0.4685, superando ligeramente al mejor modelo base clásico (Regresión Logística: AUROC 0.6874, AUPR 0.4178; XGBoost: AUROC 0.6799).
  • Gemma-7B: Mostró mejoras similares, con un macro-AUPR que aumentó de 0.25 (original) a 0.4363 (ajustado con KG-CoT), aunque su AUROC (0.6609) se mantuvo ligeramente por debajo de los mejores modelos clásicos.
  • Ablación: Eliminar la guía del KG o el filtrado redujo el desempeño, confirmando que tanto el fundamento estructural como el filtro de consistencia de etiquetas son críticos.

• Transferencia Zero-Shot a CRADLE:

  • Los modelos ajustados con MIMIC-III se aplicaron a CRADLE sin reentrenamiento.
  • La precisión mejoró de 0.40–0.51 (sin ajuste) a 0.72–0.77 (ajustado con KG-CoT).
  • Aunque el AUROC y el AUPR vieron ganancias modestas, el salto significativo en la precisión sugiere que los modelos aprendieron una frontera de decisión más robusta que generaliza a través de cohortes.

• Evaluación Humana:

  • En un estudio ciego con 115 casos, los clínicos prefirieron el modelo guiado por KG sobre el modelo base sin ajustar en el 96.5% de los casos para Claridad/Coherencia, el 94.8% para Cobertura/Relevancia y el 98.3% para Corrección/Solidez.
  • Los clínicos señalaron que los trazos guiados por KG estaban mejor organizados temporalmente, resaltaban factores de riesgo significativos y evitaban saltos sin sustento.

Significancia y Reivindicaciones

El artículo sostiene que la supervisión de CoT anclada en KG permite que los LLM pequeños de pesos abiertos aprendan un razonamiento pronóstico eficiente en datos y clínicamente fundamentado. La significancia radica en cerrar la brecha entre la predicción precisa y la justificación interpretable:

1. Fiabilidad: Al fundamentar el razonamiento en PrimeKG y filtrar por consistencia de etiquetas, el sistema mitiga las alucinaciones y asegura la plausibilidad médica.
2. Generalizabilidad: Los patrones de razonamiento aprendidos se transfieren efectivamente a cohortes fuera de la distribución (CRADLE), lo que sugiere que los modelos capturan mecanismos de enfermedad subyacentes en lugar de artefactos específicos del conjunto de datos.
3. Utilidad Clínica: El marco produce explicaciones que los clínicos encuentran más claras y confiables, facilitando potencialmente la adopción de los LLM en sistemas de apoyo a la decisión clínica donde la transparencia es primordial.

Los autores señalan modestamente que el enfoque depende de la cobertura de PrimeKG y que trabajos futuros deberían explorar vistas de EHR más ricas (notas, laboratorios) y contextos de múltiples visitas.