Centroid-Referenced Mahalanobis Matching (CRM): A Scalable, Representation-Based Framework for Causal Inference in Large Observational Studies

Abstract: Matching for causal inference can be computationally expensive at scale and can silently change the target population when overlap is limited. We propose Centroid-Referenced Mahalanobis Matching (CRM), which replaces global pairwise search with stratified sampling in two reference coordinates: each unit's Mahalanobis distance from the treated centroid and its Fisher coordinate along the treated-control mean shift. All covariates enter through the treated covariance geometry; CRM is therefore not principal-component preprocessing followed by nearest-neighbor matching. For $n$ units and $p$ pretreatment covariates, its implemented cost is $O(np^{2+p^3+n\log} n)$, simplifying to $O(np^2+n\log n)$ when $n \ge p$. We derive an error decomposition separating representation, support, discretization, and stochastic components. A pre-matching shortage fraction $\hatπ$ estimates the population support restriction $π$, which enters a gap bound under bounded treatment-effect heterogeneity. Final retention is reported separately for capacity-driven exclusions. Under representation sufficiency, smoothness, and adequate cell capacity, CRM has a conservative two-dimensional histogram mean-squared-error bound $O(n_T<sup>{-1/2})$; representation sufficiency is an additional assumption, not a consequence of ignorability given the original covariates. On Criteo, CRM retains at least 99.4% of treated units, has lower MaxSMD than corrected propensity-score matching in 31 of 36 large-scale configurations, and is roughly an order of magnitude faster. Moderate-size simulations favor some pairwise and weighting baselines on balance, locating CRM's contribution in scalability and explicit support diagnostics rather than universal finite-sample dominance.

• The paper introduces CRM, which replaces costly pairwise matching with two-dimensional centroid-referenced stratification, achieving O(np² + n log n) complexity and strong scalability in large observational studies.
• The paper shows that CRM can deliver near-complete treated-unit retention and an 8.5× median speedup on the Criteo benchmark, while corrected entropy-balancing and propensity-score baselines generally achieve better moderate-sample balance.
• The paper makes overlap restrictions explicit through the pre-matching shortage fraction π̂, but cautions that representation sufficiency, nonlinear confounding, and unresolved post-matching inference limit CRM’s guarantees and practical interpretation.

Motivation and problem statement

Matching estimators for the average treatment effect on the treated (ATT) face two structural difficulties in modern observational data. First, brute-force nearest-neighbor matching requires $O(n_T n_C)$ pairwise distance evaluations, which is prohibitive when $n_T$ reaches hundreds of thousands and $n_C$ is an order of magnitude larger. Second, and more subtly, every matching procedure implicitly restricts inference to a matched subset of treated units when overlap is limited; calipers in propensity score matching (PSM), coarsening tolerances in CEM, and iteration counts in FLAME all silently determine who is excluded and therefore which conditional average treatment effect (CATT) is actually estimated. The paper introduces Centroid-Referenced Mahalanobis Matching (CRM) to address both problems simultaneously: it eliminates global pairwise search while making the support restriction explicit before any unit is discarded.

Method

CRM replaces pairwise unit search with stratified sampling in a fixed two-dimensional reference coordinate system. Each unit $x$ is summarized by $Z(x) = (d(x), \phi(x))$, where $d(x)$ is the Mahalanobis distance from the treated centroid under the treated covariance geometry $\hat\Sigma_T$ (with ridge regularization), and $\phi(x)$ is the projection onto the Fisher direction $v = w/\|w\|_2$, where $w = \hat L^{-1}(\hat\mu_C - \hat\mu_T)$ is the whitened centroid shift obtained from the Cholesky factor of $n_T$0. Controls are then sampled within cells of an equal-frequency $n_T$1 grid built on the treated marginal distributions, with bin counts set by the Freedman–Diaconis rule ($n_T$2 capped at 200, $n_T$3).

The authors are careful to distinguish this construction from "compress-then-match" pipelines such as PCA-based or sufficient-dimension-reduction matching. All covariates enter through the covariance geometry; $n_T$4 is a reference coordinate system for distributional alignment, not a low-rank replacement followed by nearest-neighbor search. Under the homoskedastic Gaussian-class model, $n_T$5 is proportional to the linear propensity-score gradient, so CRM corrects both the propensity-score dimension (via $n_T$6) and the Mahalanobis spread (via $n_T$7) without fitting a treatment model — though the paper stresses this proportionality is a geometric analogy holding only under that model.

A key design element is the pre-matching shortage fraction $n_T$8: the fraction of treated units whose grid cells contain no controls, computed and reported before any exclusion. Calibration simulations show $n_T$9 rises sigmoidally with centroid shift magnitude, crossing 5% at shift $n_C$0 and 30% at shift $n_C$1. A within-cell nearest-neighbor refinement (CRM-NN), which selects controls by whitened full-$n_C$2-dimensional distance restricted to the cell pool, adds $n_C$3 overhead and is recommended for $n_C$4.

Theoretical properties

The theory rests on three organizing principles. First, identification shifts from high-dimensional covariate adjustment to representation adequacy: Assumption 4 (representation sufficiency, ignorability given $n_C$5) is explicitly not implied by ignorability given $n_C$6. This is stated plainly and repeatedly — the Fisher direction captures the dominant linear mean-shift component but need not remove nonlinear or higher-order confounding orthogonal to it.

Second, the main error decomposition (Theorem 2) separates four components: representation bias $n_C$7 (residual within-representation confounding), support-restriction bias $n_C$8, cell-approximation error $n_C$9, and stochastic error $x$0. A corollary bounds the ATT-to-CATT gap by $x$1 under bounded treatment-effect heterogeneity. Because this decomposition applies to any matching estimator, its practical value lies in making the support weight observable via $x$2 before discarding units.

Third, under representation sufficiency, smoothness, and adequate cell capacity, CRM inherits the conservative histogram MSE bound $x$3 in fixed representation dimension $x$4, versus $x$5 for the same argument in full covariate space — for $x$6 a rate ratio of $x$7. Empirically the fitted RMSE slope is $x$8 (i.e., MSE $x$9), faster than the bound; the authors attribute this to a variance-dominated regime under the actual Freedman–Diaconis bandwidth scaling and large control pools, and explicitly decline to assert it as a sharper theorem.

The implemented complexity is $Z(x) = (d(x), \phi(x))$0, simplifying to $Z(x) = (d(x), \phi(x))$1 when $Z(x) = (d(x), \phi(x))$2, with no pairwise-search term.

Simulation results: a candid negative finding

A notable feature of the paper is its self-correction. An earlier version used a naive 1:1 PSM baseline that silently discarded treated units on caliper collisions, depressing PSM retention (to 71% in one scenario, 21% in another) and inflating its composite error metric. Replacing it with a collision-free sorted-score matcher raises PSM retention sharply (e.g., 71% → 93%) and reverses the headline conclusion: with corrected baselines, entropy balancing (EB) or collision-free PSM attains the lowest balance–retention tradeoff (CRMSE) in all four synthetic scenarios, and CRM $Z(x) = (d(x), \phi(x))$3 is no longer the leader in any of them. The paper reframes CRM's contribution accordingly, as scalability and diagnostics rather than finite-sample balance dominance.

Within-CRM ablations remain informative. Adding the Fisher coordinate to radial-only matching reduces MaxSMD dramatically — e.g., from 0.487 to 0.126 at $Z(x) = (d(x), \phi(x))$4 in the typical scenario, and from 0.539 to 0.077 under strong elliptical covariance — isolating the value of the directional coordinate since both variants share identical whitening and binning. In the structural overlap-failure scenario (centroid shift 1.0, $Z(x) = (d(x), \phi(x))$5), no method achieves acceptable balance and retention simultaneously; CRM's diagnostic flags $Z(x) = (d(x), \phi(x))$6 before matching, whereas EB retains all units yet remains badly imbalanced (MaxSMD $Z(x) = (d(x), \phi(x))$7). A nonlinear failure-mode experiment confirms the expected limitation: confounding driven by $Z(x) = (d(x), \phi(x))$8 orthogonal to the centroid shift defeats the linear Fisher coordinate, and the Rayleigh statistic correctly flags the uninformative direction.

Against PCA-95%+NN matching, CRM achieves 2–3× lower bias under linear confounding at roughly an order of magnitude lower cost, while PCA-matching wins under covariance-driven confounding — a setting where first-moment balance is uninformative for all methods (MaxSMD ≈ 0.05–0.06 yet absolute bias ≈ 2.6–3.0).

Criteo benchmark: scale, speed, and the balance–bias decoupling

The Criteo Uplift experiment (≈13.98M observations, randomized ground-truth ATE $Z(x) = (d(x), \phi(x))$9) is the central application. Two findings stand out:

• Balance does not imply accuracy. CEM achieves the best balance on every criterion (MaxSMD $d(x)$0, eight times lower than PSM's 0.020) yet records the highest estimation error ($d(x)$1Bias$d(x)$2/SE $d(x)$3 versus 50.6–54.5 for all other methods). This empirically illustrates the bias–balance decoupling formalized in the error decomposition.
• Scale and speed. Across a corrected 36-cell grid ($d(x)$4, $d(x)$5 up to 200,000, collision-free PSM), CRM attains lower MaxSMD in 31 of 36 cells, is faster in all 36, with median speedup $d(x)$6. At the largest configuration ($d(x)$7M), PSM was not attempted while CRM completed in 6.8 minutes on a single CPU core. At the headline setting, CRM retains at least 99.4–99.5% of treated units with MaxSMD $d(x)$8 versus 0.5 s runtime against 7.0 s (PSM), 21.9 s (CEM), and 94.9 s (FLAME-lite).

The authors appropriately caveat that the runtime comparison is implementation-specific, since efficient scalar-score matchers can also exploit sorting.

Real-data applications

On the LaLonde CPS benchmark (true ATT $d(x)$91{,}794$\hat\Sigma_T$0-$\hat\Sigma_T$1 to $\hat\Sigma_T$2135$\hat\Sigma_T$3-$\hat\Sigma_T$4 bias, 67% retention) on both criteria. On NHANES 2017–2018 (daily smoking vs. systolic blood pressure; $\hat\Sigma_T$5), the shortage fraction is small ($\hat\Sigma_T$6); all three estimators yield small negative, statistically insignificant associations, and survey-weighted results shift the ATT by at most 0.62 mmHg without changing conclusions. A semi-synthetic NHANES stress test with imposed overlap violation validates the conditional-estimand interpretation against known truth.

Inference limitations

The paper is explicit that post-matching inference remains unresolved. Paired-bootstrap 95% intervals achieve sub-nominal coverage of 0.745–0.870 across tested sample sizes, consistent with Abadie and Imbens' results on bootstrap failure for nonsmooth matching estimators. An inference diagnostic shows two distinct failure modes: in mean-shift settings, standard errors underestimate Monte Carlo variability by roughly 20–25% (a calibration problem); in the covariance-driven nonlinear setting, raw CRM bias is large ($\hat\Sigma_T$7), a linear correction misses it entirely (coverage 0.000), and only a quadratic outcome-model correction detects and removes it. The proposed minimum reporting standard — six quantities including pre-matching $\hat\Sigma_T$8, final retention, characteristics of excluded units, and an explicit estimand label — is a sensible operational response, but fully valid inference accounting jointly for the learned partition, nonsmoothness, and representation bias remains open.

Limitations and open questions

Three limitations are conceded directly. First, validity hinges on representation sufficiency, which is an additional assumption not implied by ignorability given $\hat\Sigma_T$9; nonlinear or multimodal treatment assignment can leave residual confounding in $\phi(x)$0, and the quadratic-correction experiment shows how easily this escapes first-moment diagnostics. Second, CRM trades finite-sample marginal balance for retention and scalability; when minimal MaxSMD in a small overlap population is the goal, PSM or weighting is preferable. Third, bootstrap intervals are approximate summaries, not guaranteed nominal intervals. Additional open questions include: whether multi-centroid or covariance-direction extensions can address multimodal treated distributions and variance-level confounding without reintroducing tuning burden; whether a design-based variance estimator replacing the bootstrap is attainable; and how CRM behaves under GPU-accelerated backends, for which no systematic benchmark is provided.

Conclusion

CRM converts matching from a pairwise search problem into stratified distributional alignment around the treated covariance geometry, achieving $\phi(x)$1 cost, near-complete treated retention under adequate overlap, and an explicit pre-matching support diagnostic grounded in a gap bound of $\phi(x)$2. Its empirical case is deliberately narrow: corrected baselines beat it on moderate-sample balance, but it matches or exceeds collision-free PSM's balance across large-scale Criteo configurations while running roughly an order of magnitude faster, and the CEM result demonstrates concretely that marginal balance is an insufficient proxy for estimation accuracy. The paper's most durable contributions may be methodological honesty — the corrected PSM baseline reversal, the sub-nominal coverage reporting, and the insistence that $\phi(x)$3 be reported before any exclusion — rather than any claim of universal dominance.