Causally-Aware Information Bottleneck for Domain Adaptation

TL;DR

The paper addresses imputing a missing target variable $T$ in a shifted target domain by learning a mechanism-stable representation confined to the Markov blanket $M= ext{MB}(T)$, enabling zero-shot transfer under MB-invariance. It develops a DAG-aware Information Bottleneck framework with two variants: MB--GIB, a lossless, closed-form Gaussian solution equivalent to a CCA projection on $X_M$, and MB--VIB, a nonlinear, variational extension for non-Gaussian data. The authors provide identifiability and risk-preservation guarantees, finite-sample concentration results, and practical deployment guidance, plus extensive experiments on synthetic seven-node SEM, a 64-node MAGIC–IRRI gene network, and Sachs et al. single-cell data, showing robust imputations under severe distribution shifts. The proposed lightweight, structure-aware toolkit offers principled robustness for causal domain adaptation in high-dimensional settings and scalable deployment without interventional data or multi-environment supervision.

Abstract

We tackle a common domain adaptation setting in causal systems. In this setting, the target variable is observed in the source domain but is entirely missing in the target domain. We aim to impute the target variable in the target domain from the remaining observed variables under various shifts. We frame this as learning a compact, mechanism-stable representation. This representation preserves information relevant for predicting the target while discarding spurious variation. For linear Gaussian causal models, we derive a closed-form Gaussian Information Bottleneck (GIB) solution. This solution reduces to a canonical correlation analysis (CCA)-style projection and offers Directed Acyclic Graph (DAG)-aware options when desired. For nonlinear or non-Gaussian data, we introduce a Variational Information Bottleneck (VIB) encoder-predictor. This approach scales to high dimensions and can be trained on source data and deployed zero-shot to the target domain. Across synthetic and real datasets, our approach consistently attains accurate imputations, supporting practical use in high-dimensional causal models and furnishing a unified, lightweight toolkit for causal domain adaptation.

Figures (4)

• Figure 1: Causal DAG underlying the motivating example. C: stable causal drivers of $T$ (the Markov blanket); S: high-dimensional spurious proxy block whose conditional distribution shifts with D; N: nuisance block affected by D but irrelevant to $T$. The target mechanism $p(T\mid C)$ is invariant across domains.
• Figure 2: True vs. predicted (imputed) $T$ under the spurious-proxy shift. The Markov-blanket GIB remains accurate because it encodes only $C=\mathrm{MB}(T)$, for which $p_s(T\mid C)=p_t(T\mid C)$ holds exactly. In contrast, the global GIB collapses: the high-dimensional proxy block $S$ dominates the bottleneck in the source, but its relationship to $T$ reverses in the target (sign flip), yielding anti-predictive representations and large negative $R^2$.
• Figure 3: Main comparison under covariate and generalized target shift. Bars show mean $\pm$ s.e. over seeds for MB–GIB, MB–VIB, BN, IIB-style, and pure DNN. MB–GIB dominates on all metrics; the gap is largest under target shift.
• Figure 4: Heatmap summaries. Top: Ablation of VIB capacity ($z$) and compression ($\beta$) for two likelihoods. Bottom: Sensitivity to support mismatch (stretch) and missing rate for MB–VIB and MB–GIB. Lower is better (RMSE).

Theorems & Definitions (10)

• Theorem 1: Covariance concentration
• Theorem 2: MB--GIB spectral/subspace concentration
• Theorem 3: MB--GIB excess risk rate
• Corollary 1: Zero-shot target bound under MB invariance
• Corollary 2: Approximate MB invariance
• proof : Proof of Theorem \ref{['thm:cov_conc']}
• proof : Proof of Theorem \ref{['thm:gib_subspace']}
• proof : Proof of Theorem \ref{['thm:pred_rate']}
• proof : Proof of Corollary \ref{['cor:target_rate']}
• proof : Proof of Corollary \ref{['cor:approx_invar']}