Representation Learning for Semiparametric Causal Mediation Analysis under No Essential Heterogeneity

arXiv:2607.10540 · stat.ML, cs.LG · Submitted 2026-07-12 · Read on arXiv

Roberto Faleh, Sofia Morelli, Holger Brandt

stat.ML, cs.LG

Submitted: 2026-07-12

Comments: 30 pages, 2 figures

Code: https://github.com/PsychometricsMZ/UNIT

License: http://creativecommons.org/licenses/by/4.0/

The gist: 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.

Terminology

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 2000) relative to the classical approach, at no cost to bias or coverage.

Sources

Related papers