Resumen Técnico: Aprendizaje de Ordenamientos Causales para Predicción Tabular en Contexto

1. Enunciado del Problema

Las arquitecturas de aprendizaje en contexto, como los Modelos Fundacionales Tabulares (TFM), han logrado un rendimiento predictivo sólido en datos tabulares aprovechando el preentrenamiento a gran escala. Sin embargo, estos modelos son principalmente asociativos; dependen de estructuras correlacionales en lugar de una direccionalidad causal explícita. En consecuencia, sus predicciones pueden ser poco fiables bajo desplazamientos de distribución o intervenciones, ya que pueden explotar asociaciones estadísticas no causales que no generalizan.

Aunque existen métodos establecidos de descubrimiento causal para recuperar Grafos Acíclicos Dirigidos (DAG), a menudo son específicos del conjunto de datos, computacionalmente intensivos y desacoplados de las arquitecturas predictivas modernas. Por el contrario, los recientes enfoques de descubrimiento causal amortizado dependen de un entrenamiento supervisado sobre grafos sintéticos, requiriendo etiquetas causales de verdad fundamental. Existe una brecha en la integración sistemática del descubrimiento de estructura causal en arquitecturas predictivas dentro de un marco de aprendizaje en contexto no supervisado que pueda manejar datos faltantes e intervenciones.

2. Metodología: TABORDER

Los autores proponen TABORDER, una arquitectura basada en transformadores diseñada para inferir simultáneamente ordenamientos causales de variables y realizar predicciones robustas. La innovación central es el uso de atención restringida por orden causal.

Diseño de la Arquitectura

TABORDER opera como un sistema de doble transformador que procesa un conjunto de datos $D$ con entradas faltantes indicadas por una máscara $M$:

1. Codificación: Cada celda $x_i^{(r)}$ y su indicador de falta se codifican en una representación vectorial.
2. Rama de Inferencia de Orden: Un transformador estándar de atención multi-cabeza auto-procesada procesa las codificaciones para inferir una puntuación escalar $s_i$ para cada variable $X_i$. Puntuaciones más bajas indican una posición anterior en el orden causal. Estas puntuaciones se promedian a través de las filas para producir un ordenamiento global $\hat{\pi}$.
3. Rama de Predicción Restringida por Orden: Un segundo transformador modela la distribución condicional de cada celda. Crucialmente, su atención por columnas está restringida por una máscara derivada de las puntuaciones aprendidas $s$. Específicamente, la variable $X_i$ solo puede atender a variables $X_j$ donde $s_j \leq s_i$. Esto impone una factorización ordenada de la distribución conjunta:
   $$p_\theta(x^{(r)} | D, \pi) = \prod_{k=1}^d p_\theta(x_{\pi(k)}^{(r)} | x_{\pi(1:k-1)}^{(r)}, D)$$
   Esto evita que el modelo utilice características "anti-causales" (causas futuras) para predecir un objetivo, reduciendo así la sensibilidad a correlaciones espurias.

Manejo de Datos Faltantes e Incertidumbre

El modelo cuenta explícitamente con la falta en los padres causales. Bajo un supuesto de Modelo de Ruido Aditivo (ANM) ($X_i = f(X_{pa(i)}) + N_i$), el modelo decodifica una media condicional $\hat{\mu}$ y una varianza $\hat{\sigma}^2$.

• Incremento de Varianza: Si un padre causal está faltante, el modelo aumenta la varianza predictiva para la variable objetivo en una cantidad aprendible $\delta_{i,j}$. Esto captura la incertidumbre aumentada cuando las entradas relevantes no se observan.
• Objetivo de Entrenamiento: El modelo se entrena de extremo a extremo maximizando la verosimilitud de las entradas enmascaradas. El objetivo acopla el aprendizaje del orden y la predicción: los ordenamientos que explican mejor los datos enmascarados (minimizando la varianza residual) reciben una mayor verosimilitud.

Justificación Teórica

Los autores justifican el enfoque basado en verosimilitud bajo el marco del Modelo de Ruido Aditivo (ANM). Demuestran que, para modelos causales aditivos, el orden óptimo de verosimilitud es consistente con el ordenamiento topológico del DAG subyacente verdadero. Además, prueban un teorema sobre la asimetría de verosimilitud consciente de la falta: contar con la falta mediante incrementos de varianza amplía la brecha de verosimilitud entre la dirección causal correcta y las direcciones incorrectas, facilitando la identificación de ordenamientos causales incluso en presencia de datos faltantes.

3. Contribuciones Clave

1. Integración del Descubrimiento Causal y la Predicción: TABORDER cierra la brecha entre el aprendizaje de estructura causal y la predicción tabular en contexto al inferir ordenamientos causales de forma no supervisada dentro de un marco de transformador predictivo.
2. Atención Restringida por Orden: La introducción de una máscara de atención diferenciable y aprendida que restringe el flujo de información a preceder a las variables en un orden causal aprendido, asegurando que las predicciones dependan solo de causas potenciales.
3. Aprendizaje Causal Consciente de la Falta: El diseño de un mecanismo donde la falta en los padres causales aumenta explícitamente la varianza predictiva, lo que teórica y empíricamente ayuda a distinguir direcciones causales.
4. Aprendizaje No Supervisado: A diferencia de los métodos de descubrimiento causal amortizado, TABORDER no requiere grafos causales de verdad fundamental durante el entrenamiento; aprende ordenamientos óptimos únicamente mediante la maximización de verosimilitud en datos enmascarados.

4. Resultados Experimentales

Los autores evalúan TABORDER en datos sintéticos con estructuras causales conocidas y en benchmarks del mundo real.

• Inferencia de Orden Causal: En datos sintéticos generados a partir de ANM de procesos gaussianos no lineales, TABORDER logra una divergencia topológica ($d_{TOP}$) comparable a algoritmos especializados de descubrimiento causal (por ejemplo, CAM, SCORE, NOTEARS), pero infiere ordenamientos en tiempo casi instantáneo mediante una sola pasada hacia adelante. En el conjunto de datos real de señalización de proteínas Sachs, TABORDER descubre ordenamientos causales con baja divergencia ($d_{TOP} \approx 0.21$) en comparación con el gráfico de consenso, superando a otras líneas base.
• Rendimiento Predictivo e Imputación: En tareas de imputación de valores faltantes (benchmarks OpenML CC18 y CTR23), TABORDER logra un rendimiento competitivo, particularmente en tasas altas de falta ($\geq 40\%$), superando a métodos estándar como KNN, MICE y GAIN.
• Robustez ante Intervenciones: En una cadena de tres variables ($X \to Y \to Z$) con una intervención en el mecanismo $Y \to Z$, TABORDER mantiene una predicción precisa del mediador $Y$ dado $X$ y $Z$. Por el contrario, los TFM estándar (como TABPFN) y XGBoost se degradan significativamente, ya que dependen de la asociación no causal rota entre $Y$ y $Z$.
• Insights Biológicos: En el conjunto de datos Sachs bajo diversas condiciones de intervención, TABORDER infiere ordenamientos que reflejan reorganizaciones biológicamente plausibles. Por ejemplo, las intervenciones aguas arriba (por ejemplo, inhibición de PKC) causan cambios globales en el orden inferido, mientras que las intervenciones aguas abajo causan cambios localizados, alineándose con la dinámica conocida de las vías de señalización.

5. Significado y Afirmaciones

El artículo afirma que TABORDER proporciona un enfoque principiado para hacer que el aprendizaje en contexto para datos tabulares sea más robusto ante desplazamientos de distribución e intervenciones al modelar explícitamente la direccionalidad causal.

• Robustez: Al imponer que las predicciones dependan solo de predecesores en un orden causal aprendido, el modelo evita explotar correlaciones espurias que fallan bajo intervención.
• Eficiencia: A diferencia del descubrimiento causal clásico que resuelve problemas de optimización por conjunto de datos, TABORDER infiere la estructura de manera eficiente dentro de una sola pasada hacia adelante, haciéndolo escalable para tareas posteriores.
• Fundamentación Teórica: El trabajo establece que la maximización de verosimilitud en presencia de datos faltantes puede servir como un proxy para el descubrimiento causal bajo supuestos estándar de ANM, proporcionando una base teórica para el aprendizaje no supervisado de ordenamientos.

Los autores reconocen limitaciones, señalando que la implementación actual se centra en ordenamientos topológicos en lugar de DAGs completos y asume características de valor real. Sugieren que, aunque el método es robusto, las estructuras inferidas son hipótesis que requieren validación de dominio, y que el trabajo futuro podría explorar priores de entrenamiento más ricos y datos de tipo mixto.