Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction

El artículo presenta CausalLongPFN, una red de ajuste previo congelada preentrenada en modelos causales temporales sintéticos que predice resultados longitudinales contrafácticos sin entrenamiento específico del dominio, demostrando un rendimiento competitivo en evaluaciones como cáncer, VIH y MIMIC-III, al tiempo que ofrece una alternativa rentable al modelado tradicional específico de cohortes.

Lo esencial

Imagina que eres un médico intentando predecir cómo responderá un paciente a un plan de tratamiento específico durante las próximas semanas. El problema es que los pacientes son complejos: su salud cambia día tras día, sus tratamientos pasados afectan su estado actual y otros factores (como la dieta o el estrés) interfieren en los resultados. Normalmente, para hacer estas predicciones, tienes que construir un modelo informático nuevo y altamente especializado para cada nuevo grupo de pacientes que encuentras. Esto es como contratar a un nuevo arquitecto para diseñar una casa cada vez que te mudas a un nuevo vecindario. Es lento, costoso y requiere mucha información.

Este artículo presenta una nueva herramienta llamada CAUSALLONGPFN. Piensa en ella como un "Motor de Intuición Médica Universal" que ya ha aprendido las reglas del juego antes de ver a un paciente real.

Así es como funciona, desglosado en conceptos sencillos:

1. El "Campo de Entrenamiento" (Preentrenamiento Sintético)

En lugar de esperar a que lleguen los pacientes reales, los creadores construyeron un enorme "campo de entrenamiento" virtual. En este campo, simularon millones de pacientes falsos con millones de tipos de cuerpos, enfermedades y reacciones a los tratamientos diferentes. Programaron a estos pacientes falsos para que tuvieran comportamientos complejos:

• Algunos mejoran lentamente; otros colapsan rápidamente.
• Algunos tratamientos funcionan de inmediato; otros tienen un efecto retardado.
• Algunos pacientes reaccionan de manera diferente según su genética oculta.

El modelo de IA pasó todo su tiempo en este campo, aprendiendo a predecir resultados para estos pacientes falsos. No se limitó a memorizar respuestas; aprendió la lógica subyacente de cómo interactúan los tratamientos, el tiempo y la biología.

2. El "Experto Congelado" (No se necesita un nuevo entrenamiento)

Aquí está el truco de magia: una vez que el modelo terminó su campo de entrenamiento, lo congelaron. Bloquearon su cerebro. No puede aprender nada nuevo ni cambiar sus ajustes internos.

Cuando llega un nuevo grupo de pacientes reales (como pacientes con cáncer en un hospital), el modelo no empieza desde cero. No necesita ser reentrenado. En su lugar, actúa como un superpasante que lee la historia clínica.

• Las Trayectorias de Soporte: Le muestras al modelo algunos ejemplos de pacientes del hospital actual (el "soporte"). Estos son como darle al pasante unos pocos expedios de casos para que entienda el estilo específico de este hospital.
• La Consulta: Preguntas: "¿Qué le pasará a este paciente específico si le damos el Tratamiento A durante los próximos 5 días?".
• La Respuesta: El modelo utiliza instantáneamente lo que aprendió en el campo de entrenamiento, combinado con el contexto de los expedientes de casos que le mostraste, para predecir el resultado. Hace esto sin dar ni un solo paso de "descenso de gradiente" (un término técnico para el proceso habitual de reentrenamiento).

3. El "Simulador de Viaje en el Tiempo"

El modelo está diseñado para manejar datos longitudinales, lo que significa que entiende el tiempo. No solo adivina el resultado de mañana; simula el futuro paso a paso.

• Predice el Día 1.
• Toma esa predicción y la utiliza como punto de partida para el Día 2.
• Repite este proceso para ver qué sucede en el Día 5, 6 o 7.

Esto es como un simulador de vuelo que no solo muestra el despegue, sino que simula todo el recorrido del vuelo basado en las decisiones del piloto, incluso si el clima cambia a mitad del trayecto.

4. Por qué esto importa (Los Resultados)

Los autores probaron este "experto congelado" contra la forma antigua de hacer las cosas (construir un nuevo modelo para cada hospital).

• La Prueba: Utilizaron datos de cáncer, VIH, anticoagulantes (warfarina) y registros reales de la UCI.
• El Resultado: El modelo congelado funcionó tan bien como, y a veces mejor que, los modelos que fueron entrenados específicamente para cada enfermedad.
• La Gran Victoria: Funcionó especialmente bien con los datos reales de la UCI, donde no podían probar escenarios de "¿qué pasaría si?" (porque no es ético probar diferentes tratamientos en pacientes reales en una simulación). El modelo podía predecir qué es probable que suceda a continuación basándose únicamente en los datos observados.

La Conclusión

El artículo afirma que no siempre es necesario construir un modelo nuevo y personalizado para cada nuevo conjunto de datos médicos. En su lugar, puedes entrenar un modelo masivo en una amplia variedad de escenarios de "¿qué pasaría si?" (datos sintéticos) y luego usarlo como una herramienta congelada y lista para usar.

Es como tener un maestro chef que ha practicado cocinando todo tipo de cocina en una cocina virtual. Cuando traes un nuevo conjunto de ingredientes (un nuevo grupo de pacientes), no necesitan aprender a cocinar de nuevo; simplemente miran tus ingredientes e inmediatamente saben cómo preparar el plato, utilizando su vasta intuición previamente aprendida.

Nota Importante: Los autores aclaran cuidadosamente que esta herramienta es para predicción e investigación, no para tomar decisiones finales de vida o muerte en una clínica. Ayuda a los médicos a comprender los posibles resultados, pero sigue dependiendo de los mismos supuestos médicos estándar en los que cualquier otro modelo causal se apoya. Es una poderosa herramienta de investigación, no un reemplazo del juicio de un médico.

Para quienes deseen explorar la implementación o utilizar el modelo, el código está disponible en GitHub: https://github.com/Amirhossein-Zare/causal-long-pfn y los pesos del modelo preentrenado en Hugging Face: https://huggingface.co/Amirhossein-Zare/causal-long-pfn .

Resumen técnico

Resumen Técnico: Redes de Ajuste de Prior Longitudinales Causales para la Predicción de Resultados Contrafácticos

Formulación del Problema

El artículo aborda el desafío de predecir resultados potenciales bajo secuencias de tratamientos futuros en datos observacionales longitudinales. La tarea central es estimar el resultado potencial condicionado a la historia $E[Y(\bar{a})_{t+\tau} | H_t]$, donde $H_t$ representa la historia observada hasta el tiempo $t$, y $\bar{a}$ es una secuencia de tratamiento planificada.

Este problema se complica por tres factores primarios:

1. Confusión dependiente del tiempo: Las asignaciones de tratamiento en cada paso dependen de covariables que son, a su vez, resultados de tratamientos previos.
2. Dinámicas heterogéneas de pacientes: Las trayectorias individuales exhiben evoluciones de estado complejas y no lineales, así como heterogeneidad latente.
3. Limitaciones de datos: Las cohortes observacionales suelen ser demasiado pequeñas para entrenar modelos de secuencia profundos fiables desde cero para cada nuevo dominio o simulador.

Los estimadores causales longitudinales existentes (por ejemplo, RMSN, CRN, G-Net, Causal Transformer) requieren típicamente un proceso de entrenamiento supervisado separado, que incluye la selección de hiperparámetros y el modelado de la propensión, para cada nueva cohorte. Este flujo de trabajo es costoso e impracticable cuando se requiere un entrenamiento repetido específico para cada dominio.

Metodología: CAUSALLONGPFN

Los autores presentan las Redes de Ajuste de Prior Longitudinales Causales (CAUSALLONGPFN), un predictor de contexto de ajuste de prior diseñado para amortizar la predicción causal longitudinal a través de un amplio prior sobre modelos causales estructurales temporales (TSCM).

1. Preentrenamiento Sintético sobre un Prior Amplio

El modelo es preentrenado enteramente con episodios sintéticos muestreados de un prior amplio sobre TSCMs. Este prior está diseñado para abarcar una amplia clase de dinámicas causales longitudinales en lugar de replicar un único simulador. Las características clave del prior sintético incluyen:

• Grafos Temporales Causales: Dependencias contemporáneas y rezagadas dispersas con grafos instantáneos acíclicos.
• Mecanismos No Lineales: Las coordenadas de estado siguen actualizaciones autorregresivas no lineales dispersas utilizando diversas no linealidades elementales (identidad, tanh, sinusoidal, ReLU, etc.) y diversas distribuciones de ruido.
• Motivos Dinámicos: Se superponen motivos estructurados como memoria de acción, saturación, homeostasis, control de retroalimentación y canales de lectura suavizada para capturar mecanismos cualitativos como efectos retardados y retroalimentación regulatoria.
• Políticas de Comportamiento Confundido: Los tratamientos se muestrean de políticas estocásticas dependientes del estado influenciadas por la heterogeneidad latente de la unidad ($Z_i$), creando una retroalimentación entre tratamiento y confusor dependiente del tiempo.
• Modelos de Resultado: Los resultados escalares se generan mediante lecturas autorregresivas con efectos de tratamiento directos y acumulativos.

2. Arquitectura

CAUSALLONGPFN emplea una arquitectura de doble codificador:

• Codificador de Historia Causal: Un Transformer causal a nivel de trayectoria (que utiliza atención de auto-máscara) que mapea secuencias longitudinales a representaciones de la historia, asegurando que la representación en el tiempo $r$ dependa solo de la información disponible hasta ese momento.
• Codificador de Contexto PFN: Un codificador de contexto basado en Transformer que realiza la adaptación en contexto. Procesa trayectorias de soporte (tratadas como un conjunto desordenado) y un token de consulta de forma conjunta mediante auto-atención. No se asigna codificación posicional al orden de las trayectorias de soporte.
• Cabezal de Predicción de Mezcla Gaussiana: La representación de la consulta final parametriza una distribución de mezcla de 5 componentes para el resultado normalizado, proporcionando tanto predicciones puntuales como estimaciones de incertidumbre.

3. Predicción en Contexto y Despliegue (Rollout)

En el tiempo de prueba, el modelo permanece congelado. Recibe:

• Trayectorias de soporte: Ejemplos del nuevo dominio/tarea.
• Historia de la consulta: Observada hasta el tiempo $t_{obs}$.
• Secuencia de tratamiento futuro propuesta: Las intervenciones planificadas.

El modelo devuelve una distribución predictiva sin actualizaciones de gradiente, ajuste de modelo de propensión o balanceo adversarial. Para la predicción de múltiples pasos ($\tau > 1$), el modelo realiza un despliegue de inserción autorregresivo (plug-in rollout): predice la distribución del resultado de un solo paso, inserta la media de la mezcla en la historia de la consulta y repite el proceso bajo la secuencia de tratamiento especificada.

Contribuciones Clave

1. Un Modelo de Ajuste de Prior para la Predicción Causal Longitudinal: CAUSALLONGPFN es el primer modelo estilo PFN para la predicción de resultados potenciales condicionados a la historia bajo secuencias de tratamiento longitudinales planificadas. Opera como un modelo congelado que no requiere adaptación en el tiempo de prueba.
2. Un Prior Sintético sobre Tareas Causales Longitudinales: Los autores diseñan un prior de TSCM que genera tareas diversas con retroalimentación entre tratamiento y confusor, heterogeneidad latente, dinámicas no lineales, efectos retardados/acumulativos y cambios de régimen.
3. Arquitectura para la Inferencia en Contexto Longitudinal: Un novedoso doble codificador que combina un Transformer causal de historia con un codificador de contexto PFN y un cabezal de mezcla gaussiana.
4. Despliegue Contrafáctico Autorregresivo: Una extensión del predictor de un solo paso aprendido hacia la predicción de múltiples pasos mediante un despliegue de inserción determinista.
5. Evaluación Zero-Shot: Una evaluación exhaustiva de un único modelo congelado frente a líneas de base entrenadas por dominio (MSM, RMSN, G-Net, CRN, Causal Transformer, G-Transformer) en benchmarks de ramas contrafácticas y datos reales factuales.

Resultados

El modelo fue evaluado en cuatro benchmarks: crecimiento de tumores de cáncer, farmacocinética/farmacodinamia (PK/PD) de warfarina, dinámica de tratamiento de VIH y trayectorias de la UCI MIMIC-III.

• Rendimiento Balanceado por Dominio: CAUSALLONGPFN logró el mejor RMSE normalizado de un paso balanceado por dominio (0.222), superando por poco a MSM y RMSN. Para la predicción de cinco pasos, ocupó el tercer lugar general, detrás de RMSN y G-Net, pero superó a MSM, CRN y las líneas de base basadas en transformers.
• Rendimiento por Dominio:
  • MIMIC-III (Factual): El modelo ocupó el primer lugar tanto en la predicción de un paso como en la de cinco pasos, demostrando una fuerte transferencia a trayectorias clínicas del mundo real sin entrenamiento específico del dominio.
  • Benchmarks Contrafácticos (Cáncer, VIH, Warfarina): El modelo se mantuvo competitivo, ocupando el segundo o tercer lugar en tareas de un paso. Sin embargo, en tareas contrafácticas de horizonte largo (por ejemplo, predicción de cáncer de 5 pasos), los modelos recurrentes especializados en el dominio (RMSN, CRN) lograron errores más bajos, lo que sugiere una ventaja cuando existe suficiente dato del dominio objetivo para el ajuste específico.
• Calibración de Incertidumbre: El cabezal de mezcla gaussiana proporcionó información distributiva útil. La calibración varió según el dominio, siendo la warfarina la que mostró mejor calibración y MIMIC-III la que presentó intervalos más amplios debido a la mayor heterogeneidad.

Significado, Reivindicaciones y Disponibilidad

El artículo sostiene que un preentrenamiento de prior sintético amplio puede proporcionar una alternativa congelada útil frente al entrenamiento repetido y específico de dominios cuando:

• El reentrenamiento de modelos especializados es costoso o impracticable.
• Se requiere una adaptación rápida a una nueva cohorte.
• No hay supervisión contrafáctica disponible (como en las tareas de predicción factual del mundo real como MIMIC-III).

Los autores enfatizan que CAUSALLONGPFN no elimina los supuestos causales estándar (consistencia, positividad, intercambiabilidad secuencial) necesarios para interpretar datos observacionales. En su lugar, amortiza el problema de la estimación. Los resultados sugieren que un prior sintético suficientemente amplio puede capturar estructuras reutilizables en tareas de respuesta al tratamiento, convirtiendo al modelo en un fuerte predictor de contexto de propósito general. No obstante, el artículo señala modestamente que el entrenamiento específico de dominio sigue siendo valioso cuando existen suficientes datos y señales de validación del dominio objetivo, particularmente para predicciones de horizonte largo en dominios específicos.

Este trabajo se posiciona como una herramienta de investigación para el modelado de secuencias causales y la generación de hipótesis, más que como un sistema de decisión clínica autónomo, advirtiendo contra la confianza excesiva en las predicciones cuando los supuestos causales o el soporte del prior son inadecuados.

Disponibilidad de Código y Modelos:
La implementación está disponible en GitHub en https://github.com/Amirhossein-Zare/causal-long-pfn y los pesos del modelo preentrenado están disponibles en Hugging Face en https://huggingface.co/Amirhossein-Zare/causal-long-pfn .