Causally-Aware Information Bottleneck for Domain Adaptation

2026-01-07

arXiv.org

AI Summary

This work proposes a causal-aware information bottleneck framework for zero-shot cross-domain imputation of target variables under the challenging setting where labels in the target domain are entirely absent and distributional shifts exist. By constructing mechanism-stable, compact representations that preserve causally relevant information while discarding spurious variations, the method achieves robust generalization across domains. It incorporates structural priors from directed acyclic graphs (DAGs) to unify treatment of both linear Gaussian and nonlinear/non-Gaussian scenarios: in the linear case, it yields a closed-form solution equivalent to canonical correlation analysis (CCA) projection via the Gaussian Information Bottleneck (GIB); in the nonlinear case, it employs a scalable variational information bottleneck (VIB) encoder-predictor architecture. Experiments on synthetic and real high-dimensional datasets demonstrate high imputation accuracy, confirming the approach’s effectiveness and practical utility in complex causal models.

Technology Category

• Machine Learning: Causal Learning
• Reasoning under Uncertainty: Causality
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding

Application Category

• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• domain adaptation
• causal systems
• target variable imputation
• distribution shift
• missing target

Innovation

• Causally-Aware Information Bottleneck
• Domain Adaptation
• Gaussian Information Bottleneck
• Variational Information Bottleneck
• Zero-shot Imputation