Reducing distribution drift over time for longitudinal counterfactual regression

Estimating counterfactual outcomes under treatment sequences from longitudinal observational data is essential for time-series personalized decision-making. In this setting, the counterfactual regression model is trained on observations from preceding time steps and used to predict outcomes at future time steps. Therefore, it is crucial for the regression model to generalize across time steps. Existing methods typically address time-varying confounding by balancing representations of different treatment groups, while overlooking the temporal distribution drift of learned representations, which can hinder the generalization ability of the regression model over time. Motivated by this, we propose a longitudinal counterfactual regression method to reduce the temporal distribution drift of learned representations. Specifically, we characterize the temporal distribution drift using the Wasserstein distance between representations at consecutive time steps. Based on this, we establish that the future counterfactual outcome estimation error can be upper bounded by the current factual error and the distribution discrepancies measured across both treatment groups and time steps. Such a theoretical result naturally motivates us to jointly mitigate confounding bias and temporal drift. To achieve this, we propose to learn balanced representations by simultaneously reducing the distribution discrepancies across treatment groups and consecutive time steps. We conduct experiments on benchmark datasets to demonstrate the effectiveness of our method for longitudinal counterfactual regression.