Bayesian Invariance Modeling of Multi-Environment Data

Problem Statement

The paper addresses the challenge of invariant variable selection in multi-environment data. In many statistical and machine learning settings, data is collected across distinct environments (e.g., different hospitals, experimental conditions, or time periods). A standard variable selection approach identifies predictors that explain an outcome under the specific distribution of the observed data. However, this often conflates stable, causal relationships with environment-specific associations that fail to generalize to new settings.

The goal is to identify a subset of features whose conditional relationship with the outcome remains stable (invariant) across all environments, even if the marginal distributions of the features themselves vary. While previous methods have approached this via hypothesis testing (e.g., Invariant Causal Prediction, ICP) or regularized optimization, these approaches often struggle with computational scalability in high dimensions or lack a unified probabilistic framework for uncertainty quantification.

Methodology: Bayesian Invariant Prediction (BIP)

The authors propose Bayesian Invariant Prediction (BIP), a probabilistic model that casts invariant feature selection as a Bayesian inference problem.

1. Model Structure and Assumptions
BIP relies on two core assumptions regarding the data-generating process:

• Feature Heterogeneity: Feature distributions vary across environments.
• Outcome Invariance: The outcome $y$ depends on a specific subset of features (the invariant set) in the same way across all environments.

The model treats the invariant feature set as a latent variable, encoded by a binary selector vector $z \in \{0, 1\}^p$. The joint distribution of the latent selector and the observed data $D$ factorizes as:
$$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})$$
Crucially, the conditional distribution of the outcome given the selected features, $g(y|x_z)$, is pooled across all environments and shared, while the feature distributions $p_e$ are environment-specific.

2. Posterior Formulation
The model derives a simplified posterior distribution for the invariant selector $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})}$$
This expression reveals that the posterior probability of a feature set is proportional to the prior times the product of likelihood ratios between the pooled conditional and the local environment-specific conditionals. If a set $z$ is truly invariant, the pooled conditional matches the local conditionals, maximizing the ratio. If $z$ is not invariant, the mismatch reduces the posterior probability.

3. Inference Algorithms

• Exact Inference (BIP): For low-dimensional settings, the authors enumerate all $2^p$ candidate feature sets to compute the exact posterior. This involves estimating local and pooled conditionals (e.g., via maximum likelihood in a linear Gaussian family) and calculating the likelihood ratios.
• Variational Inference (BIP-VI): To handle high-dimensional data where exact enumeration is intractable, the authors develop BIP-VI. This algorithm approximates the posterior using a mean-field variational family $q_\phi(z) = \prod \text{Bernoulli}(z_j | \text{sigmoid}(\phi_j))$. It optimizes the Evidence Lower Bound (ELBO) using stochastic gradient ascent and unbiased gradient estimators (U2G) for discrete variables. Unlike exact methods, BIP-VI does not require exhaustive search or pre-screening of features.

Theoretical Contributions

The paper establishes rigorous theoretical guarantees for the BIP framework:

• Posterior Consistency: Under assumptions of unique invariance and consistent estimation of conditionals, the BIP posterior concentrates on the true invariant feature selector $z^*$ as the number of environments ($E$) and sample size per environment ($n$) approach infinity.
• Contraction Rates: The authors prove that the posterior contracts exponentially fast with respect to $n$ and $E$. The rate depends on two factors:
  1. Prior Knowledge: Informative priors that down-weight non-invariant sets accelerate contraction.
  2. Environment Heterogeneity: Greater heterogeneity between environments (measured by the discrepancy between local and pooled conditionals) leads to faster posterior contraction.
• Finite Sample and Relaxations: The theory extends to finite samples and relaxed assumptions (e.g., non-unique invariant sets or model misspecification). In these cases, the posterior concentrates on the "most approximately invariant" selectors consistent with the model and prior.

Experimental Results

The authors evaluate BIP and BIP-VI on synthetic and real-world datasets, comparing them against ICP, Hidden-ICP, and EILLS.

1. Synthetic Studies

• Low Dimensions ($p=10$): BIP and BIP-VI recover the true invariant set with high accuracy, outperforming or matching existing methods as the number of environments increases. The posterior mass concentrates on the true set as predicted by theory.
• High Dimensions ($p=450$): Existing methods (ICP, EILLS) require pre-screening (e.g., Lasso) to reduce the feature space, which often leads to poor performance if the screening misses true invariant features. In contrast, BIP-VI operates on the full feature space without pre-screening. It achieves significantly higher exact discovery rates and scalability compared to screened versions of other methods.

2. Real-World Application: Gene Perturbation
The method is applied to a yeast gene-perturbation dataset ($p=6,169$ genes) to identify regulators of a target gene.

• Performance: BIP-VI achieves the highest precision among invariant prediction methods, outperforming ICP-screen and EILLS-screen. While its recall is moderate (due to the conservative nature of the posterior and limited environments), it identifies a larger number of correct predictions than ICP-screen (371 vs. ~67 in a specific comparison).
• Validation: Predictions made by BIP-VI align closely with top findings from previous literature (Meinshausen et al., 2016) and are validated against held-out knockout experiments.

Significance and Claims

The paper claims that BIP offers a probabilistic perspective on invariant prediction that bridges the gap between hypothesis testing and optimization. Its primary significance lies in:

1. Uncertainty Quantification: Unlike point-estimate methods, BIP provides a full posterior distribution over invariant sets, allowing for principled uncertainty quantification.
2. Scalability: BIP-VI overcomes the exponential complexity of exact invariant prediction, enabling the analysis of high-dimensional data without relying on potentially lossy pre-screening steps.
3. Theoretical Insight: The work provides formal proofs of consistency and contraction rates, explicitly linking the speed of learning to the heterogeneity of the environments.

The authors note limitations, including the reliance on fixed estimates for conditional models (which may introduce bias in small samples) and the use of mean-field variational inference, which may struggle with multi-modal posteriors in settings with very few environments. Future work is suggested to integrate conditional model estimation directly into the probabilistic framework and explore more expressive variational families.