Representation Learning Preserving Ignorability and Covariate Matching for Treatment Effects

TL;DR

This work tackles causal effect estimation from observational data plagued by hidden confounding and covariate distribution mismatch. It introduces a sequential representation-learning framework that first enforces invariance across anchor-driven domains via Inter-Domain Gradient Matching (IDGM) and then applies covariate matching using an Integral Probability Metric (IPM) before estimating effects with a CFR-based model. The approach leverages partial causal-graph knowledge (an anchor variable) to generate multiple environments, derives a theoretical bound showing that approximate invariance yields an interval around the true causal effect, and demonstrates superior or competitive performance on tabular benchmarks (IHDP, Jobs, Cattaneo) and an image-based crowd-management dataset, outperforming strong neural baselines. The combination of IDGM and covariate matching, together with the provable connection to valid adjustment sets, provides a robust and scalable method for treatment effect estimation under non-ignorable confounding with distribution shift. The work contributes a practical, theory-backed pathway for reliable causal inference in real-world, multi-domain settings with limited graph information.

Abstract

Estimating treatment effects from observational data is challenging due to two main reasons: (a) hidden confounding, and (b) covariate mismatch (control and treatment groups not having identical distributions). Long lines of works exist that address only either of these issues. To address the former, conventional techniques that require detailed knowledge in the form of causal graphs have been proposed. For the latter, covariate matching and importance weighting methods have been used. Recently, there has been progress in combining testable independencies with partial side information for tackling hidden confounding. A common framework to address both hidden confounding and selection bias is missing. We propose neural architectures that aim to learn a representation of pre-treatment covariates that is a valid adjustment and also satisfies covariate matching constraints. We combine two different neural architectures: one based on gradient matching across domains created by subsampling a suitable anchor variable that assumes causal side information, followed by the other, a covariate matching transformation. We prove that approximately invariant representations yield approximate valid adjustment sets which would enable an interval around the true causal effect. In contrast to usual sensitivity analysis, where an unknown nuisance parameter is varied, we have a testable approximation yielding a bound on the effect estimate. We also outperform various baselines with respect to ATE and PEHE errors on causal benchmarks that include IHDP, Jobs, Cattaneo, and an image-based Crowd Management dataset.

Figures (5)

• Figure 1: Our proposed architecture. Here, $\mathcal{X}^{-X_t}$ denotes covariates without the anchor variable, i.e., an instance of the features recorded in the dataset, $t$ is the treatment variable, and $Y_0$ and $Y_1$ are the factual outcomes according to the group to which the instance belongs to. $W$ are the weights of the initial Matching Representation network, which are updated by the FISH Algorithm shi2021gradient -- a second order linear approximation of IDGM. The data is divided into domains and then the gradient inner product is used to update $\tilde{W}$. Once the final weights, $\tilde{W}$ are learnt across domains and across batches, we pass on $\Phi_{W}(\mathcal{X}^{-X_t})$ as covariates in a new projection space to the CFR model, $\Gamma_{\Theta}$shalit2017estimating. The hypothesis layers $h_1$ and $h_2$ learn the treated and the control distributions separately, while the $IPM$ is applied to the treated and the control distributions in the projection space. Both objectives, invariance shi2021gradient, and covariate matching shalit2017estimating are achieved before estimating the causal effect.
• Figure 2: We showcase the ablation study performed to understand the behaviour of the IPM scaling hyperparameter $\alpha$ and the Fish update hyperparameter, $\epsilon$. As we see, even though the difference in the ATE errors is quite low, we outperform all neural baselines while having a lower PEHE error. This means that performing gradient matching is advantageous before performing covariate matching.
• Figure 3: An example of a causal graph pre-intervention and post-intervention.
• Figure 4: Example of covariates from the dataset showcased in \ref{['sec:img']} by takeuchi2021grab
• Figure 5: SC-CFR architecture, with spatial convolutional layers and pooling for image processing along with our proxy for gradient matching. We sequentially apply convolutions and pooling layers followed by gradient matching and covariate matching as shown in Figure \ref{['fig:ourarch']}.

Theorems & Definitions (6)

• Definition 1
• Definition 2
• Definition 3
• Theorem 1
• Theorem 2
• proof