Causal Representation Learning with Optimal Compression and Complex Treatments

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 $\alpha$, 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 $\mathcal{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.

When a doctor wants to know whether a patient will benefit from a high dose versus a low dose of a medication — or which of many available treatment options is best for that specific individual — they ideally want predictions grounded in evidence. Machine learning models can help by learning from observational data such as medical records, but they face a fundamental challenge: patients who receive different treatments are often systematically different to begin with, creating misleading patterns. To fix this, models learn compressed summaries of patient data that erase treatment-related differences. But compress too aggressively and you destroy the very information needed for accurate predictions — a genuine tension with no obvious resolution. Prior work could handle simple "treated vs. untreated" comparisons, but real-world settings often involve many treatments simultaneously — dozens of drug dosages, multiple therapy options — where the math breaks down: the number of constraints explodes, tuning becomes unstable, and the right compression level is unknown. This paper shows that the "how much to compress" question has a principled statistical answer, not merely a trial-and-error one. We derive a procedure that automatically finds the optimal compression level from data — eliminating expensive manual tuning — and remains computationally efficient no matter how many treatments are involved. Experiments confirm that our approach significantly outperforms existing methods, bringing reliable individualized treatment effect estimation within reach for the complex, multi-treatment scenarios that arise throughout medicine, economics, and public policy.