# Causal Representation Learning with Optimal Compression under Complex Treatments

arXiv: https://arxiv.org/abs/2603.11907 · HTML: https://mlsanity.com/paper/2603.11907/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 9 of its 20 references that are in the map.
- First posted on arXiv 12 Mar 2026.

## Builds on (most cited first, 9 of 9)

- [Learning Representations for Counterfactual Inference](https://mlsanity.com/paper/1605.03661/) (2016)
- [On Inductive Biases for Heterogeneous Treatment Effect Estimation](https://mlsanity.com/paper/2106.03765/) (2021)
- [Generalization Bounds and Representation Learning for Estimation of Potential Outcomes and Causal Effects](https://mlsanity.com/paper/2001.07426/) (2020)
- [VCNet and Functional Targeted Regularization For Learning Causal Effects of Continuous Treatments](https://mlsanity.com/paper/2103.07861/) (2021)
- [Perfect Match: A Simple Method for Learning Representations For Counterfactual Inference With Neural Networks](https://mlsanity.com/paper/1810.00656/) (2018)
- [Estimation of causal effects with multiple treatments: a review and new ideas](https://mlsanity.com/paper/1701.05132/) (2017)
- [Bounds on Representation-Induced Confounding Bias for Treatment Effect Estimation](https://mlsanity.com/paper/2311.11321/) (2023)
- [Geodesic Causal Inference](https://mlsanity.com/paper/2406.19604/) (2024)
- [Generalization Bounds for Causal Regression: Insights, Guarantees and Sensitivity Analysis](https://mlsanity.com/paper/2405.09516/) (2024)

## Abstract (arXiv)

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.