An Analysis of Causal Effect Estimation using Outcome Invariant Data Augmentation
Uzair Akbar, Niki Kilbertus, Hao Shen, Krikamol Muandet, Bo Dai
TL;DR
The work addresses causal effect estimation from observational data by reinterpreting data augmentation (DA) as a pathway to simulate interventions on the treatment mechanism when the outcome is invariant to the augmentation. It introduces IV-like (IVL) regression to relax IV assumptions while still mitigating confounding, and shows that composing parameterized DA with IVL can realize worst-case, adversarial interventions that further reduce bias. The authors provide population-level theory in linear SEMs, supported by finite-sample simulations and real-data experiments (optical-device and Colored MNIST) demonstrating improved causal estimates and robust prediction across interventions. The approach broadens the accessibility of causal inference by leveraging commonly available DA in the absence of valid IVs, while acknowledging the need for domain knowledge and careful tuning of regularization parameters. Overall, outcome-invariant DA combined with IVL regression offers a principled, practically grounded framework for reducing confounding bias and enhancing generalization across treatment interventions.
Abstract
The technique of data augmentation (DA) is often used in machine learning for regularization purposes to better generalize under i.i.d. settings. In this work, we present a unifying framework with topics in causal inference to make a case for the use of DA beyond just the i.i.d. setting, but for generalization across interventions as well. Specifically, we argue that when the outcome generating mechanism is invariant to our choice of DA, then such augmentations can effectively be thought of as interventions on the treatment generating mechanism itself. This can potentially help to reduce bias in causal effect estimation arising from hidden confounders. In the presence of such unobserved confounding we typically make use of instrumental variables (IVs) -- sources of treatment randomization that are conditionally independent of the outcome. However, IVs may not be as readily available as DA for many applications, which is the main motivation behind this work. By appropriately regularizing IV based estimators, we introduce the concept of IV-like (IVL) regression for mitigating confounding bias and improving predictive performance across interventions even when certain IV properties are relaxed. Finally, we cast parameterized DA as an IVL regression problem and show that when used in composition can simulate a worst-case application of such DA, further improving performance on causal estimation and generalization tasks beyond what simple DA may offer. This is shown both theoretically for the population case and via simulation experiments for the finite sample case using a simple linear example. We also present real data experiments to support our case.