Resumen Técnico: Aprendizaje de Representación Causal a partir de Datos de Redes (GraCE-VAE)

Planteamiento del Problema

Los métodos actuales de aprendizaje de representación causal a menudo asumen que las variables observadas son vectores de características no estructurados, ignorando el rico contexto relacional inherente en muchos dominios científicos (por ejemplo, interacciones proteína-proteína, membresías de vías-genes). Mientras que los métodos de inferencia basados en grafos existentes pueden utilizar estructuras de red, típicamente asumen que el grafo causal es conocido o solo consideran intervenciones observables en nodos conocidos. Por el contrario, los métodos de desentrañamiento causal que aprenden modelos estructurales causales (SCM) latentes a partir de datos observacionales e intervencionales suelen fallar al no aprovechar las redes de interacción conocidas entre las entidades medidas. Esto crea una brecha en el aprendizaje de estructuras causales latentes y efectos de intervención mientras se explota el contexto de la red observada.

El artículo aborda el desafío de descubrir factores causales latentes y sus interdependencias a partir de datos de red enriquecidos con información de intervención, específicamente en entornos donde las variables están asociadas con una red de interacción conocida (por ejemplo, vías biológicas).

Metodología: GraCE-VAE

Los autores proponen GraCE-VAE (Graph-Aware Causal Effects Variational AutoEncoder), un marco que integra Redes Neuronales de Grafos (GNN) con modelado generativo profundo para descubrir grafos causales latentes a partir de datos de redes heterogéneas.

Arquitectura

GraCE-VAE emplea una división de codificador-decodificador diseñada para separar el uso del contexto estructural del modelo causal generativo:

1. Codificador Sensible al Grafo (Graph-Aware Encoder):

   • Utiliza una GNN (específicamente instanciada como GraphSAGE) para procesar un grafo de contexto heterogéneo que contiene entidades observadas primarias ($X$, por ejemplo, genes) y entidades auxiliares ($H$, por ejemplo, vías).
   • El codificador condiciona sobre los atributos observados de $X$, la vista auxiliar $H$ y las estructuras de adyacencia ($A_{XX}, A_{XH}, A_{HH}$).
   • Produce incrustaciones (embeddings) sensibles a la estructura que parametrizan la distribución posterior sobre las variables de ruido exógeno ($Z$), infundiendo efectivamente el contexto de la red en el proceso de inferencia amortizada.
   • Crucialmente, el grafo de contexto se utiliza solo para la inferencia; no entra en el decodificador.

2. Decodificador Causal (SCM Latente):

   • El decodificador instancia un Modelo Estructural Causal (SCM) sobre variables latentes $U$.
   • Aprende un Grafo Acíclico Dirigido (DAG) que representa la estructura causal entre los factores latentes.
   • Modelado de Intervención: Las intervenciones se modelan mediante un proceso de dos etapas. Un codificador de intervención mapea los indicadores de intervención a una selección suave sobre las variables latentes (identificando el objetivo) y una propuesta de modificación del mecanismo (por ejemplo, un desplazamiento aditivo). El decodificador luego aplica esta modificación al mecanismo causal del nodo latente seleccionado mientras deja los demás sin cambios.
   • Una función de mezcla $f_\theta$ mapea las variables causales latentes $U$ de vuelta al espacio observado $X$.

Objetivo de Aprendizaje

El modelo se entrena utilizando una suma ponderada de tres componentes:

1. Término ELBO: Reconstruye las distribuciones de datos observacionales e intervencionales.
2. Término de Alineación de Distribución: Utiliza la Discrepancia de Máximo Medio (MMD) para alinear las distribuciones contrafactuales generadas por el modelo con las distribuciones intervencionales empíricas.
3. Regularización: Incluye penalizaciones de dispersión (sparsity) sobre la matriz de adyacencia del DAG aprendido y regularizadores basados en grafos.

Identificabilidad

El artículo establece que GraCE-VAE hereda las garantías de identificabilidad de los VAE de discrepancia causal previos (específicamente CMVAE [34]). Bajo los supuestos de:

• Fidelidad intervencional lineal.
• Disponibilidad de datos tanto observacionales como intervencionales.
• Las muestras son i.i.d. dentro de cada régimen de intervención.
• Cada variable latente está sujeta a al menos una intervención.
• Un supuesto de vista auxiliar donde la información de la vía $H$ ayuda a la inferencia pero no proporciona información generativa adicional para $X$ una vez que $U$ es conocido.

El DAG causal latente y los objetivos de intervención son identificables hasta la clase de equivalencia estándar (permutaciones de factores equivalentes).

Contribuciones Clave

1. Marco para Contexto Estructurado: Propone un autoencoder variacional de grafos que aprende conjuntamente un grafo causal latente y efectos de intervención, guiado por conocimiento previo de redes heterogéneas (entidades medidas y nodos de grupo). A diferencia de los métodos que asumen variables totalmente observadas o grafos causales dados, GraCE-VAE descubre la estructura latente.
2. Clarificación Teórica: Clarifica que, debido a que la clase del modelo generativo coincide con los VAE de discrepancia causal previos (el codificador sensible al grafo solo modifica la inferencia amortizada), GraCE-VAE hereda las garantías de identificabilidad existentes sin requerir nuevos teoremas para una familia de decodificadores diferente.
3. Validación Empírica: Demuestra a través de experimentos en tres conjuntos de datos de perturbación CRISPR (norman, replogle-small, replogle-large) que el aprovechamiento del contexto de red estructurado mejora la predicción de los resultados intervencionales, particularmente para combinaciones de perturbaciones no vistas. Los estudios de ablación confirman que tanto el codificador basado en grafos como el decodificador causal son esenciales para el rendimiento.

Resultados Experimentales

Los experimentos se realizaron en tres conjuntos de datos que involucran perturbaciones genéticas:

• Intervenciones Simples: GraCE-VAE logró un fuerte rendimiento de $R^2$, superando a líneas base como CMVAE (que ignora la estructura) y SENA. Fue competitivo o superior a CMVAE-multihot y VGAE (una ablación sin el decodificador causal).
• Doble Intervención No Vista: El modelo fue entrenado en intervenciones simples y evaluado en dobles intervenciones no vistas. GraCE-VAE superó significativamente a CMVAE y SENA en la predicción de estas combinaciones complejas, logrando un $R^2$ promedio de 0.7845 comparado con el 0.7382 de CMVAE. Esto resalta la sinergia entre la codificación del contexto estructurado de la GNN y la capacidad del decodificador causal para generalizar.
• Estudios de Ablación:
  • Diseño de GNN: Se probaron diferentes arquitecturas de GNN (GCN, GAT, GraphSAGE); GraphSAGE funcionó mejor. Aumentar la profundidad de las capas no produjo mejoras adicionales.
  • Niveles de Contexto: La incorporación de información de aristas más rica (gene-gene, pathway-gene, pathway-pathway) generalmente mejoró la predicción de intervención simple, aunque las ganancias se estancaron o disminuyeron ligeramente para las dobles intervenciones, sugiriendo un compromiso entre la riqueza del contexto y el sobreajuste.
  • Grafos Corruptos: Reemplazar el 50% de las aristas biológicas con aristas aleatorias degradó consistentemente el rendimiento, confirmando que la topología biológica significativa impulsa las ganancias.
  • Priors: Agregar incrustaciones de modelos fundacionales de célula única (scFM) o codificadores de intervención alternativos no superó consistentemente la configuración por defecto de GraCE-VAE, lo que sugiere que las ganancias provienen de la integración específica de la estructura del grafo y el modelado causal más que de la simple adición de más priors.

Significado y Reivindicaciones

El artículo afirma que GraCE-VAE ofrece un nuevo enfoque para descubrir factores causales latentes en sistemas de red complejos al unificar la información relacional y la inferencia causal basada en intervenciones. Su principal significancia radica en:

• Cerrar una Brecha: Aborda la falta de métodos que aprenden estructuras causales latentes mientras explotan las redes de contexto observadas entre las entidades medidas.
• Generalización: Demuestra una capacidad superior para extrapolar a combinaciones de intervenciones no vistas (por ejemplo, dobles perturbaciones) mediante el aprovechamiento del contexto biológico estructurado, una capacidad crítica para el descubrimiento científico.
• Solidez Teórica: Mantiene garantías de identificabilidad rigurosas bajo supuestos estándar, demostando que la incorporación de la estructura del grafo en el codificador no compromete la recuperabilidad teórica del modelo causal latente.

Los autores señalan limitaciones, incluyendo la dependencia de la cobertura de datos intervencionales, el uso de formas funcionales de intervención simplificadas (desplazamientos aditivos) y el supuesto específico de la vista auxiliar requerido para la identificabilidad. Concluyen que, si bien los DAGs aprendidos proporcionan hipótesis biológicamente plausibles (por ejemplo, alineándose con vías conocidas de ciclo celular o diferenciación), deben verse como espacios de hipótesis para la simulación contrafactual más que como validaciones causales definitivas sin una validación externa adicional.