Representation Learning for Semiparametric Causal Mediation Analysis under No Essential Heterogeneity
2026-07-12 • Machine Learning
Machine Learning
AI summaryⓘ
The authors propose a new method called UNIT to estimate how a treatment affects an outcome through a mediator variable. Their approach uses deep learning to better model how treatment effects vary with individual characteristics, then applies a statistical technique called G-estimation to get reliable estimates even when some confounding factors are unmeasured. They show that improving the first-stage modeling with a tool called TARNet leads to more precise estimates of the mediation effect without increasing errors. Simulations confirm that their method reduces uncertainty in the estimates compared to traditional approaches.
structural mediationG-estimationno essential heterogeneityTARNetconditional average treatment effectmediator-outcome confoundingdeep representation learningheterogeneous treatment effectplug-in estimatorcausal inference
Authors
Roberto Faleh, Sofia Morelli, Holger Brandt
Abstract
We propose a two-stage estimator for structural mediation parameters that combines deep representation learning with G-estimation under the "no essential heterogeneity" (NEH) assumption. We call the method UNIT. In the first stage,TARNet estimates the heterogeneous effect of a randomized treatment on a mediator by learning a shared covariate representation across treatment arms.The resulting conditional average treatment effect (CATE) estimate provides a plug-in approximation to the heterogeneity-dependent component of the weight function entering the G-estimating equation of Zheng and Zhou (2015), which identifies the structural parameters even in the presence of unmeasured mediator-outcome confounding. We show that more accurate first-stage representation learning can yield a more informative plug-in weight and thereby improve the precision of the structural parameter estimator. In simulations with non-Gaussian covariates and nonlinear mediator effects, TARNet weights reduce the Stage-2 standard error of the mediation coefficient by a factor of $1.45$ to $1.51$ (median across replications, $n \ge 2000$) relative to the classical approach, at no cost to bias or coverage.