Resumen Técnico: Use What You Know: Causal Foundation Models with Partial Graphs

Declaración del Problema

La inferencia causal tradicionalmente depende de un proceso de dos pasos: primero, establecer supuestos de identificación (a menudo mediante un grafo causal) y, segundo, seleccionar un estimador a medida adaptado a esos supuestos. Aunque los recientes Modelos de Fundación Causales (CFM, por sus siglas en inglés) basados en Redes de Ajuste de Datos Previos (PFN) ofrecen un enfoque unificado al amortizar el descubrimiento e inferencia causal, actualmente sufren una limitación crítica: no pueden incorporar conocimiento de dominio en el momento de la prueba.

Los CFM existentes:

1. Restringen el prior de entrenamiento a estructuras causales específicas (por ejemplo, solo casos de ajuste de puerta trasera/back-door), lo que requiere un reentrenamiento costoso para nuevos supuestos.
2. Marginalizan sobre todas las estructuras posibles, lo que conduce a estimaciones innecesariamente conservadoras y predicciones imprecisas cuando hay conocimiento causal parcial disponible pero se ignora.

El desafío central abordado es cómo permitir que un CFM de propósito general se condicione a información causal parcial (como un grafo causal parcialmente conocido o relaciones ancestrales) durante la inferencia sin sacrificar la flexibilidad del paradigma de los modelos de fundación.

Metodología

1. Representación del Conocimiento Causal: Matrices Ancestrales Parciales (PAMs)

Los autores proponen el uso de Matrices Ancestrales Parcialmente Conocidas (PAMs) como la señal de condicionamiento. A diferencia de las matrices de adyacencia que requieren el conocimiento de los aristas causales directas ($A_{ij}=1$), las PAMs codifican relaciones ancestrales ($T_{ij}=1$ si $z_i$ es un ancestro de $z_j$).

• Ventajas: El conocimiento ancestral suele ser más robusto y fácil de obtener que los mecanismos causales directos (por ejemplo, saber que fumar causa cáncer es más fácil que especificar la vía biológica exacta).
• Parcialidad: La matriz $\tilde{T} \in \{-1, 0, 1\}^{K \times K}$ permite:
  • $1$: Ancestro conocido.
  • $-1$: No ancestro conocido (por ejemplo, descartado por orden temporal).
  • $0$: Relación desconocida.
    Esta representación subsume los casos de ningún conocimiento, conocimiento parcial y conocimiento completo.

2. Mecanismos de Condicionamiento Arquitectónico

El artículo introduce una arquitectura de CFM basada en Transformer modificada para ingerir las PAMs a través de dos mecanismos complementarios:

• Sesgo de Atención Estructural (Structural Attention Biasing):
  • Sesgo Suave (Soft Bias): Sesgos escalares aprendibles ($\beta_{anc}, \beta_{non-anc}$) que se añaden a los logits de atención. Un sesgo positivo fomenta la atención hacia los ancestros conocidos; un sesgo negativo desincentiva la atención hacia los no ancestros conocidos. Esto permite al modelo priorizar a los padres causales manteniendo la robustez ante información incompleta.
  • Enmascaramiento Duro (Hard Masking): Una alternativa donde la atención a los no ancestros conocidos se establece estrictamente en $-\infty$.
• Modulación de Características Estructurales (GCN):
  • Un codificador de Redes Convolucionales de Grafos (GCN) procesa la PAM para generar incrustaciones (embeddings) de nodos conscientes de la estructura.
  • Estas incrustaciones se inyectan en cada bloque del Transformer mediante Normalización de Capa Adaptativa (AdaLN), modulando las representaciones ocultas de las variables según su posición estructural global en el grafo causal.

3. Entrenamiento y Generación de Datos

• Prior: Los autores construyeron un prior de Modelo de Fundación Tabular (TFM) que es nativamente causal. Este genera conjuntos de datos sintéticos muestreando Modelos Causales Estructurales (SCMs) con diversos mecanismos (MLP, XGBoost, diversas no linealidades) y distribuciones de ruido.
• Entrenamiento Amortizado: El modelo se entrena para manejar distintos grados de información faltante. Durante el entrenamiento, las entradas en la PAM se ocultan aleatoriamente (se establecen en 0), obligando al modelo a aprender un único conjunto de pesos que pueda adaptarse a la información ancestral completa, parcial o nula.

Resultados Clave

1. Efectividad del Condicionamiento

Experimentos en priors lineales-Gaussianos y no lineales complejos demuestran que:

• Soft Attention + GCN es la estrategia de condicionamiento más efectiva, superando consistentemente a los baselines sin condicionamiento de grafo, de enmascaramiento duro y de enfoques de solo GCN.
• El condicionamiento en matrices ancestrales tiene un rendimiento comparable al condicionamiento en matrices de adyacencia completas, validando el uso del conocimiento ancestral como un proxy práctico.

2. Robustez a la Información Parcial

• Eficiencia de la Amortización: Un único modelo entrenado con PAMs amortizadas (ocultas aleatoriamente) funciona casi tan bien como los modelos especializados entrenados exclusivamente con información completa o nula.
• Ignorancia vs. Errores de Especificación: Un hallazgo crítico es que ocultar las aristas inciertas (marcándolas como "desconocidas") es mucho menos perjudicial que especificar erróneamente (cambiar la dirección). El rendimiento se degrada gradualmente a medida que se ocultan las aristas, pero cae abruptamente cuando se especifican incorrectamente. Esto resalta el valor práctico de las matrices parciales.

3. Benchmarks Semi-Sintéticos

En el benchmark RealCause (IHDP, ACIC, CPS, PSID):

• Proporcionar información ancestral parcial que codifica la unconfoundedness (ausencia de confusión) mejora significamente el rendimiento (medido por $\sqrt{PEHE}$ y el error relativo de ATE) en comparación con modelos sin condicionamiento de grafo.
• El CFM condicionado por grafo supera a los baselines especializados como Do-PFN y CausalFM en la mayoría de los datasets.
• Aunque CausalPFN logra el mejor rendimiento general en RealCause, los autores atribuyen esto a un alineamiento del prior (los datos de entrenamiento de CausalPFN imitan el proceso generador de datos de RealCause) más que al condicionamiento de grafo. Una vez que se introduce el condicionamiento de grafo, su modelo supera a CausalFM y Do-PFN.

4. Comparación con Modelos Especializados

Cuando se viola el criterio de puerta trasera (un escenario donde los modelos especializados como CausalPFN y CausalFM dependen de supuestos mal especificados), el CFM condicionado por grafo muestra una ventaja sustancial, demostrando su capacidad para manejar casos no identificables aprovechando la estructura parcial disponible.

Significación y Reivindicaciones

El artículo afirma abordar un obstáculo central en el desarrollo de modelos de fundación causales "todo en uno": la capacidad de responder consultas causales de manera basada en datos y aprovechando eficazmente cualquier cantidad de experiencia de dominio.

• Marco Unificado: Demuestra que un único CFM puede condicionarse de manera flexible a datos observacionales y a niveles variables de conocimiento causal (ninguno, parcial o completo) sin necesidad de reentrenamiento.
• Utilidad Práctica: Al utilizar Matrices Ancestrales Parciales, este enfoque cierra la breancia entre la rigidez de los estimadores causales tradicionales (que requieren grafos precisos) y la naturaleza de "caja negra" de los modelos de fundación actuales (que ignoran el conocimiento de dominio).
• Alcance Modesto: Los autores reconocen que, si bien las ganancias de rendimiento absoluto en algunas métricas son pequeñas (atribuidas a los pequeños efectos causales en sus priors), el método establece un fundamento factible y efectivo para futuros modelos de estimación de efectos causales robustos. Señalan que el trabajo actual se centra en variables totalmente observadas y aún no aborda los confundidores ocultos, dejando eso para desarrollos futuros "todo en uno".