Resumen Técnico: TabPFN-CFM

Declaración del Problema

Determinar la estructura causal y predecir los resultados causales a partir de datos observacionales sigue siendo un desafío significativo en el aprendizaje automático y la estadística. Los métodos existentes suelen abordar estas tareas de forma aislada: algunos se centran en las predicciones estructurales (aprendizaje del grafo causal), mientras que otros se centran en las predicciones de resultados (estimación del efecto de las intervenciones). Además, muchos enfoques tienen dificultades para manejar variables de confusión no observadas, lo que puede conducir a grafos no identificables y estimaciones de resultados sesgadas. Adicionalmente, los modelos actuales suelen carecer de la capacidad de soportar todo el espectro de consultas causales definidas por la Jerarquía Causal de Pearl: observacional, intervencional y contrafactual.

Metodología

Los autores presentan TabPFN-CFM, un modelo fundacional causal diseñado para predecir conjuntamente la estructura causal y los resultados a partir de datos observacionales. El enfoque se basa en el marco de Prior sobre Funciones (PFN, por sus siglas en inglés) Bayesiano.

Marco Central

El modelo asume que el sistema subyacente sigue un Modelo Causal Estructural (SCM) $\psi = \{U, V, F\}$, donde $U$ representa las variables no observadas, $V$ las variables observadas y $F$ las ecuaciones estructurales. El modelo está entrenado para aproximar la distribución posterior $P(\psi | D_{obs}, G_{est})$, donde $D_{obs}$ es el conjunto de datos observacionales y $G_{est}$ es el conocimiento previo opcional del grafo.

Generación de Datos y Entrenamiento

• Priors Sintéticos: El modelo se entrena en conjuntos de datos sintéticos generados a partir de una distribución previa sobre SCMs. Este prior incluye estructuras de grafos aleatorios (DAGs), nodos faltantes, diversas distribuciones de ruido y ecuaciones estructurales definidas por redes neuronales con pesos y no linealidades aleatorias.
• Formulación de la Tarea: Cada muestra de entrenamiento consiste en un conjunto de ajuste ($D_{fit}$), una consulta de predicción ($D_{pred}$) y una muestra contrafactual ($D_{causal}$). El modelo está entrenado para manejar tres tipos de consultas simultáneamente:
  1. Consulta Observacional: Predicción de $P(y^* | x^*, T=t^*)$.
  2. Consulta Intervencional: Predicción de $P(y^* | x^*, do(T=1-t^*))$.
  3. Consulta Contrafactual: Predicción de $P(y^* | x^*, do(T=1-t^*), y=y^*)$.
• Manejo de Grafos: El modelo utiliza Grafos Mixtos Dirigidos Acíclicos (ADMGs) para representar la confusión no observada mediante aristas bidireccionales. Esto permite al modelo predecir estructuras de confusión en lugar de asumir que todas las variables son observadas.

Arquitectura

TabPFN-CFM extiende la arquitectura Do-PFN con condicionamiento de grafo explícito:

• Codificador (Encoder): Codifica el conjunto de ajuste, la consulta de predicción y la estructura previa del ADMG (matriz de adyacencia $A$, matriz bidireccional $C$ y matriz ancestral $R$). La codificación del grafo proporciona acceso directo a las relaciones de padres, hijos, ancestros y descendientes.
• Backbone de Transformer: Utiliza mecanismos de atención por filas y columnas. Crucialmente, el módulo de predicción atiende a las codificaciones del grafo, permitiendo que las predicciones incorporen el prior del grafo.
• Salidas (Outputs):
  • Predicción de Resultados: Una cabeza produce logits sobre cubetas (buckets) de resultados discretizados.
  • Predicción de Estructura: Un decodificador genera predicciones elemento a elemento para la matriz de adyacencia dirigida ($\hat{A}$), la matriz de correlación bidireccional ($\hat{C}$) y la matriz ancestral ($\hat{R}$).
• Eficiencia de Entrenamiento: Los autores implementaron varias optimizaciones arquitectónicas y de entrenamiento (por ejemplo, optimizador Muon, activación ReLU2, pre-normalización de capa) que mejoraron la velocidad de entrenamiento aproximadamente 3-4 veces en comparación con las configuraciones base.

Función de Pérdida

El modelo se optimiza utilizando una función de pérdida conjunta:
$$L = L_{pred} + \lambda \cdot L_{graph}$$
Donde $L_{pred}$ es la pérdida de entropía cruzada para la predicción de resultados, y $L_{graph}$ es la pérdida de entropía cruzada binaria para las predicciones estructurales (matrices de adyacencia, ancestral y de confusión). Pruebas teóricas en el apéndice demuestran que incluir la entrada del grafo como una condición solo puede mejorar la precisión de la posterior (o dejarla sin cambios) y que minimizar la log-probabilidad es equivalente a minimizar la divergencia KL respecto a la verdadera posterior.

Principales Contribuciones

1. Modelo Fundacional Causal Unificado: TabPFN-CFM es el primer modelo que predice conjuntamente tanto la estructura causal como los resultados a través de los tres niveles de la Jerarquía Causal de Pearl dentro de un único marco de modelo fundacional.
2. Manejo de Confusión No Observada: Al emplear ADMGs, el modelo predice explícitamente las variables de confusión no observadas, abordando una limitación en los métodos actuales de aprendizaje de estructura causal mediante aprendizaje profundo.
3. Inferencia Condicionada por Grafo: El modelo aprovecha las estructuras de grafo conocidas cuando están disponibles para mejorar la precisión de la predicción, manteniéndose robusto cuando el grafo es desconocido.
4. Mejoras de Eficiencia: El procedimiento de entrenamiento refinado y los cambios arquitectónicos resultan en una convergencia de entrenamiento casi 4 veces más rápida en comparación con los bases anteriores.

Resultados Experimentales

El modelo fue evaluado en conjuntos de datos sintéticos (tanto in-distribution como out-of-distribution) y conjuntos de datos del mundo real (Amazon Sales y Law School Admissions).

• Predicción de Resultados:
  • En SEMs sintéticos fuera de la distribución (OOD), TabPFN-CFM logró un Error Cuadrático Medio (MSE) más bajo que los meta-aprendices establecidos (S-, T-, X-, DR-learners) y Do-PFN en entornos intervencionales y contrafactuales.
  • En el conjunto de datos de Law School Admissions, el modelo mejoró significativamente la precisión en entornos contrafactuales al utilizar información contrafactual, superando a los modelos base que no pueden realizar tales predicciones.
• Predicción de Estructura:
  • TabPFN-CFM superó a los modelos base fuertes (FCI, GES, LiNGAM, PC, AVICI) en la predicción de matrices de adyacencia, ancestrales y de confusión.
  • Específicamente, en datos sintéticos OOD, el modelo alcanzó un AUROC de 0.889 para la adyacencia y de 0.828 para las matrices de confusión, significativamente superior al siguiente mejor ejecutor (AVICI con 0.623 y 0.563 respectivamente).
• Condicionamiento de Grafo: Proporcionar la estructura real del grafo como entrada mejoró el rendimiento en entornos intervencionales y contrafactuales, pero tuvo un impacto insignificante en las predicciones observacionales, lo cual es consistente con la teoría causal.

Significado

El artículo postula que TabPFN-CFM representa un paso adelante en la unificación de las tareas de inferencia causal. Al entrenar un único modelo sobre un diverso prior de sistemas causales, el enfoque demuestra que los modelos fundacionales pueden generalizarse a problemas causales del mundo real sin sobreajustarse a distribuciones prior específicas. La capacidad de manejar la confusión no observada y soportar consultas contrafactuales sugiere un camino hacia un razonamiento causal más robusto en el aprendizaje automático. Los autores señalan que, si bien el trabajo actual se centra en estructuras acíclicas, las extensiones futuras podrían abordar entornos de series temporales y causales cíclicos.