Bayesian Invariance Modeling of Multi-Environment Data

Probleemstelling

Het artikel behandelt de uitdaging van invariante variabele selectie in multi-omgevingsdata. In veel statistische en machine learning-omgevingen worden gegevens verzameld over verschillende omgevingen (bijv. verschillende ziekenhuizen, experimentele condities of tijdsperioden). Een standaard benadering voor variabele selectie identificeert voorspellers die een uitkomst verklaren onder de specifieke distributie van de geobserveerde data. Dit leidt echter vaak tot een verwarring tussen stabiele, causale relaties en omgevingsspecifieke associaties die niet generaliseren naar nieuwe settings.

Het doel is om een subset van kenmerken te identificeren waarvan de conditionele relatie met de uitkomst stabiel (invariant) blijft over alle omgevingen, zelfs als de marginale distributies van de kenmerken zelf variëren. Hoewel eerdere methoden dit hebben benaderd via hypothesetoetsing (bijv. Invariant Causal Prediction, ICP) of geregulariseerde optimalisatie, worstelen deze benaderingen vaak met computationele schaalbaarheid in hoge dimensies of missen ze een verenigd probabilistisch kader voor onzekerheidskwantificering.

Methodologie: Bayesian Invariant Prediction (BIP)

De auteurs stellen Bayesian Invariant Prediction (BIP) voor, een probabilistisch model dat invariante kenmerkselectie beschouwt als een Bayesiaans inferentieprobleem.

1. Modelstructuur en Aannames
BIP rust op twee kernaannames met betrekking tot het datagenererende proces:

• Kenmerkheterogeniteit: De distributies van kenmerken variëren per omgeving.
• Uitkomstinvariantie: De uitkomst $y$ hangt in alle omgevingen op dezelfde manier af van een specifieke subset van kenmerken (de invariante set).

Het model behandelt de invariante kenmerkset als een latente variabele, gecodeerd door een binaire selectorvector $z \in \{0, 1\}^p$. De gezamenlijke distributie van de latente selector en de geobserveerde data $D$ factoriseert als:
$$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})$$
Cruciaal is dat de conditionele distributie van de uitkomst gegeven de geselecteerde kenmerken, $g(y|x_z)$, over alle omgevingen wordt gepoolde en gedeeld, terwijl de kenmerkdistributies $p_e$ omgevingsspecifiek zijn.

2. Posterior Formulering
Het model leidt een vereenvoudigde posterior distributie af voor de invariante 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})}$$
Deze expressie laat zien dat de posterior waarschijnlijkheid van een kenmerkenset proportioneel is aan de prior vermenigvuldigd met het product van de likelihood-ratio's tussen de gepoolde conditionele en de lokale omgevingsspecifieke conditionalen. Als een set $z$ werkelijk invariant is, komt de gepoolde conditional overeen met de lokale conditionalen, waardoor de ratio wordt gemaximaliseerd. Als $z$ niet invariant is, vermindert de mismatch de posterior waarschijnlijkheid.

3. Inferentie-algoritmen

• Exacte Inferentie (BIP): Voor laagdimensionale settings enumereren de auteurs alle $2^p$ kandidaat-kenmerkensets om de exacte posterior te berekenen. Dit omvat het schatten van lokale en gepoolde conditionalen (bijv. via maximum likelihood in een lineair Gaussisch gezin) en het berekenen van de likelihood-ratio's.
• Variational Inference (BIP-VI): Om hoogdimensionale data te verwerken waar exacte enumeratie onuitvoerbaar is, ontwikkelen de auteurs BIP-VI. Dit algoritme benadert de posterior met behulp van een mean-field variational familie $q_\phi(z) = \prod \text{Bernoulli}(z_j | \text{sigmoid}(\phi_j))$. Het optimaliseert de Evidence Lower Bound (ELBO) met behulp van stochastic gradient ascent en unbiased gradient estimators (U2G) voor discrete variabelen. In tegenstelling tot exacte methoden vereist BIP-VI geen uitputtende zoektocht of pre-screening van kenmerken.

Theoretische Bijdragen

Het artikel stelt rigoureuze theoretische garanties vast voor het BIP-kader:

• Posterior Consistentie: Onder aannames van unieke invariantie en consistente schatting van de conditionalen, concentreert de BIP-posterior zich op de ware invariante kenmerkselector $z^*$ naarmate het aantal omgevingen ($E$) en de steekproefomvang per omgeving ($n$) naar oneindig gaan.
• Contractieraten: De auteurs bewijzen dat de posterior exponentieel snel convergeert met betrekking tot $n$ en $E$. De snelheid hangt af van twee factoren:
  1. Voorkennis: Informatieve priors die niet-invariante sets minder zwaar wegen, versnellen de contractie.
  2. Omgevingsheterogeniteit: Grotere heterogeniteit tussen omgevingen (gemeten door de discrepantie tussen lokale en gepoolde conditionalen) leidt tot een snellere posterior contractie.
• Eindige Steekproeven en Relaxaties: De theorie strekt zich uit naar eindige steekproeven en versoepelde aannames (bijv. niet-unieke invariante sets of model-misspecificatie). In deze gevallen concentreert de posterior zich op de "meest benaderend invariante" selectors die consistent zijn met het model en de prior.

Experimentele Resultaten

De auteurs evalueren BIP en BIP-VI op synthetische en real-world datasets, waarbij ze deze vergelijken met ICP, Hidden-ICP en EILLS.

1. Synthetische Studies

• Lage Dimensies ($p=10$): BIP en BIP-VI herstellen de ware invariante set met hoge nauwkeurigheid en presteren beter dan of gelijk aan bestaande methoden naarmate het aantal omgevingen toeneemt. De posterior massa concentreert zich op de ware set zoals voorspeld door de theorie.
• Hoge Dimensies ($p=450$): Bestaande methoden (ICP, EILLS) vereisen pre-screening (bijv. Lasso) om de kenmerkruimte te verkleinen, wat vaak leidt tot matige prestaties als de screening de ware invariante kenmerken mist. In tegenstelling hiertoe werkt BIP-VI op de volledige kenmerkruimte zonder pre-screening. Het bereikt aanzienlijk hogere exacte ontdekkingspercentages en schaalbaarheid vergeleken met de gescreende versies van andere methoden.

2. Real-World Applicatie: Gen-perturbatie
De methode wordt toegepast op een gist-gen-perturbatie dataset ($p=6.169$ genen) om regulatoren van een doelgen te identificeren.

• Prestaties: BIP-VI bereikt de hoogste precisie onder de invariante predictiemethoden en presteert beter dan ICP-screen en EILLS-screen. Hoewel de recall matig is (door de conservatieve aard van de posterior en het beperkte aantal omgevingen), identificeert het een groter aantal correcte voorspellingen dan ICP-screen (~371 versus ~67 in een specifieke vergelijking).
• Validatie: Voorspellingen gedaan door BIP-VI komen nauw overeen met de belangrijkste bevindingen uit de literatuur (Meinshausen et al., 2016) en worden gevalideerd tegenover uitgesloten knockout-experimenten.

Betekenis en Claims

Het artikel beweert dat BIP een probabilistisch perspectief biedt op invariante predictie dat de kloof overbrugt tussen hypothesetoetsing en optimalisatie. De primaire betekenis ligt in:

1. Onzekerheidskwantificering: In tegen tegenstelling tot point-estimate methoden, biedt BIP een volledige posterior distributie over invariante sets, wat een gefundeerde onzekerheidskwantificering mogelijk maakt.
2. Schaalbaarheid: BIP-VI overwint de exponentiële complexiteit van exacte invariante predictie, waardoor analyse van hoogdimensionale data mogelijk is zonder te vertround op potentieel informatieverliesende pre-screening stappen.
3. Theoretisch Inzicht: Het werk biedt formele bewijzen voor consistentie en contractieraten, waarbij de snelheid van leren expliciet wordt gekoppeld aan de heterogeniteit van de omgevingen.

De auteurs merken beperkingen op, waaronder de afhankelijkheid van vaste schattingen voor de conditionele modellen (wat bias kan introduceren bij kleine steekproeven) en het gebruik van mean-field variational inference, wat moeite kan hebben met multi-modale posteriors in settings met zeer weinig omgevingen. Toekomstig werk suggereert het integreren van de schatting van de conditionele modellen direct in het probabilistische kader en het verkennen van expressievere variational families.