[Paper Review] Neural Score Matching for High-Dimensional Causal Inference

Oscar Clivio, Fabian Falck+3 more|arXiv (Cornell University)|Mar 1, 2022

Bayesian Modeling and Causal Inference4 citations

TL;DR

This paper proposes neural score matching, a novel method for high-dimensional causal inference that uses neural networks to learn low-dimensional balancing scores, enabling effective matching in high-dimensional covariate spaces. The approach reduces imbalance and improves treatment effect estimation, outperforming traditional methods on semi-synthetic benchmarks including ACIC2016, News, and IHDP datasets.

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.

Motivation & Objective

• To address the curse of dimensionality in traditional matching methods used for causal inference.
• To develop a theoretically grounded method for learning multivariate balancing scores using neural networks.
• To enable effective matching in high-dimensional settings without relying on outcome regression.
• To improve treatment effect estimation and covariate balance in observational studies with complex, high-dimensional covariates.
• To provide a scalable, non-parametric alternative to propensity score matching and exact matching in high-dimensional data.

Proposed method

• The method uses a neural network to learn a low-dimensional balancing score from high-dimensional covariates, derived from the intermediate layers of a propensity score model.
• Matching is performed in the learned balancing score space using a distance metric, rather than in the original covariate space.
• Theoretical bounds are derived to relate imbalance in the original covariate space to imbalance in the lower-dimensional score space.
• The approach extends to non-balancing scores by analyzing their deviation from the balancing condition.
• The method is trained end-to-end to minimize imbalance while preserving treatment effect estimation accuracy.
• Neural score matching avoids outcome regression by focusing solely on covariate balancing through learned representations.

Experimental results

Research questions

• RQ1 Can neural networks learn low-dimensional balancing scores that preserve covariate balance in high-dimensional settings?
• RQ2 How does imbalance in the original covariate space relate to imbalance in the learned score space?
• RQ3 Can neural score matching outperform traditional matching methods in terms of treatment effect estimation and balance metrics?
• RQ4 Does the method remain robust when the number of covariates exceeds the sample size or when covariates are high-dimensional?
• RQ5 How does the choice of network layer for the balancing score affect performance compared to full model outputs?

Key findings

• Neural score matching achieved the lowest calibration error on the ACIC2016 dataset, with in-sample error of 0.156 ± 0.005 and hold-out error of 0.311 ± 0.011.
• On the IHDP dataset, the method achieved an ATT error of 0.156 ± 0.005 (in-sample) and 0.311 ± 0.011 (hold-out), outperforming LogReg PS and PCA-based methods.
• The method significantly reduced sample imbalance compared to random matching and baseline methods, with hold-out imbalance of 0.442 ± 0.009 on ACIC2016.
• Neural score matching showed competitive performance on the News dataset, with lower ATT error and imbalance than PCA and random matching.
• Theoretical results establish that imbalance in the original space is bounded by imbalance in the score space, justifying the use of neural networks for dimensionality reduction in matching.
• The method demonstrated robustness across multiple benchmarks, with consistent improvements in balance and estimation accuracy despite high-dimensional covariates.