Neural Score Matching for High-Dimensional Causal Inference

TL;DR

This work addresses high-dimensional causal inference by replacing scalar propensity scores with learned, multivariate balancing scores obtained via neural networks (neural score matching). It provides theoretical results that bound covariate imbalance in the original space by imbalance in a lower-dimensional balancing score space and derives bounds for both IPM-based metrics and non-balancing scores. The proposed two-stage approach learns a balancing score layer to enable low-dimensional matching, then uses the propensity-score head for potential outcome estimation, and it demonstrates competitive performance on semi-synthetic datasets (ACIC 2016, News) in terms of calibration, ATT estimation, and balance. The study highlights the practical potential of learned representations for matching in high-dimensional domains while outlining limitations and avenues for future work, including more complex architectures, miscalibration handling, and broader applicability to CATE and CEM integration.

Abstract

Traditional methods for matching in causal inference are impractical for high-dimensional datasets. They suffer from the curse of dimensionality: exact matching and coarsened exact matching find exponentially fewer matches as the input dimension grows, and propensity score matching may match highly unrelated units together. To overcome this problem, we develop theoretical results which motivate the use of neural networks to obtain non-trivial, multivariate balancing scores of a chosen level of coarseness, in contrast to the classical, scalar propensity score. We leverage these balancing scores to perform matching for high-dimensional causal inference and call this procedure neural score matching. We show that our method is competitive against other matching approaches on semi-synthetic high-dimensional datasets, both in terms of treatment effect estimation and reducing imbalance.

Figures (8)

• Figure 1: An illustration of neural score matching. In the first stage, a propensity score model is fitted to obtain low-dimensional balancing scores. In the second stage, samples are matched (to one neighbour) in the balancing score space based on a given distance metric. Matched samples are subsequently used to estimate the ATT.
• Figure S2: Calibration error boxplots on the ACIC2016 dataset: in-sample (left) and hold-out (right), without (up) and with (bottom) outliers. The data points underlying this figure refer to the average calibration error across a dataset version, corresponding to a single draw of the random seed, and a training seed.
• Figure S3: ATT error boxplots on the ACIC2016 dataset: in-sample (left) and hold-out (right), without (up) and with (bottom) outliers. The data points underlying this figure refer to the ATT computed on a dataset version, corresponding to a single draw of the random seed, and a training seed.
• Figure S4: Sample imbalance boxplots on the ACIC2016 dataset: in-sample (left) and hold-out (right), without (up) and with (bottom) outliers. The data points underlying this figure refer to sample imbalance computed on a dataset version, corresponding to a single draw of the random seed, and a training seed.
• Figure S5: ATT error boxplots on the News dataset: in-sample (left) and hold-out (right), without (up) and with (bottom) outliers. The data points underlying this figure refer to the ATT computed on a dataset version, corresponding to a single draw of the random seed, and a training seed.

Theorems & Definitions (15)

• Definition 1
• Proposition 1
• Corollary 1.1
• Proposition 2
• Proposition 3
• Proposition 4
• Corollary 4.1
• Proposition 5
• Proposition 6
• Proposition 7