Causal Representation Learning with Optimal Compression under Complex Treatments

Wanting Liang; Haoang Chi; Zhiheng Zhang

cs.LGMar 12, 2026

Causal Representation Learning with Optimal Compression under Complex Treatments

Authors: Wanting Liang, Haoang Chi, Zhiheng Zhang

Organizations: Shanghai University of Finance and Economics, Shanghai, China.

Abstract

Estimating Individual Treatment Effects (ITE) in multi-treatment scenarios faces two critical challenges: the Hyperparameter Selection Dilemma for balancing weights and the Curse of Dimensionality in computational scalability. This paper derives a novel multi-treatment generalization bound and proposes a theoretical estimator for the optimal balancing weight $α$ , eliminating expensive heuristic tuning. We investigate three balancing strategies: Pairwise, One-vs-All (OVA), and Treatment Aggregation. While OVA achieves superior precision in low-dimensional settings, our proposed Treatment Aggregation ensures both accuracy and O(1) scalability as the treatment space expands. Furthermore, we extend our framework to a generative architecture, Multi-Treatment CausalEGM, which preserves the Wasserstein geodesic structure of the treatment manifold. Experiments on semi-synthetic and image datasets demonstrate that our approach significantly outperforms traditional models in estimation accuracy and efficiency, particularly in large-scale intervention scenarios.

Figures & tables

Panel of Figure 1 — Figure 1: Performance Comparison (PEHE). (a) All methods outperform baseline at $K=4$ . (b) At $K=20$ , Pairwise degrades while Aggregation remains robust. (See Appendix D.3 , Figure 5 for training efficiency)

Panel of Figure 2 — Figure 2: Performance and Efficiency Analysis on Digits Dataset. (a) The estimated ADRF (blue) closely tracks the ground truth (red). (b) Dual-axis comparison showing PEHE error (bars) and Training Time (line).

Panel of Figure 3 — Figure 3: Geometric Validation on Hierarchical Treatments. (a) The learned embeddings spontaneously recover the underlying tree structure, placing the Root centrally and separating the L/R branches. (b) Counterfactual interpolation from Leaf LL( $Y=-3$ ) to Leaf RR ( $Y=+3$ ) respects the causal topology by passing through the Root’s effect region ( $Y\approx 0$ ), whereas the linear baseline (grey) ignores the structure.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Detailed Architecture of Multi-Treatment CausalEGM. The left panel shows the original binary design, while the right panel highlights our proposed extensions (Embedding Layer and Softmax Activation) for complex treatment regimes.

Figure 5: Training Time Efficiency at $K=20$ . The Pairwise strategy incurs distinct computational costs due to $\binom{K}{2}$ constraints.

Figure 6: Latent Space Topology (Rotated MNIST). The model learns an isometric embedding of the treatments, recovering the cyclic order $0^{\circ}\to 45^{\circ}\to\dots\to 315^{\circ}\to 0^{\circ}$ .

Figure 7: Global Interpolation ( $0^{\circ}\to 180^{\circ}$ ). The geodesic path correctly models the monotonic decrease in outcome $Y$ as the digit rotates from upright to upside-down, tracking the cosine function.

Figure 8: Local Interpolation ( $0^{\circ}\to 315^{\circ}$ ). Despite the large gap in discrete indices ( $0$ vs $7$ ), the model recognizes the topological proximity, yielding a smooth transition consistent with the periodic boundary condition.

Causal Effect Estimation

Causal effect estimation measures how outcomes would change under interventions, using experiments or observational data and explicit assumptions.

See all 264 papers

Barycentric Fused Gromov-Wasserstein Balancing for Causal Inference under Multiple Treatments
Aug 22, 2026Yuki Murakami, Takumi Hattori, Kohsuke KubotaCausal Effect Estimation Causal Inference
NTT DOCOMO, INC., Tokyo, Japan · Yokohama City University, Kanagawa, Japan
Supervised Learning Objectives for High-Dimensional Treatment Effects
May 26, 2026Nikita Dhawan, Arnav Paruthi, Andrew Kim +3Causal Effect Estimation Causal Inference
University of Toronto, Vector Institute · University of Toronto · Vanguard
TLRNet: Estimating Individual Treatment Effect based on Local Information and Single Learner Structure
Jul 23, 2026Ali Haghpanah Jahromi, Mohammad Taheri, Zohreh AzimifarCausal Effect Estimation Conditional Average Treatment Effect Estimation
Department of Electrical and Computer Engineering, University of Shiraz, Shiraz, Iran

Aug 22, 2026J/K move · Enter open · S save

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

Yuki Murakami, Takumi Hattori, Kohsuke Kubota

NTT DOCOMO, INC., Tokyo, Japan · Yokohama City University, Kanagawa, Japan

Causal Effect Estimation Causal Inference Conditional Average Treatment Effect Estimation Optimal Transport

Open paper AskPDF arXiv

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.

Abstract

Causal Effect Estimation

Explore similar work

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

Supervised Learning Objectives for High-Dimensional Treatment Effects

TLRNet: Estimating Individual Treatment Effect based on Local Information and Single Learner Structure