Estimating Long-term Heterogeneous Dose-response Curve: Generalization Bound Leveraging Optimal Transport Weights

TL;DR

This work tackles the problem of estimating the long-term heterogeneous dose-response curve (HDRC) under unobserved confounding and continuous treatment by leveraging data from short-term experiments and long-term observational data. It introduces an optimal transport (OT) based reweighting framework to align short-term outcomes across sources, enabling identifiability of the HDRC under the LU assumption, and derives a generalization bound on counterfactual prediction error using the reweighted distribution. Building on these theory results, the authors propose LEARN, a three-module estimator that combines OT weighting, balanced representation learning, and a varying-coefficient long-term predictor to handle continuous treatments. Empirical results on synthetic and semi-synthetic datasets show that LEARN outperforms baselines and demonstrate effective confounding mitigation, stability to batch size, and improved utility for personalized long-term decision-making.

Abstract

Long-term treatment effect estimation is a significant but challenging problem in many applications. Existing methods rely on ideal assumptions, such as no unobserved confounders or binary treatment, to estimate long-term average treatment effects. However, in numerous real-world applications, these assumptions could be violated, and average treatment effects are insufficient for personalized decision-making. In this paper, we address a more general problem of estimating long-term Heterogeneous Dose-Response Curve (HDRC) while accounting for unobserved confounders and continuous treatment. Specifically, to remove the unobserved confounders in the long-term observational data, we introduce an optimal transport weighting framework to align the long-term observational data to an auxiliary short-term experimental data. Furthermore, to accurately predict the heterogeneous effects of continuous treatment, we establish a generalization bound on counterfactual prediction error by leveraging the reweighted distribution induced by optimal transport. Finally, we develop a long-term HDRC estimator building upon the above theoretical foundations. Extensive experiments on synthetic and semi-synthetic datasets demonstrate the effectiveness of our approach.

Figures (3)

• Figure 1: Causal graphs of experimental and observational data.
• Figure 1: Hyperparameter sensitivity. (a) Strength of IPM. (b) Strength of short-term loss between observational and experimental datas. (c) Strength of negative entropy regularization.
• Figure 2: Model architecture of the proposed LEARN.

Theorems & Definitions (17)

• Proposition 1
• Theorem 1
• Theorem 2
• Definition 1
• Theorem 3
• Theorem 4
• Proposition 1
• proof
• Theorem 1
• proof