# Representation Learning for Semiparametric Causal Mediation Analysis under No Essential Heterogeneity

arXiv: https://arxiv.org/abs/2607.10540 · HTML: https://mlsanity.com/paper/2607.10540/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 8 of its 42 references that are in the map.
- First posted on arXiv 12 Jul 2026.

## Builds on (most cited first, 8 of 8)

- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Learning Representations for Counterfactual Inference](https://mlsanity.com/paper/1605.03661/) (2016)
- [Direct and Indirect Effects](https://mlsanity.com/paper/1301.2300/) (2013)
- [Nonparametric Estimation of Heterogeneous Treatment Effects: From Theory to Learning Algorithms](https://mlsanity.com/paper/2101.10943/) (2021)
- [Atlantic Causal Inference Conference (ACIC) Data Analysis Challenge 2017](https://mlsanity.com/paper/1905.09515/) (2019)
- [Structural Nested Models and G-estimation: The Partially Realized Promise](https://mlsanity.com/paper/1503.01589/) (2015)
- [Robust thin-plate splines for multivariate spatial smoothing](https://mlsanity.com/paper/2207.13931/) (2022)
- [RAPSEM: Identifying Latent Mediators Without Sequential Ignorability via a Rank-Preserving Structural Equation Model](https://mlsanity.com/paper/2509.23935/) (2025)

## Abstract (arXiv)

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.