Beyond Flatland: A Geometric Take on Matching Methods for Treatment Effect Estimation

TL;DR

GeoMatching addresses the core challenge of covariate confounding in observational causal inference by incorporating the geometry of the data manifold. It learns a latent representation and a Riemannian metric to compute geodesic distances, guiding nearest-neighbor matching along the manifold rather than in raw Euclidean space. The approach yields more accurate TE estimates, demonstrating robustness to increasing input dimensionality, presence of outliers, and benefits in semi-supervised settings across synthetic Swissroll, semi-synthetic Mocap, IHDP, and Lalonde datasets. This geometry-aware matching framework offers a principled way to reduce extrapolation bias and opens avenues for integrating differential geometry with causal discovery and broader causal inference methods. The work highlights practical improvements in TE estimation and provides a foundation for future optimization of geodesic computations and latent representations in causal tasks.

Abstract

Matching is a popular approach in causal inference to estimate treatment effects by pairing treated and control units that are most similar in terms of their covariate information. However, classic matching methods completely ignore the geometry of the data manifold, which is crucial to define a meaningful distance for matching, and struggle when covariates are noisy and high-dimensional. In this work, we propose GeoMatching, a matching method to estimate treatment effects that takes into account the intrinsic data geometry induced by existing causal mechanisms among the confounding variables. First, we learn a low-dimensional, latent Riemannian manifold that accounts for uncertainty and geometry of the original input data. Second, we estimate treatment effects via matching in the latent space based on the learned latent Riemannian metric. We provide theoretical insights and empirical results in synthetic and real-world scenarios, demonstrating that GeoMatching yields more effective treatment effect estimators, even as we increase input dimensionality, in the presence of outliers, or in semi-supervised scenarios.

Figures (22)

• Figure 1: Key idea of the paper. We propose GeoMatching, a geometry-aware matching framework that pairs cases and controls along the manifold of confounders $X$. A) Assumed causal graph in this work; B) Different sub-causal graphs for covariates $X$ often induce different geometric data structures; latent variable $R$ represents the manifold structure; C) GeoMatching learns a latent Riemannian metric $\mathbf{G}$ (represented as colorbar) which captures the geometric structure and uncertainty of the data, enabling matched pairs to be faithful to the manifold structure. D) Geometry-aware matched samples translate into more accurate estimates of individual (ITE) and average treatment effects (ATE).
• Figure 4: Additional Results for (Semi)-Synthetic Datasets. (1st row) Swissroll data; (2nd row) Motion Capture data. 1st column is the same as in the main text to facilitate comparison; 2nd column corresponds to Precision in Estimation of Heterogeneous Treatment Effects (PEHE) averaged over all seeds, and 3rd column shows Individual Treatment Effect (ITE) Absolute Error for a single seed.
• Figure 6: Impact of dimensionality $D$ on TE estimation. As $D$ increases, matching all covariates perfectly becomes impossible.
• Figure 7: Sensitivity Analysis of GeoMatching w.r.t different Hyperparameters. We consider variability for a single random seed in ATE Absolute Error performance w.r.t different latent representations (PCA, Isomap, and NSM latent space), latent dimensionality, and $\sigma$, the parameter that controls the curvature of the Riemannian manifold.
• Figure : (a) Euclidean in $\mathcal{X}$

Theorems & Definitions (6)

• Definition 1
• Definition 2
• Definition 3
• Definition 4
• Definition 5
• Definition 6