Dynamic Inter-treatment Information Sharing for Individualized Treatment Effects Estimation

TL;DR

Estimating individualized treatment effects $\tau(x)$ from observational data is challenged by data scarcity and a lack of end-to-end inter-treatment information sharing. The authors introduce HyperITE, a framework that employs a hypernet to generate weights for ITE learners, enabling dynamic information sharing across treatment groups during training and yielding variants like HyperTLearner and HyperTARNet. Empirical results across IHDP, ACIC-2016, and Twins show that HyperITE typically improves PEHE over strong baselines, with larger gains on smaller datasets, demonstrating the value of end-to-end shared learning in causal inference. This work offers a general, practical approach to enhance ITE estimation in limited-data settings and broadens the applicability of neural approaches for personalized decision-making in real-world domains.

Abstract

Estimation of individualized treatment effects (ITE) from observational studies is a fundamental problem in causal inference and holds significant importance across domains, including healthcare. However, limited observational datasets pose challenges in reliable ITE estimation as data have to be split among treatment groups to train an ITE learner. While information sharing among treatment groups can partially alleviate the problem, there is currently no general framework for end-to-end information sharing in ITE estimation. To tackle this problem, we propose a deep learning framework based on `\textit{soft weight sharing}' to train ITE learners, enabling \textit{dynamic end-to-end} information sharing among treatment groups. The proposed framework complements existing ITE learners, and introduces a new class of ITE learners, referred to as \textit{HyperITE}. We extend state-of-the-art ITE learners with \textit{HyperITE} versions and evaluate them on IHDP, ACIC-2016, and Twins benchmarks. Our experimental results show that the proposed framework improves ITE estimation error, with increasing effectiveness for smaller datasets.

Figures (8)

• Figure 1: An overview of the architectures and gradient flows for T-Learner and HyperTLearner, where $\hat{\mu}(x;w_0) \text{ and } \hat{\mu}(x;w_1)$ have exactly the same architecture in both but different training process.
• Figure 2: Architectures and gradient flows for TARNet and HyperTARNet, where $\hat{\mu}(x;w_0) \text{ and } \hat{\mu}(x;w_1)$ have exactly the same architecture in both but different training process. HyperTARNet is employed to train and share information between PO heads, while $\phi(x)$ is learnt similarly to TARNet.
• Figure 3: Effect of dataset size on the performance of selected learners using ACIC-2016 and Twins datasets (shaded region shows one standard error).
• Figure 4: An overview of the architectures and gradient flows for SNet+ and HyperSNet+, where $\hat{\mu}(x;w_0), \hat{\mu}(x;w_1) \text{ and } \hat{\pi}(x;w_p)$ have exactly the same architecture in both but different training process.
• Figure 5: Effect of hypernet's type on performance of HyperTARNet and HyperTLearner using PEHE-in (left) and PEHE-out (right) performance metrics on IHDP, ACIC-2016 and Twins datasets.