Bayesian Invariance Modeling of Multi-Environment Data

Planteamiento del Problema

El artículo aborda el desafío de la selección de variables invariantes en datos de múltiples entornos. En muchos contextos estadísticos y de aprendizaje automático, los datos se recolectan a través de distintos entornos (por ejemplo, diferentes hospitales, condiciones experimentales o periodos de tiempo). Un enfoque estándar de selección de variables identifica los predictores que explican un resultado bajo la distribución específica de los datos observados. Sin embargo, esto a menudo confunde las relaciones causales estables con asociaciones específicas del entorno que no logran generalizarse a nuevos escenarios.

El objetivo es identificar un subconj conjunto de características cuya relación condicional con el resultado permanezca estable (invariante) en todos los entornos, incluso si las distribuciones marginales de las propias características varían. Aunque métodos anteriores han abordado esto mediante pruebas de hipótesis (por ejemplo, Invariant Causal Prediction, ICP) u optimización regularizada, estos enfoques suelen tener dificultades con la escalabilidad computacional en altas dimensiones o carecen de un marco probabilístico unificado para la cuantificación de la incertidumbre.

Metodología: Predicción Invariante Bayesiana (BIP)

Los autores proponen la Predicción Invariante Bayesiana (BIP), un modelo probabilístico que plantea la selección de características invariantes como un problema de inferencia bayesiana.

1. Estructura del Modelo y Supuestos
BIP se basa en dos supuestos fundamentales respecto al proceso generador de datos:

• Heterogeneidad de las Características: Las distribaciones de las características varían entre entornos.
• Invarianza del Resultado: El resultado $y$ depende de un subconjunto específico de características (el conjunto invariante) de la misma manera en todos los entornos.

El modelo trata el conjunto de características invariantes como una variable latente, codificada por un vector selector binario $z \in \{0, 1\}^p$. La distribución conjunta del selector latente y los datos observados $D$ se factoriza como:
$$p(z, D) = p(z) \prod_{e=1}^E \prod_{i=1}^{n_e} p_e(x_{z, ei}) \, g(y_{ei} | x_{z, ei}) \, p_e(x_{-z, ei} | x_{z, ei}, y_{ei})$$
Crucialmente, la distribución condicional del resultado dado las características seleccionadas, $g(y|x_z)$, se agrupa (pooled) a través de todos los entornos y es compartida, mientras que las distribaciones de las características $p_e$ son específicas de cada entorno.

2. Formulación de la Posterior
El modelo deriva una distribución posterior simplificada para el selector invariante $z$:
$$p(z | D) \propto p(z) \prod_{e=1}^E \prod_{i=1}^{n_e} \frac{g(y_{ei} | x_{z, ei})}{p_e(y_{ei} | x_{z, ei})}$$
Esta expresión revela que la probabilidad posterior de un conjunto de características es proporcional al prior multiplicado por el producto de las razones de verosimilitud entre la condicional agrupada y las condicionales locales de cada entorno. Si un conjunto $z$ es verdaderamente invariante, la condicional agrupada coincide con las condicionales locales, maximizando la razón. Si $z$ no es invariante, el desajuste reduce la probabilidad posterior.

3. Algoritmos de Inferencia

• Inferencia Exacta (BIP): Para entornos de baja dimensión, los autores enumeran todos los $2^p$ conjuntos de características candidatos para calcular la posterior exacta. Esto implica estimar las condicionales locales y agrupadas (por ejemplo, mediante máxima verosimilitud en una familia Gaussiana lineal) y calcular las razones de verosimilitud.
• Inferencia Variacional (BIP-VI): Para manejar datos de alta dimensión donde la enumeración exacta es intratable, los autores desarrollan BIP-VI. Este algoritmo aproxima la posterior utilizando una familia variacional de campo medio $q_\phi(z) = \prod \text{Bernoulli}(z_j | \text{sigmoid}(\phi_j))$. Optimiza el Límite Inferior de la Evidencia (ELBO) mediante ascenso de gradiente estocástico y estimadores de gradiente sin sesgo (U2G) para variables discretas. A diferencia de los métodos exactos, BIP-VI no requiere una búsqueda exhaustiva ni un pre-filtrado de características.

Contribuciones Teóricas

El artículo establece garantías teóricas rigurosas para el marco BIP:

• Consistencia de la Posterior: Bajo supuestos de invariancia única y estimación consistente de las condicionales, la posterior de BIP se concentra en el selector de características invariantes real $z^*$ a medida que el número de entornos ($E$) y el tamaño de la muestra por entorno ($n$) tienden al infinito.
• Tasas de Contracción: Los autores demuestran que la posterior se contrae exponencialmente rápido con respecto a $n$ y $E$. La tasa depende de dos factores:
  1. Conocimiento Previo (Prior): Los priors informativos que penalizan los conjuntos no invariantes aceleran la contracción.
  2. Heterogeneidad del Entorno: Una mayor heterogeneidad entre entornos (medida por la discrepancia entre las condicionales locales y la agrupada) conduce a una contracción de la posterior más rápida.
• Muestras Finitas y Relajaciones: La teoría se extiende a muestras finitas y supuestos relajados (por ejemplo, conjuntos invariantes no únicos o error de especificación del modelo). En estos casos, la posterior se concentra en los selectores "más aproximadamente invariantes" consistentes con el modelo y el prior.

Resultados Experimentales

Los autores evalúan BIP y BIP-VI en conjuntos de datos sintéticos y del mundo real, comparándolos con ICP, Hidden-ICP y EILLS.

1. Estudios Sintéticos

• Bajas Dimensiones ($p=10$): BIP y BIP-VI recuperan el conjunto invariante real con alta precisión, superando o igualando a los métodos existentes a medida que aumenta el número de entornos. La masa de la posterior se concentra en el conjunto real tal como predice la teoría.
• Altas Dimensiones ($p=450$): Los métodos existentes (ICP, EILLS) requieren un pre-filtrado (por ejemplo, Lasso) para reducir el espacio de características, lo que a menudo conduce a un desempeño deficiente si el filtrado omite características invariantes reales. En contraste, BIP-VI opera sobre el espacio completo de características sin pre-filtrado. Logra tasas de descubrimiento exacto y escalabilidad significativamente mayores en comparación con las versiones filtradas de otros métodos.

2. Aplicación en el Mundo Real: Perturbación de Genes
El método se aplica a un conjunto de datos de perturbación de genes de levadura ($p=6,169$ genes) para identificar reguladores de un gen objetivo.

• Desempeño: BIP-VI logra la mayor precisión entre los métodos de predicción invariante, superando a ICP-screen y EILLS-screen. Aunque su recall es moderado (debido a la naturaleza conservadora de la posterior y al número limitado de entornos), identifica un mayor número de predicciones correctas que ICP-screen (371 frente a ~67 en una comparación específica).
• Validación: Las predicciones realizadas por BIP-VI se alinean estrechamente con los principales hallazgos de la literatura previa (Meinshausen et al., 2016) y se validan contra experimentos de knockout no utilizados en el entrenamiento.

Significación y Reivindicaciones

El artículo afirma que BIP ofrece una perspectiva probabilística sobre la predicción invariante que cierra la brecha entre las pruebas de hipótesis y la optimización. Su principal significación radica en:

1. Cuantificación de la Incertidumbre: A diferencia de los métodos de estimación puntual, BIP proporciona una distribución posterior completa sobre los conjuntos invariantes, lo que permite una cuantificación de la incertidumbre fundamentada.
2. Escalabilidad: BIP-VI supera la complejidad exponencial de la predicción invariante exacta, permitiendo el análisis de datos de alta dimensión sin depender de pasos de pre-filtrado que pueden ser perjudiciales.
3. Perspectiva Teórica: El trabajo proporciona pruebas formales de consistencia y tasas de contracción, vinculando explícitamente la velocidad de aprendizaje con la heterogeneidad de los entornos.

Los autores señalan limitaciones, incluyendo la dependencia de estimaciones fijas para los modelos condicionales (lo que puede introducir sesgo en muestras pequeñas) y el uso de inferencia variacional de campo medio, que puede tener dificultades con las posteriores multimodales en entornos con muy pocos entornos. Se sugiere como trabajo futuro integrar la estimación del modelo condicional directamente en el marco probabilístico y explorar familias variacionales más expresivas.