Propensity Weighted federated learning for treatment effect estimation in distributed imbalanced environments

Estimating treatment effects from observational data in medicine using causal inference is a very relevant task due to the abundance of observational data and the ethical and cost implications of conducting randomized experiments or experimental interventions. However, how could we estimate the effect of a treatment in a hospital that has very restricted access to treatment? In this paper, we want to address the problem of distributed causal inference, where hospitals not only have different distributions of patients, but also different treatment assignment criteria. Furthermore, it is necessary to take into account that due to privacy restrictions, personal patient data cannot be shared between hospitals. To address this problem, we propose an adaptation of the federated learning algorithm FederatedAveraging to one of the most advanced models for the prediction of treatment effects based on neural networks, TEDVAE. Our algorithm adaptation takes into account the shift in the treatment distribution between hospitals and is therefore called Propensity Weighted FederatedAveraging (PW FedAvg). As the distributions of the assignment of treatments become more unbalanced between the nodes, the estimation of causal effects becomes more challenging. The experiments show that PW FedAvg manages to reduce errors in the estimation of individual causal effects when imbalances are large, compared to Vanilla FedAvg and other federated learning-based causal inference algorithms based on the application of federated learning to linear parametric models, Gaussian Processes and Random Fourier Features. © 2024 The Author(s)