Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

An Expectation-Maximization Algorithm for Domain Adaptation in Gaussian Causal Models

Mohammad Ali Javidian

TL;DR

The proposed DAG-aware first-order EM algorithm improves target imputation accuracy over a fit-on-source Bayesian network and a Kiiveri-style EM baseline, with the largest gains under pronounced domain shift.

Abstract

We study the problem of imputing a designated target variable that is systematically missing in a shifted deployment domain, when a Gaussian causal DAG is available from a fully observed source domain. We propose a unified EM-based framework that combines source and target data through the DAG structure to transfer information from observed variables to the missing target. On the methodological side, we formulate a population EM operator in the DAG parameter space and introduce a first-order (gradient) EM update that replaces the costly generalized least-squares M-step with a single projected gradient step. Under standard local strong-concavity and smoothness assumptions and a BWY-style \cite{Balakrishnan2017EM} gradient-stability (bounded missing-information) condition, we show that this first-order EM operator is locally contractive around the true target parameters, yielding geometric convergence and finite-sample guarantees on parameter error and the induced target-imputation error in Gaussian SEMs under covariate shift and local mechanism shifts. Algorithmically, we exploit the known causal DAG to freeze source-invariant mechanisms and re-estimate only those conditional distributions directly affected by the shift, making the procedure scalable to higher-dimensional models. In experiments on a synthetic seven-node SEM, the 64-node MAGIC-IRRI genetic network, and the Sachs protein-signaling data, the proposed DAG-aware first-order EM algorithm improves target imputation accuracy over a fit-on-source Bayesian network and a Kiiveri-style EM baseline, with the largest gains under pronounced domain shift.

An Expectation-Maximization Algorithm for Domain Adaptation in Gaussian Causal Models

TL;DR

The proposed DAG-aware first-order EM algorithm improves target imputation accuracy over a fit-on-source Bayesian network and a Kiiveri-style EM baseline, with the largest gains under pronounced domain shift.

Abstract

We study the problem of imputing a designated target variable that is systematically missing in a shifted deployment domain, when a Gaussian causal DAG is available from a fully observed source domain. We propose a unified EM-based framework that combines source and target data through the DAG structure to transfer information from observed variables to the missing target. On the methodological side, we formulate a population EM operator in the DAG parameter space and introduce a first-order (gradient) EM update that replaces the costly generalized least-squares M-step with a single projected gradient step. Under standard local strong-concavity and smoothness assumptions and a BWY-style \cite{Balakrishnan2017EM} gradient-stability (bounded missing-information) condition, we show that this first-order EM operator is locally contractive around the true target parameters, yielding geometric convergence and finite-sample guarantees on parameter error and the induced target-imputation error in Gaussian SEMs under covariate shift and local mechanism shifts. Algorithmically, we exploit the known causal DAG to freeze source-invariant mechanisms and re-estimate only those conditional distributions directly affected by the shift, making the procedure scalable to higher-dimensional models. In experiments on a synthetic seven-node SEM, the 64-node MAGIC-IRRI genetic network, and the Sachs protein-signaling data, the proposed DAG-aware first-order EM algorithm improves target imputation accuracy over a fit-on-source Bayesian network and a Kiiveri-style EM baseline, with the largest gains under pronounced domain shift.
Paper Structure (58 sections, 5 theorems, 112 equations, 3 figures, 4 tables)

This paper contains 58 sections, 5 theorems, 112 equations, 3 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. The Classical EM Algorithm and Early Variants.
  4. EM for Covariance Structure and Structural Equation Models.
  5. EM Variants, First-Order EM, and Tutorial Overviews.
  6. Modern Statistical Guarantees: Population vs. Sample-Level Analyses.
  7. Other Notable EM-Related Advances.
  8. Problem Statement
  9. Data and missingness.
  10. Gaussian DAG / SEM model.
  11. Shift classes and invariances.
  12. Methodology and Theoretical Results
  13. Population-EM Operator under a Known Causal DAG
  14. Gaussian SEM parameterization.
  15. The imputation step (E-step).
  16. ...and 43 more sections

Key Result

Lemma 1

Fix $\vartheta^{(r)}$ and consider $\widehat{Q}(\cdot\mid \vartheta^{(r)})$ as a function of $b_t$ with $\sigma_t^2$ fixed at $(\sigma_t^2)^{(r)}$. Then $\widehat{Q}(\cdot\mid \vartheta^{(r)})$ is concave and gradient-Lipschitz (i.e., $L^{(r)}$-smooth) in $b_t$, with $L^{(r)} \;=\; \lambda_{\max}\!\ and therefore constitutes a GEM step Dempster1977Wu1983.

Figures (3)

  • Figure 1: True vs. predicted (imputed) $T$ under covariate shift (top row) and a local mechanism shift at $T$ (bottom row). The DAG-aware 1st-order EM achieves near-perfect recovery of $T$ in this example, while the fit-on-source baseline degrades under mechanism shift.
  • Figure 2: Causal DAG underlying the shared data-generating process across source and target domains for the motivating example. C$_1$, C$_2$: context; Z, X: intermediates; T: target (systematically unobserved in target domain); P, Y: downstream. The causal structure is invariant across domains.
  • Figure 3: True versus predicted HT under strong interventions for three methods: (a) fit-on-source baseline, (b) Kiiveri EM, (c) our 1st-order EM.

Theorems & Definitions (11)

  • Remark 1: Local vs. basin-of-attraction guarantees (BWY-style)
  • Lemma 1: GEM ascent for the one-step update
  • Theorem 1: Population contraction for EM and block gradient-EM
  • Lemma 2: Curvature constants for the active linear-Gaussian mechanism
  • Proposition 1: Sufficient condition for gradient stability
  • Lemma 3: Lipschitz conditional-mean map for one-missing-node Gaussian SEM
  • proof : Proof of Lemma \ref{['lem:gem_ascent']}
  • proof : Proof of Theorem \ref{['thm:population_contraction_param']}
  • proof : Proof of Lemma \ref{['lem:lambda_mu']}
  • proof : Proof of Proposition \ref{['prop:gamma_bound']}
  • ...and 1 more