Bayesian Invariance Modeling of Multi-Environment Data

Problema

O artigo aborda o desafio da seleção de variáveis invariantes em dados de múltiplos ambientes. Em muitos contextos estatísticos e de aprendizado de máquina, os dados são coletados através de distintos ambientes (por exemplo, diferentes hospitais, condições experimentais ou períodos de tempo). Uma abordagem padrão de seleção de variáveis identifica preditores que explicam um resultado sob a distribuição específica dos dados observados. No entanto, isso frequentemente confunde relações causais estáveis com associações específicas de um ambiente que falham em generalizar para novos cenários.

O objetivo é identificar um subconjunto de características cuja relação condicional com o resultado permaneça estável (invariante) em todos os ambientes, mesmo que as distribuições marginais das próprias características variem. Embora métodos anteriores tenham abordado isso via testes de hipóteses (ex: Invariant Causal Prediction, ICP) ou otimização regularizada, essas abordagens frequentemente enfrentam dificuldades de escalabilidade computacional em altas dimensões ou carecem de um arcabouço probabilístico unificado para quantificação de incerteza.

Metodologia: Predição Invariante Bayesiana (BIP)

Os autores propõem a Predição Invariante Bayesiana (BIP), um modelo probabilístico que trata a seleção de características invariantes como um problema de inferência bayesiana.

1. Estrutura do Modelo e Suposições
A BIP baseia-se em duas suposições fundamentais sobre o processo gerador de dados:

• Heterogeneidade de Características: As distribuições das características variam entre os ambientes.
• Invariância do Resultado: O resultado $y$ depende de um subconjunto específico de características (o conjunto invariante) da mesma forma em todos os ambientes.

O modelo trata o conjunto de características invariantes como uma variável latente, codificada por um vetor seletor binário $z \in \{0, 1\}^p$. A distribuição conjunta do seletor latente e dos dados observados $D$ fatoriza 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, a distribuição condicional do resultado dado as características selecionadas, $g(y|x_z)$, é agrupada (pooled) entre todos os ambientes e compartilhada, enquanto as distribuições das características $p_e$ são específicas de cada ambiente.

2. Formulação da Posterior
O modelo deriva uma distribuição posterior simplificada para o seletor 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 expressão revela que a probabilidade posterior de um conjunto de características é proporcional à priori vezes o produto das razões de verossimilhança entre os condicionais agrupados e os condicionais locais específicos do ambiente. Se um conjunto $z$ for verdadeiramente invariante, o condicional agrupado coincidirá com os condicionais locais, maximizando a razão. Se $z$ não for invariante, o descompasso reduz a probabilidade posterior.

3. Algoritmos de Inferência

• Inferência Exata (BIP): Para cenários de baixa dimensão, os autores enumeram todos os $2^p$ conjuntos de características candidatos para calcular a posterior exata. Isso envolve estimar os condicionais locais e agrupados (por exemplo, via máxima verossimilhança em uma família Gaussiana linear) e calcular as razões de verossimilhança.
• Inferência Variacional (BIP-VI): Para lidar com dados de alta dimensão onde a enumeração exata é intratável, os autores desenvolvem a BIP-VI. Este algoritmo aproxima a posterior usando uma família variacional de campo médio $q_\phi(z) = \prod \text{Bernoulli}(z_j | \text{sigmoid}(\phi_j))$. Ele otimiza o Limite Inferior de Evidência (ELBO) usando gradiente estocástico ascendente e estimadores de gradiente não enviesados (U2G) para variáveis discretas. Diferente dos métodos exatos, a BIP-VI não requer busca exaustiva ou pré-seleção (pre-screening) de características.

Contribuições Teóricas

O artigo estabelece garantias teóricas rigorosas para o arcabouço BIP:

• Consistência da Posterior: Sob suposições de invariância única e estimativa consistente de condicionais, a posterior da BIP concentra-se no verdadeiro seletor de características invariantes $z^*$ conforme o número de ambientes ($E$) e o tamanho da amostra por ambiente ($n$) tendem ao infinito.
• Taxas de Contração: Os autores provam que a posterior contrai exponencialmente rápido em relação a $n$ e $E$. A taxa depende de dois fatores:
  1. Conhecimento Prévio: Prioris informativas que penalizam conjuntos não invariantes aceleram a contração.
  2. Heterogeneidade do Ambiente: Uma maior heterogeneidade entre ambientes (medida pelo descompasso entre os condicionais locais e o agrupado) leva a uma contração mais rápida da posterior.
• Amostras Finitas e Relaxações: A teoria estende-se a amostras finitas e suposições relaxadas (ex: conjuntos invariantes não únicos ou erro de especificação do modelo). Nesses casos, a posterior concentra-se nos seletores "mais aproximadamente invariantes" consistentes com o modelo e a priori.

Resultados Experimentais

Os autores avaliam a BIP e a BIP-VI em conjuntos de dados sintéticos e do mundo real, comparando-as com ICP, Hidden-ICP e EILLS.

1. Estudos Sintéticos

• Baixas Dimensões ($p=10$): A BIP e a BIP-VI recuperam o conjunto invariante verdadeiro com alta precisão, superando ou igualando métodos existentes conforme o número de ambientes aumenta. A massa da posterior concentra-se no conjunto verdadeiro, conforme previsto pela teoria.
• Altas Dimensões ($p=450$): Métodos existentes (ICP, EILLS) exigem pré-seleção (ex: Lasso) para reduzir o espaço de características, o que frequentemente leva a um desempenho ruim se a pré-seleção perder características invariantes verdadeiras. Em contraste, a BIP-VI opera no espaço total de características sem necessidade de pré-seleção. Ela alcança taxas de descoberta exata e escalabilidade significativamente maiores em comparação às versões com pré-seleção de outros métodos.

2. Aplicação no Mundo Real: Perturbação Gênica
O método é aplicado a um conjunto de dados de perturbação gênica de levedura ($p=6.169$ genes) para identificar reguladores de um gene alvo.

• Desempenho: A BIP-VI alcança a maior precisão entre os métodos de predição invariante, superando ICP-screen e EILLS-screen. Embora seu recall seja moderado (devido à natureza conservadora da posterior e ao número limitado de ambientes), ela identifica um número maior de previsões corretas do que o ICP-screen (371 vs. ~67 em uma comparação específica).
• Validação: As previsões feitas pela BIP-VI alinham-se estreitamente com as principais descobertas da literatura anterior (Meinshausen et al., 2016) e são validadas contra experimentos de nocaute (knockout) não utilizados no treinamento.

Significância e Alegações

O artigo afirma que a BIP oferece uma perspectiva probabilística sobre a predição invariante que une a lacuna entre testes de hipóteses e otimização. Sua principal significância reside em:

1. Quantificação de Incerteza: Ao contrário de métodos de estimativa pontual, a BIP fornece uma distribuição posterior completa sobre os conjuntos invariantes, permitindo uma quantificação de incerteza fundamentada.
2. Escalabilidade: A BIP-VI supera a complexidade exponencial da predição invariante exata, permitindo a análise de dados de alta dimensão sem depender de etapas de pré-seleção potencialmente prejudiciais.
3. Insight Teórico: O trabalho fornece provas formais de consistência e taxas de contração, ligando explicitamente a velocidade de aprendizado à heterogeneidade dos ambientes.

Os autores observam limitações, incluindo a dependência de estimativas fixas para modelos condicionais (o que pode introduzir viés em amostras pequenas) e o uso de inferência variacional de campo médio, que pode ter dificuldade com posteriors multimodais em cenários com muito poucos ambientes. Sugere-se que trabalhos futuros integrem a estimativa de modelos condicionais diretamente no arcabouço probabilístico e explorem famílias variacionais mais expressivas.