Causal Representation Learning with Optimal Compression under Complex Treatments

TL;DR

This paper introduces a new theoretical framework and estimator for optimal treatment balancing in multi-treatment causal inference, improving scalability and accuracy in complex, large-scale scenarios.

Contribution

It derives a novel multi-treatment generalization bound, proposes an estimator for optimal balancing weights, and extends to a generative model preserving treatment manifold structure.

Findings

Treatment Aggregation achieves O(1) scalability with large treatment spaces.

Proposed methods outperform traditional models in accuracy and efficiency.

Multi-Treatment CausalEGM effectively models treatment manifolds in experiments.

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 $\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 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.