PaperlayerView exact revision on arXiv

Source-linked AI summary

Barycentric Fused Gromov-Wasserstein Balancing for Causal Inference under Multiple Treatments

Yuki Murakami, Takumi Hattori, Kohsuke Kubota

arXiv:2608.22024v1stat.MEcs.AIcs.LGstat.ML

TL;DR

Pairwise balancing for multiple treatments can leave residual imbalance, distort local proximity structures, and incur quadratic computational cost. CIHSI-Net addresses these challenges with BFG-WB, aligning treatment representations to a shared Wasserstein barycenter using FGW; simulations and a real-world marketing application support its effectiveness and practical utility.

  • Problem

    Pairwise balancing methods struggle with global alignment, consistent local proximity preservation, and scalability when estimating heterogeneous single and interaction treatment effects under multiple treatments.

  • Method

    CIHSI-Net uses BFG-WB to align each treatment-pattern representation to a shared Wasserstein barycenter with FGW, combining feature alignment with local geometric preservation.

  • Results

    CIHSI-Net consistently outperforms state-of-the-art baselines in simulation studies and demonstrates practical utility on real-world marketing data.

  • Takeaways & Limitations

    BFG-WB provides scalable global alignment while preserving local proximity structures for heterogeneous single and interaction effect estimation in multi-treatment settings.

  • Takeaways & Limitations

    The proposed framework remains bounded by pairwise balancing's unresolved multiple-treatment limitations, including inconsistent local geometry and quadratic cost with the number of treatment patterns.

Abstract

from arXiv · show

Estimating heterogeneous single and interaction treatment effects from observational data under multiple simultaneous treatments is crucial for decision-making. To mitigate estimation variance, previous studies balance representation distributions between every pair of treatment patterns. However, such pairwise balancing scales quadratically with the number of treatment patterns and fails to preserve consistent local proximity structures across patterns, which degrades counterfactual estimation. To address these challenges, we propose the Causal Inference for Heterogeneous Single and Interaction Treatment Effects Network (CIHSI-Net), a deep learning framework built on a novel Barycentric Fused Gromov-Wasserstein Balancing (BFG-WB) objective. BFG-WB aligns the representation distribution of each treatment pattern with a shared Wasserstein barycenter, achieving global alignment while reducing the computational complexity from quadratic to linear, and its Fused Gromov-Wasserstein discrepancy preserves the local proximity structures essential for reliable heterogeneous effect estimation. Simulation studies show that CIHSI-Net consistently outperforms state-of-the-art baselines, and an application to real-world marketing data demonstrates its practical utility in complex multi-treatment scenarios.

1 Introduction

Multi-treatment observational studies need accurate heterogeneous single and interaction effect estimates, but pairwise representation balancing leaves global, structural, and computational challenges. CIHSI-Net addresses these challenges with BFG-WB and reports strong simulation and marketing-data results.

  • Motivation: Observational estimation of heterogeneous single and interaction effects matters for applications such as healthcare and marketing.Randomized trials can be infeasible because of ethical concerns, high costs, and combinatorial treatment options.
  • Limitations of prior work: Pairwise balancing can leave residual imbalance, inconsistent neighborhood structures, and quadratic complexity across treatment patterns.Reducing discrepancy for one pair may increase discrepancies for others, while geometric distortions can degrade effect estimation.
  • Proposed framework: CIHSI-Net aligns every treatment-pattern representation with a shared Wasserstein barycenter using the BFG-WB objective.Its FGW discrepancy accounts for both feature values and geometry.
  • Proposed framework: BFG-WB targets global alignment and local proximity preservation while reducing balancing complexity from quadratic to linear.The objective is designed to reduce distributional discrepancies without sacrificing local structure.
  • Results: CIHSI-Net consistently outperforms state-of-the-art baselines for heterogeneous single and interaction effect estimation in simulations and shows practical utility on marketing data.Ablation studies support the contribution of BFG-WB.

2 Related Work

Prior work uses representation balancing and architecture-based methods for multiple treatments, but pairwise balancing does not provide a unified, scalable solution that also preserves cross-pattern geometry. CIHSI-Net is introduced to address this open challenge.

  • Representation balancing: Representation balancing reduces discrepancies between treatment-group representations to mitigate estimation variance from selection bias.Optimal Transport regularization globally aligns representation distributions across treatment groups.
  • Representation balancing: Local-structure methods preserve neighbors within distributions or guide cross-group alignment using neighborhood geometries.These extensions are motivated by counterfactual inference's reliance on local smoothness.
  • Open challenge: Pairwise balancing can let one alignment degrade another, incurs quadratic cost, and promotes geometric consistency only within individual pairs.These limitations leave residual imbalance and inconsistent structure across treatment patterns.
  • Multiple-treatment methods: Separate outcome networks may be unstable for rare treatment patterns, while latent-variable methods can lose robustness under generative assumptions.The related-work discussion contrasts these architecture-driven approaches with pseudo-sample data augmentation.
  • Open challenge: CIHSI-Net is proposed as a framework to fulfill the need for scalable global balancing with local geometry preservation.The following section defines its problem setting and causal estimands.

3 Preliminaries

The paper formulates conditional single and interaction effects under multiple binary treatments using potential outcomes and standard observational assumptions. It motivates representation balancing while highlighting the need to preserve local geometry and introduces a framework using FGW.

  • Problem setting: The dataset contains independent units with covariates, multiple simultaneous binary treatments, and factual outcomes, while potential outcomes exist for every treatment pattern.Only each unit's outcome under its observed treatment is available.
  • Identification: SUTVA, ignorability, and the stated assumptions support identification of the causal estimands from observed data.The paper states that CASE and CAIE are identifiable under Assumptions 1–3.
  • Causal estimands: CASE measures the conditional marginal effect of applying one treatment alone relative to no treatment.It extends the standard conditional average treatment effect to the multi-treatment setting.
  • Causal estimands: CAIE measures interaction among a treatment subset on the additive scale, such as the joint effect's deviation from summed single effects.For two treatments, it compares the joint treatment effect with the effects of each treatment separately.
  • Representation balancing: Wasserstein balancing reduces distributional discrepancies but does not explicitly preserve within-distribution local proximity structures.Optimal transport can map neighboring points to scattered locations, potentially causing structural collapse.
  • Representation balancing: Disrupted neighborhoods can degrade counterfactual estimation because local smoothness links similar representations with similar outcomes.Accurate estimation therefore benefits from both global alignment and neighborhood consistency.
  • Motivation for the framework: In multiple-treatment settings, pairwise structural distortions can vary across alignments and destabilize reference neighbors for heterogeneous effect inference.The paper proposes simultaneously achieving global alignment and local structure preservation with FGW.

4 CIHSI-Net: The Proposed Framework

CIHSI-Net uses BFG-WB to align every treatment-pattern representation with a shared Wasserstein barycenter while preserving local proximity structures through FGW. This star-shaped design supports error control and reduces balancing complexity from quadratic to linear in the number of treatment patterns.

  • Barycentric Alignment: BFG-WB aligns each treatment-pattern representation distribution with a shared Wasserstein barycenter to reduce global distributional imbalance.The barycenter acts as a common anchor, resolving optimization conflicts that can arise when separate treatment-pattern pairs are balanced independently.
  • Fused Discrepancy: FGW combines feature-based Wasserstein alignment with structure-based Gromov-Wasserstein preservation under one optimal transport plan.The Wasserstein component evaluates feature values, while the Gromov-Wasserstein component evaluates within-distribution geometric structure.
  • Local Structure Preservation: The FGW structural term penalizes matching nearby points to pairs with dissimilar distances, preserving neighborhood consistency across treatment patterns.This mechanism helps prevent counterfactual prediction from using neighborhoods that mix units with different heterogeneity patterns.
  • Computational Efficiency: BFG-WB requires at most (Lb + 1) L = O(L) evaluations, compared with O(L^2) for standard pairwise balancing.Here L is the number of treatment patterns and Lb is a fixed barycenter-update iteration bound independent of L.
  • Theoretical Analysis: Under stated assumptions, CASE and CAIE estimation errors are theoretically upper-bounded by factual prediction error plus distributional discrepancy terms.The bounds explicitly involve rare-treatment probabilities and the combinatorial treatment-pattern structure; FGW preserves local geometry supporting the representation smoothness assumption.

5 Simulation Experiments

The simulation experiments compare CIHSI-Net with representative multi-treatment causal inference baselines using CASE and CAIE integrated squared errors. Across settings with and without treatment interactions, CIHSI-Net achieves the lowest reported errors and retains its advantage as the number of treatments increases.

  • Experimental Design: CIHSI-Net is evaluated on simulations with complex interactions and without interactions to assess accuracy and robustness under selection bias and heterogeneous effects.The experiments compare estimation accuracy and use ablations to assess component contributions.
  • Experimental Setup: The study compares CIHSI-Net with TECE-VAE, NCoRE, and CISI-Net using 60% training, 10% validation, and 30% test data.Performance is assessed with CASE and CAIE integrated squared errors averaged over 100 independent runs.
  • Accuracy Comparison: Across all simulation settings, CIHSI-Net achieves the lowest ϵCASE and ϵCAIE estimation errors among the proposed and baseline methods.Its advantage is most pronounced for ϵCAIE, while high accuracy is maintained whether interaction effects are present.
  • Scalability: CIHSI-Net maintains its CASE and CAIE estimation advantage as K increases to eight, demonstrating scalability.The paper also reports the best overall estimation performance in the semi-synthetic setting.
  • Ablation Study: The Wasserstein barycenter combined with FGW consistently achieves the lowest average CASE and CAIE estimation errors in the ablation study.Barycentric alignment outperforms pairwise alternatives, and FGW improves performance when combined with the barycenter.

6 Application to Multiple Marketing Promotions

CIHSI-Net is applied to mobile-payment promotions to estimate heterogeneous single and interaction effects across user groups stratified by pre-promotion usage. The results show channel- and usage-dependent effects, including changing three-way interactions.

  • Real-world setting: The dataset contains three simultaneous promotions, 71 covariates, standardized payment outcomes, and 11 user groups stratified by pre-promotion payment amount.The promotions are two offline campaigns from one merchant group (CP1, CP2) and one online campaign from another (CP3).
  • Single-treatment effects: All estimated CASEs are positive, with online promotion CP3 strongest among low-usage users and offline promotions showing modest increases for higher-usage groups.These patterns indicate that strongly affected user groups differ across promotion channels and designs.
  • Interaction effects: The same-group interaction τCAIE({1, 2}) remains positive but declines with usage, whereas cross-group interactions show mixed patterns.The bottom panel is used to examine heterogeneous interaction structures across usage strata.
  • Interaction effects: The three-way interaction τCAIE({1, 2, 3}) shifts from negative for low-usage users to positive for high-usage users.The authors suggest this may reflect divided attention among low-usage users and complementary effects among regular users.
  • Robustness: Sensitivity and uncertainty analyses indicate that the qualitative CASE and CAIE patterns remain stable across BFG-WB hyperparameters α and η and broadly consistent under bootstrap analysis.The supplied passage truncates the final bootstrap statement but explicitly reports broad consistency.

7 Conclusion

The study introduces CIHSI-Net with BFG-WB to address limitations of pairwise balancing in multi-treatment causal inference. Experiments report improved estimation of heterogeneous single and interaction effects, while marketing data illustrate practical utility.

  • Contribution: BFG-WB aligns treatment representations globally through a Wasserstein barycenter while preserving local proximity structures through the FGW discrepancy.The conclusion presents this combination as the framework’s central design.
  • Empirical findings: Simulation experiments show that CIHSI-Net outperforms baselines in estimating heterogeneous single and interaction effects.The conclusion reports this result across simulation datasets without specifying a single metric value.
  • Empirical findings: A real-world marketing application illustrates the framework’s practical utility in multi-treatment settings.The application concerns heterogeneous effects from multiple promotions.

Barycentric Fused Gromov-Wasserstein Balancing

The framework connects identifiable causal estimands and counterfactual-loss bounds to barycenter-based representation balancing. BFG-WB controls discrepancies to a shared barycenter while FGW preserves local geometry, reducing evaluation complexity relative to pairwise balancing.

  • Identifiability: Under Assumptions 1–3, the conditional expected potential outcome μ(x, t) is identifiable, making CASE and CAIE identifiable linear combinations of these outcomes.CASE and CAIE are constructed from identifiable μ(x, t) terms for specific treatment vectors.
  • Computational complexity: BFG-WB requires at most a linear number of OT-based evaluations in the number of treatment patterns, whereas standard pairwise balancing requires a quadratic number.The barycenter-update and post-estimation FGW evaluations each scale linearly when the update budget is fixed.
  • Setup: The representation map, task embedding, and outcome network form the predictor input used to estimate each treatment-specific potential outcome.The predictor is defined as h([ϕ(x), t_w(t)]).
  • Theoretical motivation: Under the stated assumptions, counterfactual loss is bounded by factual loss plus Wasserstein discrepancies between treatment-specific representation distributions.Lemma 2 relates loss differences to W1, and Lemma 3 extends this relation to expected counterfactual loss.
  • Barycentric balancing: Distances to a common Wasserstein barycenter control aggregate discrepancies across multiple treatment patterns, motivating the barycenter-based regularization term Lϕ.The bound converts pairwise discrepancy control into treatment-to-barycenter discrepancy control.

A.3.3 Main Proofs: Theorems 1 and 2

The main proofs bound CASE and CAIE estimation errors under stated assumptions by combining factual prediction losses with distributional discrepancies. Barycenter-based bounds connect these theoretical guarantees to the BFG-WB training objective, subject to idealized representation assumptions.

  • Theorem 1: CASE: Theorem 1 provides an upper bound for CASE under Assumptions 1–3, auxiliary conditions, Lipschitz loss, and uniform treatment-pattern weights.The theorem applies to each individual treatment k.
  • Proof strategy: The proof derives estimation-error bounds from factual prediction-loss terms and then upper-bounds remaining pairwise Wasserstein terms using a shared barycenter.This links the theoretical bounds to the BFG-WB regularization used in training.
  • Theorem 2: CAIE: Theorem 2 provides an upper bound for CAIE under the same assumptions, using coefficients that reflect the combinatorial structure of the interaction set S.The interaction set must contain at least two treatments.
  • Assumptions and scope: The analysis treats invertibility of the representation map and Lipschitz loss in representation space as idealized conditions adopted from representation-based causal inference.The smoothness assumption is more plausible when nearby representations correspond to units with similar outcomes.

B Data Generating Process

The simulations generate two datasets with covariate-dependent treatment assignment and vary whether the outcome includes treatment interactions. Both scenarios use three treatments and 50,000 observations.

  • Two simulation datasets share functional forms for covariate generation, treatment assignment, and outcome generation, differing only in whether outcome interactions are included.Simulation 1 includes interaction effects, while the second scenario excludes them.
  • The simulations evaluate CASE and CAIE estimation under selection bias induced by covariate-dependent treatment assignment.Indicator variables H model realistic treatment assignment dependence on covariates.
  • The number of treatments is fixed at K = 3, and the sample size is N = 50,000.
  • The outcome-generation parameter l determines whether interaction-effect terms are included: l = 1 includes them and l = 0 excludes them.

C Implementation Details

The implementation specifies CIHSI-Net’s architecture, BFG-WB computation, training protocol, baseline comparisons, and scalability evaluations. Results indicate lower estimation errors and substantially shorter training times as treatment counts grow.

  • Architecture and training: CIHSI-Net uses representation learning, task embedding, and outcome prediction subnetworks built from fully connected layers with Leaky ReLU activations.The representation network outputs 64-dimensional representations and applies batch normalization because BFG-WB is sensitive to feature scaling.
  • BFG-WB computation: During training, CIHSI-Net estimates a free-support Wasserstein barycenter for each mini-batch using 16 uniformly weighted support points.Treatment-pattern-specific empirical representation distributions are used, with uniformly weighted barycenter coefficients.
  • Optimization workflow: The optimization workflow takes observed data and mapping networks as input and outputs the learning objective, with α, η, β, λ_t, and w_t controlling regularization and weighting.
  • Training protocol and baselines: CIHSI-Net and baselines are trained with comparable conditions, including Adam optimization, early stopping, and 60%/10%/30% train-validation-test splits.CIHSI-Net uses a learning rate of 10^-4, batch size 256, and L2 regularization of 10^-5 unless otherwise specified.
  • Scalability: O(L^2) optimal-transport computations for pairwise balancing are reduced to O(L) with BFG-WB.The computational reduction is accompanied by shorter training times as treatment numbers increase.
  • Scalability: At K = 8, CISI-Net requires 1529.83 [s/epoch], whereas CIHSI-Net requires 48.32 [s/epoch].At K = 5, the corresponding times are 138.54 [s/epoch] and 22.71 [s/epoch].
  • Scalability: Across K ∈{2, 4, 6, 8}, CIHSI-Net achieves the lowest average CASE and CAIE estimation errors over 50 trials.The performance gap relative to CISI-Net becomes larger as K increases.
  • Ablation and discrepancy analysis: The combination of the Wasserstein barycenter and FGW discrepancy achieves the lowest reported CASE and CAIE errors.

D.6 Evaluation on Semi-Synthetic Dataset

The semi-synthetic evaluation combines Dataset A’s real covariates and treatment assignments with known nonlinear heterogeneous outcomes. CIHSI-Net performs strongly against baselines, especially for interaction-effect estimation.

  • Data construction: The semi-synthetic dataset preserves Dataset A’s real covariate distribution and treatment imbalance while providing known ground-truth CASE and CAIE values.
  • Results: CIHSI-Net achieves the strongest overall performance under Dataset A’s empirical covariates and observed treatment assignments.CISI-Net and NCoRE are best for some individual CASE errors, but CIHSI-Net achieves the best results for all CAIE errors and the lowest averaged errors.
  • Dataset and preprocessing: Dataset A contains three simultaneous promotions, producing 2^3 = 8 treatment patterns, with payment amount standardized as the outcome.The promotions include two offline promotions from one merchant group and one online promotion from another.
  • Heterogeneity analysis: Users are stratified into 11 groups according to pre-promotion payment usage to analyze heterogeneous treatment effects.Usage intervals span 5,000 JPY, with users above 50,000 JPY aggregated into one high-usage group.

E.2 Sensitivity and Uncertainty Analysis on Dataset A

Sensitivity and uncertainty analyses assess whether Dataset A’s estimated causal-effect patterns depend on BFG-WB hyperparameters. The qualitative patterns remain stable across configurations and bootstrap intervals.

  • Sensitivity analysis: The analysis varies BFG-WB regularization strength α ∈{0.5, 1.0, 1.5} and FGW trade-off parameter η ∈{0.5, 0.6, 0.7}.Nine configurations are evaluated with other model settings fixed.
  • Uncertainty analysis: Figure A4 reports estimated CASE and CAIE across user groups, with shaded regions representing bootstrap confidence intervals.Rows encode α values and columns encode η values.
  • Findings: The overall treatment-effect patterns remain stable across hyperparameter settings, although some individual confidence intervals include zero.
  • Findings: The three-way interaction retains its sign reversal from negative for low-usage users to positive for high-usage users across configurations.
  • Dataset B comparison: For Dataset B, both estimated single-treatment effects are positive, CP5 exceeds CP4, and both effect sizes decrease with higher pre-promotion usage.
  • Dataset B comparison: The estimated CP4–CP5 interaction is consistently negative across user groups, with smaller absolute magnitude at higher usage levels.The passage interprets this pattern as consistent with spending reallocation between competing merchants.
Loading 2608.22024v1…