PaperlayerView exact revision on arXiv

Source-linked AI summary

Diverse Geometries, Frozen Weights: Robust Heterogeneous Treatment-Effect Estimation via Causal Expert Ensembles

Ali Haghpanah Jahromi, Mohammad Taheri, Zohreh Azimifar

arXiv:2609.29974v2cs.LG

TL;DR

Estimating heterogeneous treatment effects from observational data is difficult because suitable inductive biases vary across regimes and unit-level effects are unavailable for ordinary validation. GeoACE combines complementary experts with frozen, validation-learned task-level weights. The ensemble beat every constituent on six of seven PEHE benchmarks, but established estimators remained stronger on ACIC and TWINS.

  • Problem

    CATE methods must address varying overlap, treatment imbalance, prognostic structure, and sample size, while individual treatment effects are unavailable for ordinary validation.

    Source: §1.1 · p. 1, §1.2 · p. 3, §2.5 · p. 5

  • Method

    GeoACE combines experts with distinct overlap and outcome geometries and learns task-level weights from internal validation predictions before expert refitting and test evaluation.

    Source: §4.2 · p. 8, §9 · p. 21

  • Results

    The five-expert ensemble beat every constituent on six of seven PEHE benchmarks, while established estimators remained stronger on ACIC and TWINS.

    Source: §8 · p. 20

  • Takeaways & Limitations

    The results support controlled expert diversity and leakage-safe aggregation as a benchmark-dependent robustness strategy, not universal dominance.

    Source: §8 · p. 20, §9 · p. 21

    The validation criteria are observable surrogates rather than direct CATE risk, and task-level weights cannot adapt to subregions within a dataset.

    Source: §8 · pp. 20–21

Abstract

from arXiv · show

Estimating heterogeneous treatment effects from observational data is difficult because the most appropriate inductive bias varies with overlap, treatment imbalance, prognostic structure, and sample size. We introduce the Geometry-Diverse Anchor-Correction Expert Ensemble (GeoACE), a five-expert framework that combines a common anchor-correction estimator with complementary overlap-aware and outcome-guided geometries. Its task-level ensemble weights are learned only from internal validation predictions, frozen before test evaluation, and then applied to experts refitted on the complete development sample. The fifth expert, O-Phi-ACE, constructs an outcome-free, overlap-aware statistical projection from covariates and treatment assignment and replaces the anchor input with this lower-dimensional geometry. We evaluate GeoACE against 11 comparators on eight benchmark protocols. Adding O-Phi-ACE reduced mean sqrt(PEHE) relative to the four-expert ensemble on all seven benchmarks with individual-effect truth, winning 998 of 1,225 paired tasks; the change on JOBS policy risk was negligible. The five-expert ensemble ranked first on IHDP100, IHDPA, and IHDPB and second on NEWS, differing from the NEWS leader by 0.13%. Across the seven sqrt(PEHE) benchmarks it obtained the lowest observed average rank (3.714), although the omnibus Friedman and Iman-Davenport tests were not significant (p=0.328 and p=0.330). Using the same five frozen experts, inverse-DR weighting was consistently better than winner-take-all selection, convex DR fitting, R-stacking, and causal Q-aggregation in benchmark-balanced analyses, but was statistically indistinguishable from equal weighting and DR ridge shrinkage. The evidence therefore supports geometry-diverse expert libraries and leakage-free aggregation as a robustness strategy, not universal superiority of either GeoACE or one weighting rule.

1 Introduction

Observational studies must estimate treatment-effect variation without observing both potential outcomes for any unit, while treatment assignment and support vary across covariate profiles. GeoACE addresses this challenge with five complementary experts and validation-learned, frozen task-level weights, while retaining standard identification assumptions.

1.1 Context and Motivation

  • No single estimator is uniformly well matched to variation in overlap, prognostic structure, treatment imbalance, and sample size across datasets.The individual unit-level treatment-effect contrast is not jointly observed, complicating direct validation and model selection.Source: §1.1 · p. 2, §1.2 · pp. 2–3

1.2 Problem Definition and Technical Challenges

  • Treatment–control imbalance can make factual prediction unreliable for counterfactual prediction, while limited overlap weakens support for one potential-outcome surface.The introduction also identifies unmeasured confounding, rare-subgroup sparsity, and spillovers as challenges or scope boundaries.Source: §1.2 · p. 2

1.3 Contributions

  • GeoACE combines five experts on a shared anchor–correction backbone, including OΦ-ACE’s outcome-free, treatment-aware projection, and freezes task-level validation weights before refitting and test evaluation.The paper compares factual, doubly robust, ridge-regularized, equal-weight, selection, stacking, and Q-aggregation controls using the same frozen expert predictions.Source: §1.3 · p. 3
  • Across eight benchmark protocols, the results support controlled expert diversity and leakage-safe aggregation, but remain benchmark-dependent rather than establishing universal dominance.The ensemble’s test-isolated workflow separates weight selection, expert refitting, and final evaluation.Source: §1.1 · p. 2, §1.3 · p. 3

2 Related Work

CATE estimators encode different ways to share information, balance treatment arms, model heterogeneity, and handle overlap, with no single geometry suited to every regime. GeoACE studies whether a controlled expert set and validation-based aggregation can address this model-selection problem without test leakage.

2.1 Methodological Families and Their Trade-offs

  • Outcome regressions, propensity-based balancing, meta-learners, and orthogonal learners make different trade-offs in interpretability, information sharing, arm imbalance, and nuisance sensitivity.Orthogonal and doubly robust corrections can be noisy under weak overlap and remain dependent on identification assumptions.Source: §2.1 · pp. 3–4

2.2 Local, Bayesian, and Representation-Based Estimators

  • Tree, Bayesian, and neural approaches capture local, nonlinear, or shared response structure, but their benefits depend on support, model assumptions, and how treatment-related information is retained or balanced.Specialized subgroup, dynamic-sharing, and pairwise methods may add complexity and become unstable with small samples.Source: §2.2 · p. 4

2.3 Latent and Prior-Fitted Models

  • Latent generative and prior-fitted estimators infer missing responses or reuse simulated-task priors, making accuracy dependent on structural assumptions, proxy quality, or prior–target correspondence.These approaches do not relax ignorability or overlap requirements.Source: §2.3 · pp. 4–5

2.4 CATE Model Selection and Aggregation

  • Because individual effects are unavailable for ordinary validation, observable pseudo-outcome losses can guide selection or aggregation but may be noisy when overlap or nuisance estimation is weak.GeoACE evaluates its inverse-DR rule against equal weighting, selection, convex DR fitting, R-stacking, causal Q-aggregation, and DR ridge using the same frozen predictions.Source: §2.4 · p. 5

2.5 Comparative Synthesis

  • Estimator families target different combinations of sharing, balance, overlap, nuisance robustness, and effect-aligned training, but these targets do not guarantee protection against unmeasured confounding or structural positivity violations.GeoACE uses task-level simplex weights from observable validation predictions rather than winner-take-all selection or an individual-level gate; its validation criteria remain noisy causal surrogates.Source: §2.5 · p. 5

3 Problem Setup

The paper targets the conditional average treatment effect, identified from treatment-specific conditional outcome functions under stated causal assumptions. It learns task-level expert weights within development data, then uses convex aggregation to combine expert predictions.

3.1 Observed Data and Target Estimand

  • The target is the CATE, τ(x)=µ1(x)−µ0(x), with ATE, ATT, potential-outcome error, and policy quantities treated as complementary evaluation criteria.Identification relies on consistency, no interference, conditional exchangeability, and positivity over the region where effects are estimated.Source: §3.1 · p. 6

3.2 Data Partitioning and Information Boundaries

  • The held-out test set is excluded from nuisance estimation, anchor construction, checkpoint selection, and expert weighting; counterfactual outcomes, when available, are accessed only after predictions are finalized.Because individual effects are unavailable in standard validation data, the paper uses factual prediction error and a doubly robust pseudo-outcome as complementary proxy criteria.Source: §3.2 · pp. 6–7

3.3 Expert Predictions and Convex Aggregation

  • Each of five experts predicts both treatment-specific outcomes and their contrast, and a task-level simplex weight vector combines those predictions across individuals.Aggregating the two response surfaces and then taking their difference is equivalent to aggregating expert-specific CATE estimates; task-level weights avoid a covariate-dependent gate.Source: §3.3 · p. 7

4 Proposed Method

GeoACE combines five anchor–correction experts with different inductive biases using task-level weights learned from internal validation predictions, frozen before test evaluation, and applied after refitting on all development data. Its shared backbone and varied geometries provide a common starting point while allowing specialized corrections.

4.1 Overview

  • Validation predictions determine one task-level weight vector, which is frozen before the five experts are refitted on all development data and evaluated on the test set.The workflow separates fitting, validation-based selection, refitting, and final test evaluation.Source: §4.1 · p. 8, §4.4.1 · p. 14

4.2 Common Anchor–Correction Backbone

  • Each expert starts from structured, low-dimensional arm-specific outcome anchors and learns treatment-specific corrections, with an anchor-deviation penalty limiting unsupported departures.The shared backbone controls architectural variation, so expert differences primarily reflect their causal inductive biases.Source: §4.2 · pp. 8–9

4.3 Complementary Causal Experts

  • GΦ-ACE uses pooled outcome-guided geometry, AΦ-ACE uses arm-specific outcome-guided geometry, and OΦ-ACE uses outcome-free overlap geometry.The experts share the causal target but impose caution through different inductive biases.Source: §4.3 · p. 10
  • ACE is the unmodified reference, while OW-ACE changes the factual-loss weights to emphasize covariate regions supported by both treatment arms.OW-ACE does not create overlap or change the CATE estimand.Source: §4.3.1 · p. 10, §4.3.2 · p. 10

4.3.3 Outcome-Guided and Outcome-Free Overlap Geometry

  • The geometry ranks directions by overlap-supported variation and, when included, outcome signal relative to measured treatment imbalance; it does not create overlap or address unmeasured confounding.GΦ-ACE favors shared prognostic structure, whereas AΦ-ACE can retain directions relevant to only one treatment arm.Source: §4.3.3 · p. 11
  • OΦ-ACE removes outcome utility and supplies a lower-dimensional overlap-aware representation to the anchor, while the neural correction continues to receive the original standardized covariates.This makes its geometry distinct from loss reweighting and the outcome-guided constructions.Source: §4.3.3 · pp. 11–12

4.4.1 Inverse Doubly Robust Risk Weighting

  • The primary selector converts experts’ validation doubly robust CATE errors into simplex weights using an inverse-risk rule.The selector uses ACE nuisance predictions and clips propensity scores for stability; inverse factual RMSE and DR ridge are prespecified alternatives.Source: §4.4.1 · pp. 13–14

4.5 Selection, Refit, and Evaluation

  • Refitting uses the complete development sample with training durations and ensemble weights held fixed, but the convex ensemble need not outperform its best constituent.Finite-sample clipping, nuisance estimation, and reuse of validation observations limit the population-level interpretation of the DR criterion.Source: §4.5 · pp. 14–15

5 Experimental Setup

The evaluation spans eight fixed benchmark protocols varying in sample size, dimensionality, outcome type, treatment prevalence, overlap, and response complexity. Methods use common task partitions and held-out predictions, with benchmark-appropriate metrics and choices fixed before final evaluation.

5.1 Datasets

  • Eight fixed protocols span semi-synthetic settings with individual-effect truth and JOBS, an experimental benchmark without individual counterfactual labels.The protocols include three distinct IHDP constructions, NEWS, TWINS, JOBS, and ACIC2016 and ACIC2017 tasks.Source: §5.1 · p. 15

5.4 Implementation Details

  • All methods use the same task partitions, and preprocessing is fit on the fitting subset; ensemble weights and selected training durations are frozen before refitting and one-time test evaluation.Predictions are persisted before test truth is accessed, and the same evaluator is applied to comparators and proposed estimators.Source: §5.1 · p. 15, §5.4 · p. 17

5.2 Baselines

  • The comparison includes eleven reported comparator variants spanning meta-learning, response-surface, orthogonal, neural, subgroup, and pairwise approaches.CFRNet-MMD and CFRNet-WASS count separately because their discrepancy penalties produce different fitted estimators.Source: §5.2 · pp. 15–16

5.3 Evaluation Metrics

  • For benchmarks with individual-effect truth, sqrt(PEHE) is primary; JOBS instead uses absolute ATT error and policy risk because unit-level counterfactual truth is unavailable.ATE error measures average-effect calibration and is interpreted jointly with PEHE, not as a measure of recovered heterogeneity.Source: §5.3 · pp. 16–17

6 Results

GeoACE combines experts whose strengths vary across benchmarks, achieving its best cross-benchmark average rank without statistically significant omnibus superiority. Its performance is strong on IHDP and NEWS but weaker on ACIC and TWINS, while JOBS policy performance is comparable to the four-expert ensemble.

6.1 Performance of the Individual Experts and Their Ensemble

  • The ensemble outperformed every individual expert on six of seven PEHE benchmarks; on JOBS, its policy risk was identical to OΦ-ACE’s.The strongest individual expert varied by benchmark, and the ensemble matched the four-expert policy risk nearly exactly on JOBS.Source: §6.1 · pp. 17–18, Table 3 · p. 18

6.2 Comparison with Established Estimators

  • GeoACE ranked first on all three IHDP protocols and second on NEWS, 0.00220 behind PairNet (0.13% relative), but was less competitive on ACIC and TWINS.Table 4 compares GeoACE with the strongest comparator on each PEHE benchmark.Source: Table 4 · p. 18, §6.2 · p. 18
  • 3.714 was GeoACE’s lowest observed average rank across seven PEHE benchmarks, but omnibus tests were nonsignificant (p=0.328; p=0.330).The average-rank result is descriptive and does not establish overall superiority.Source: §6.2 · p. 18

6.3 Policy Learning on JOBS

  • On JOBS, GeoACE’s mean policy risk was 0.23758 (policy value 0.76242), versus 0.23755 for the four-expert ensemble; adding OΦ-ACE was practically neutral.JOBS is evaluated on policy value rather than unit-level PEHE.Source: §6.3 · p. 19

7 Additional Analysis

The analyses find that OΦ-ACE adds useful prediction geometry, although expert contributions and aggregation performance vary by benchmark. The conclusions remain descriptive: neither every expert nor every secondary metric improves uniformly.

7.1 Contribution of OΦ-ACE

  • Adding OΦ-ACE as a fifth expert to the otherwise identical four-expert ensemble reduced mean PEHE on all seven compatible benchmarks and won 998 of 1,225 paired tasks.On TWINS and ACIC2017, the fifth expert improves PEHE but can worsen mean-effect calibration; the ablation does not establish gains on every secondary metric.Source: §7.1 · p. 19

7.2 Leave-One-Expert-Out Sensitivity

  • Removing OΦ-ACE raised mean PEHE on all seven truth-available benchmarks by 0.57%–3.83%, while its JOBS policy-risk effect was essentially zero.OW-ACE helped on five benchmarks, but removing it improved NEWS; the remaining experts had smaller, sign-changing effects after reweighting.Source: §7.2 · pp. 19–20

7.3 Aggregation-Rule Controls

  • Against the same frozen experts, inverse-DR had lower benchmark-balanced excess loss than Best-DR (6.97%), Convex-DR (3.69%), R-stacking (3.30%), and causal Q-aggregation (4.90%).Equal-5 and Ridge-DR differed from GeoACE by −0.12% and −0.09%, respectively, with intervals spanning zero; the four alternatives’ intervals excluded zero.Source: §7.3 · p. 20

7.4 Sensitivity of the Cross-Benchmark Ranking

  • Excluding TWINS, GeoACE again had the lowest observed average rank (3.167), but Friedman and Iman–Davenport tests remained nonsignificant (p=0.255; p=0.250).Removing TWINS improved the numerical rank without changing the inferential conclusion.Source: §7.4 · p. 20

8 Discussion

The discussion attributes GeoACE’s empirical coverage to experts with benchmark-dependent strengths, while emphasizing that neither the ensemble nor its weighting rule is universally superior. Its use also entails identification, validation, and computational constraints.

Discussion

  • Expert strengths exchange rank across datasets, and the five-expert ensemble beats every constituent on six of seven PEHE benchmarks, yet established estimators remain stronger on ACIC and TWINS.GeoACE leads the IHDP protocols and nearly matches the NEWS leader.Source: §8 · p. 20
  • Inverse-DR weighting outperforms hard selection and less-regularized validation objectives, but is effectively tied with equal weighting and DR ridge shrinkage.The favorable average rank is descriptive, not evidence of statistical dominance over every alternative.Source: §8 · p. 20
  • Observational conclusions depend on consistency, conditional exchangeability, and positivity; expert diversity cannot recover information lost to unmeasured confounding or structural non-overlap.The task-level weights also cannot adapt to subregions within a dataset, and improved PEHE can coexist with worse ATE or ATT calibration.Source: §8 · pp. 20–21
  • Fitting five experts and refitting them after freezing weights and training durations costs more than using a single estimator; the selector itself estimates only four free simplex parameters.Expert fits can run in parallel, but runtime, memory, and prediction storage remain higher.Source: §8 · p. 21
  • Future work could test regional gates, weight uncertainty, multi-valued or continuous treatments, and external validation against randomized reference evidence.Such validation could help determine when overlap geometry is preferable to outcome-guided geometry.Source: §8 · p. 21

9 Conclusion

GeoACE combines complementary experts with a task-level validation weight vector that is frozen before refitting and test evaluation. The conclusion supports controlled diversity and leakage-safe aggregation, but not universal superiority or an advantage of inverse-DR over uniform or ridge-shrunk averaging.

9 Conclusion

  • GeoACE preserves several causal inductive biases without an individual-level gate; results across eight protocols support the expert library, especially on IHDP and NEWS, but not universal superiority.The evidence also does not distinguish inverse-DR weighting from uniform or ridge-shrunk averaging.Source: §9 · p. 21

A Validation Selectors Beyond the Primary Rule

  • The alternative validation selectors use the primary rule’s validation predictions and exclude test outcomes and simulated treatment-effect truth; the final ensemble retains inverse-DR weights unless stated otherwise..Source: A Validation Selectors Beyond the… · p. 25

A.1 Inverse Factual-RMSE Weighting

  • Inverse-RMSE weighting uses a=1 and ϵ=10^-8 for smooth normalization, avoiding winner-take-all selection; factual-response accuracy need not imply accurate contrast estimation.The passage describes this as useful when the validation sample is modest.Source: §A.1 · p. 25

A.2 Simplex-Constrained Nonnegative DR Ridge

  • The nonnegative DR ridge selector constrains weights to the simplex and uses λ2=0.01 and λu=0.01, with the latter shrinking toward equal weighting.Ridge regularization stabilizes weights when expert predictions are strongly correlated; no test criterion tunes the solution.Source: §A.2 · p. 25

B Mathematical Properties and Limits

  • Anchor-correction can reduce finite-sample variance when the anchor captures useful outcome structure, but correlated anchor error remains shared; Φ geometries do not create common support.Validation distributes weight among controlled biases rather than assuming one is uniformly preferable.Source: B Mathematical Properties and Limits · pp. 25–26

B.1 Convex Aggregation

Convex aggregation bounds ensemble CATE error by the weighted expert errors, but does not guarantee improvement over the best expert. The section explains validation surrogates and refitting, scopes the guarantees, and constructs geometry matrices from overlap, outcome association, and treatment imbalance.

B.1 Convex Aggregation

  • Convex aggregation bounds pointwise ensemble squared error by the weighted sum of expert squared errors.The bound relies on nonnegative simplex weights; it does not guarantee improvement over the best expert, which requires useful experts with imperfectly correlated errors.Source: §B.1 · p. 26

B.2 Validation Surrogates and Fixed Refit

  • Factual RMSE is a stable but indirect selector, while population DR risk targets CATE risk up to an additive noise term under stated conditions.Estimated nuisances, propensity clipping, weak overlap, and limited validation size can make the DR pseudo-outcome noisy; cross-fitting would impose stronger independence.Source: §B.2 · p. 26
  • GeoACE freezes one task-level simplex weight vector before refitting, but refit drift can make those pre-refit weights suboptimal.The joint DR solution is regularized toward equal weighting, and individual experts and equal weighting remain explicit comparators.Source: §B.2 · p. 27

B.3 Scope

  • GeoACE cannot identify effects under unmeasured confounding, recover unsupported counterfactuals, guarantee improvement over the best expert, or vary weights across covariate regions.Its role is to address task-level uncertainty over controlled inductive biases, with performance evaluated across repeated benchmark tasks.Source: §B.3 · p. 27

C Computational Construction of the Geometry Matrices

  • The geometry construction standardizes fitting covariates, estimates and clips propensity scores, and uses overlap-weighted covariate moments to build its matrices.Outcome-guided variants incorporate overall or arm-specific outcome associations and measured imbalance; the outcome-free O-Phi-ACE instead uses only fitting-subset covariates and treatment, replacing the anchor input with a retained half-dimensional projection.Source: C Computational Construction of the… · pp. 27–28
  • The geometry experts solve a symmetric generalized eigenproblem, ranking directions by overlap-supported variation and any included outcome association relative to measured treatment-group separation.The leading ceil(p/2) eigenvectors form the projection; G-Phi-ACE whitens projected scores, while A-Phi-ACE does not.Source: C Computational Construction of the… · p. 28

D.1 Leave-One-Expert-Out Results

The leave-one-expert-out analysis compares the full and reduced systems using validation-only refitting and predictions locked before test truth is opened. O-Phi-ACE shows a positive contribution on six of seven PEHE benchmarks, while JOBS has no significant expert-removal contrast after adjustment.

D.1 Leave-One-Expert-Out Results

  • The full and reduced systems differ only in the available expert column and resulting validation-frozen simplex weights, with test truth locked until predictions are fixed.Figure 2 reports paired mean increase in primary loss after removing an expert, divided by full-ensemble mean loss; positive values favor retaining that expert.Source: §D.1 · pp. 28–29, Figure 2 · p. 29
  • O-Phi-ACE has a positive paired contribution on six of seven PEHE benchmarks after within-benchmark Holm adjustment; no expert-removal contrast is significant on JOBS.IHDP100 has a positive paired-location test, but its bootstrap interval for mean absolute change crosses zero, making that case distribution-sensitive.Source: §D.1 · pp. 28–29, Figure 2 · p. 29

D.2 Controls for the Aggregation Rule

Aggregation controls use the same five frozen expert predictions and internal validation data, isolating combination rules from expert training. GeoACE outperforms several selection and flexible-fit controls across benchmarks, while GeoACE, equal weighting, and Ridge-DR have smaller differences that change sign across protocols.

D.2 Controls for the Aggregation Rule

  • The controls hold five expert prediction columns fixed and freeze each rule's validation-based weights or selected expert before test truth is opened.Equal-5, Best-DR, Convex-DR, R-stacking, Causal-Q, and Ridge-DR use respectively uniform weights, validation DR selection, simplex DR regression, residualized R-loss, Q-aggregation, and DR regression shrunk toward uniform weights.Source: §D.2 · p. 29
  • GeoACE significantly improves on Best-DR in seven benchmarks, Convex-DR in six, R-stacking in five, and causal Q-aggregation in six.NEWS favors Convex-DR and R-stacking, while differences among GeoACE, Equal-5, and Ridge-DR change sign across benchmarks.Source: §D.2 · pp. 29–30, Figure 3 · p. 30
  • Figure 3 compares benchmark-specific primary losses relative to Equal-5; GeoACE, Equal-5, and Ridge-DR cluster more tightly than hard selection and several flexible validation fits.Negative relative values favor the indicated rule, and the less stable controls vary across protocols.Source: Figure 3 · p. 30

Related papers

Loading 2609.29974v2…