[논문 리뷰] Causal Representation Learning with Optimal Compression under Complex Treatments

Wanting Liang, Haoang Chi외 1명|arXiv (Cornell University)|2026. 03. 12.

Advanced Causal Inference Techniques인용 수 0

한 줄 요약

본 논문은 bound-guided, bound-optimized 균형화 접근법을 다중 치료 인과 표현 학습에 도입하고, 세 가지 균형화 전략(Pairwise, One-vs-All, 및 Treatment Aggregation)과 최적 압축 수준을 선택하는 이층(bilevel) 방법을 제안하며, Wasserstein-geodesic counterfactuals로의 확장을 다룹니다.

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.

연구 동기 및 목표

• 다중 치료 인과 표현 학습에서 편향-정보의 trade-off를 동기부여하고 해결한다.
• 일반화 경계로 균형 가중치의 휴리스틱 튜닝을 제거하고 최적의 alpha에 대한 일관 추정치를 제공한다.
• 치료 수가 증가함에 따라 안정성을 분석하고 확장 가능한 계산 특성을 가진 세 가지 균형화 전략을 제안한다.
• Wasserstein 지오데식 구조를 보존하는 생성 아키텍처로 프레임워크를 확장하여 counterfactual 보간을 가능하게 한다.

제안 방법

• 사실적 위험과 표현 불균형을 연결하는 다중 치료 일반화 경계를 도출한다.
• 균형 가중치 alpha를 포함하는 제약 최적화(또는 패널티화) 문제로 표현 학습을 공식화한다.
• 세 가지 균형화 전략을 도입한다: Pairwise (O(K^2) 복잡도), One-vs-All (O(K) 복잡도), 그리고 HSIC를 통한 Treatment Aggregation (K에 대해 O(1))
• alpha를 선택하기 위한 경험적, bound-driven 이층 절차를 정의한다(BOAB: Bound-Optimized Adaptive Balancing).
• counterfactual 보간을 위해 Wasserstein 지오데식 구조를 보존하는 Multi-Treatment CausalEGM 아키텍처로 확장한다.
• alpha 추정기의 유한 표본 정확도와 점근적 정규성에 대한 이론적 결과 및 K에 따른 안정성 스케일링을 제공한다.

실험 결과

연구 질문

• RQ1 다중 치료 학습에서 아이덴티티 표현의 불변성과 정보 보존 사이의 균형을 휴리스틱 튜닝 없이 최적화하려면 어떻게 해야 하는가?
• RQ2 다양한 균형화 전략(Pair, OVA, Aggregation)이 치료 수 K에 따라 어떻게 스케일링되며 ITE 추정 편향과 분산에 어떤 영향을 미치는가?
• RQ3 bound 기반 alpha 추정기에 대한 유한 표본 및 점근적 보장을 제공하고 K에 따른 안정성을 정량화할 수 있는가?
• RQ4 Wasserstein-geodesic 일관 프레임워크가 치료 간 의미 있는 counterfactual 보간을 가능하게 하는가?
• RQ5 대규모 개입 설정에서 추정 정확도와 효율성의 실증적 이득은 무엇인가?

주요 결과

• 다중 치료 일반화 경계는 ITE 오차가 사실적 위험 плюс 표현 불균형 항 plus 복잡도 항에 의해 한정된다고 보여준다.
• 최적의 alpha는 bound 기반 목표를 최소화하는 이층 절차를 통해 추정 가능하며 휴리스틱 튜닝을 피한다.
• Treatment Aggregation은 K에 대해 O(1) 수준의 균형화 복잡도를 달성하여 차원의 저주를 완화한다.
• One-vs-All 균형화는 중간 규모의 K에서 최상의 성능을 보이고, Aggregation은 K가 커질수록 견고하고 확장 가능하게 남아 실험에서 최대 K=20까지 시연되었다.
• 반-합성 및 이미지 유사 데이터 세트에 대한 실증 결과는 베이스라인 대비 PEHE를 개선했고, Aggregation은 대규모 K에서도 성능을 유지하였다(예: 대규모 실험).