X-FEMR: A Token-level Explainable Approach for Electronic Health Records Foundation Models using Transformer-based Models

Resumen Técnico: X-FEMR

Planteamiento del Problema
Los Modelos Fundacionales para Registros de Salud Electrónicos (FEMRs, por sus siglas en inglés), como ETHOS y CLMBR-T-Base, han demostrado un sólido rendimiento al convertir trayectorias longitudinales de pacientes en representaciones generalizables para tareas de predicción clínica. Sin embargo, estos modelos funcionan como "cajas negras", lo que plantea preocupaciones significativas sobre el sesgo, la interpretabilidad y la confianza clínica. Las técnicas existentes de IA Explicable (XAI), incluyendo la atribución de características (p. ej., SHAP) y los métodos basados en gradientes (p. ej., mapas de saliencia), a menudo no son adecuadas para arquitecturas de escala fundacional. Estos métodos tienen dificultades con las representaciones distribuidas de alta dimensión, las interacciones complejas de contexto largo y la naturaleza agnóstica a la tarea de los espacios latentes de los FEMR. Por consiguiente, existe una falta de marcos de explicabilidad a nivel de token diseñados específicamente para los FEMR que puedan alinearse con el razonamiento clínico. Los modelos sustitutos actuales (p. ej., regresión lineal, árboles de decisión) a menudo asumen relaciones de datos lineales y suaves, lo cual no logra capturar la dinámica temporal irregular, heterogénea y no lineal inherente a los datos de EHR. En consecuencia, hay una carencia de marcos de explicabilidad a nivel de token diseñados específicamente para los FEMR que puedan alinearse con el razonamiento clínico.

Metodología
Los autores proponen X-FEMR, el primer enfoque de explicabilidad a nivel de token para los FEMR, utilizando un modelo sustituto basado en transformadores (transformer-based). La metodología consta de tres componentes primarios:

1. Recolección y Preprocesamiento de Datos:

   • El estudio utiliza el conjunto de datos EHRSHOT, organizando los registros de pacientes en secuencias de eventos clínicos ordenadas en el tiempo siguiendo el Modelo de Datos Común OMOP (OMOP-CDM).
   • Cada evento se caracteriza por cuatro características: un código médico (p. ej., SNOMED, LOINC), un valor numérico (p. ej., resultados de laboratorio), un campo de texto y una marca de tiempo.
   • Los pares de entrada-salida se extraen del FEMR objetivo (CLMBR-T-Base) para dos tareas de predicción específicas: la predicción de la Estancia Hospitalaria (LOS) y la transferencia a la UCI.

2. Desarrollo del Modelo Sustituto:

   • Se entrena un modelo sustituto basado en Transformer para aproximar el comportamiento de entrada-salida del FEMR. A diferencia de los sustitutos lineales tradicionales, esta arquitectura captura dependencias temporales complejas y patrones no lineales.
   • Codificación de Eventos: Cada evento clínico se codifica mediante embeddings de códigos aprendidos, embeddings de texto y un mecanismo de inyección de compuerta (gated injection) para los valores numéricos. Las dinámicas temporales se capturan calculando el intervalo de tiempo hasta la predicción ($\Delta t$), el cual es transformado logarítmicamente y normalizado.
   • Estrategias de Supervisión: El sustituto se entrena bajo dos configuraciones:
     • Supervisión de etiquetas duras (Hard-label Supervision): Predicción de salidas binarias (p. ej., transferencia vs. no transferencia) para capturar la frontera de decisión final.
     • Supervisión de etiquetas blandas (Soft-label Supervision): Predicción de salidas de probabilidad bruta para capturar la confianza del modelo y la lógica de decisión matizada.
   • El estudio compara estas configuraciones para determinar cuál refleja mejor la lógica de decisión subyacente del FEMR.

3. Explicabilidad a Nivel de Token y Evaluación:

   • Análisis SHAP: El modelo sustituto entrenado se analiza mediante SHAP (SHapley Additive exPlanations) para cuantificar la contribución de los tokens individuales (eventos clínicos) a las predicciones del modelo.
   • Métrica de Alineación Clínica: Para evaluar la relevancia clínica de estas explicaciones, los autores introducen una métrica novedosa llamada Relación de Eventos Validados Clínicamente (R). Esta métrica cuantifica la correspondencia entre los tokens identificados como importantes por el sustituto y un conjunto de características clínicamente validadas (p. ej., códigos ICD-10 específicos, signos vitales como frecuencia cardíaca y presión arterial) establecidas en la literatura médica.
   • $R = \frac{\sum_{f_j \in F_c} c(f_j; D)}{\sum_{f_j \in F} c(f_j; D)}$, donde $F_c$ es el subconjunto de tokens clínicamente relevantes y $c(f_j; D)$ es la frecuencia de ocurrencia del token $f_j$ en el conjunto de datos.

Resultados Clave

• Rendimiento del Sustituto: El modelo sustituto basado en transformer entrenado con supervisión de etiquetas duras aproximó estrechamente el comportamiento predictivo del FEMR, superando al CLMBR-T-Base original en la tarea de LOS (AUPRC: 0.4434 vs. 0.4173) y logrando un rendimiento comparable en la tarea de transferencia a la UCI. La supervisión de etiquetas blandas produjo un rendimiento inferior, probablemente debido a la dificultad de modelar señales positivas débiles cerca del umbral de decisión.
• Importancia de los Tokens: El análisis SHAP reveló que los valores numéricos (p. ej., mediciones de laboratorio) y los deltas de tiempo fueron las características más influyentes, contribuyendo de manera más significativa que el texto o los códigos.
• Alineación Clínica: Los tokens identificados como importantes se alinearon bien con el conocimiento clínico.
  • Para la tarea de LOS, los 5 tokens más frecuentes (Frecuencia cardíaca, SBP, Temperatura corporal, DBP, Frecuencia respiratoria) representaron el 29% de los eventos totales, con una Relación de Eventos Validados Clínicamente ($R$) de 0.3093.
  • Para la tarea de Transferencia a la UCI, los 5 tokens principales (Frecuencia cardíaca, Temperatura corporal, SBP, DBP, Saturación de oxígeno) representaron el 24% de los eventos totales, con una $R$ de 0.2717.
  • Estos resultados indican que las explicaciones del modelo sustituto se concentran en los predictores clínicamente validados de los resultados de los pacientes.

Significancia y Reivindicaciones
El artículo afirma presentar el primer enfoque de explicabilidad a nivel de token específicamente para los FEMR. Su importancia radica en:

1. Cerrar la Brecha de Interpretabilidad: Va más allá de las explicaciones tradicionales a nivel de característica para proporcionar información granular a nivel de token sobre cómo los FEMR aprovechan elementos específicos del historial del paciente (p. ej., signos vitales específicos en momentos específicos) para las predicciones.
2. Marco Validado: Al introducir la Relación de Eventos Validados Clínicamente, los autores proporcionan un método cuantitativo para evaluar si las explicaciones del modelo se alinean con el conocimiento médico establecido, abordando la falta de métricas de evaluación para la explicabilidad de los FEMR.
3. Utilidad Práctica: El enfoque ofrece un marco práctico para generar IA clínica interpretable y confiable, ayudando a los clínicos a comprender el razonamiento detrás de las recomendaciones de la IA y potencialmente mejorando los métodos de puntuación clínica o las guías de práctica.

Los autores reconocen limitaciones, incluyendo la complejidad de imitar trayectorias de pacientes largas y ruidosas, y la posible pérdida de información detallada al combinar múltiples canales de características en un solo token de evento. Sugieren que trabajos futuros podrían explorar estrategias de tokenización alternativas que traten el tiempo como un componente de modelado distinto.