# Federated Learning for Estimating Heterogeneous Treatment Effects

arXiv: https://arxiv.org/abs/2402.17705 · HTML: https://mlsanity.com/paper/2402.17705/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 70 references that are in the map.
- First posted on arXiv 27 Feb 2024.

## Builds on (most cited first, 30 of 35)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [The Future of Digital Health with Federated Learning](https://mlsanity.com/paper/2003.08119/) (2020)
- [Estimation and Inference of Heterogeneous Treatment Effects using Random Forests](https://mlsanity.com/paper/1510.04342/) (2015)
- [Federated Learning Based on Dynamic Regularization](https://mlsanity.com/paper/2111.04263/) (2021)
- [Exploiting Shared Representations for Personalized Federated Learning](https://mlsanity.com/paper/2102.07078/) (2021)
- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Cross-stitch Networks for Multi-task Learning](https://mlsanity.com/paper/1604.03539/) (2016)
- [Quasi-Oracle Estimation of Heterogeneous Treatment Effects](https://mlsanity.com/paper/1712.04912/) (2017)
- [Learning Representations for Counterfactual Inference](https://mlsanity.com/paper/1605.03661/) (2016)
- [Recursive Partitioning for Heterogeneous Causal Effects](https://mlsanity.com/paper/1504.01132/) (2015)
- [Personalized Federated Learning using Hypernetworks](https://mlsanity.com/paper/2103.04628/) (2021)
- [FedBE: Making Bayesian Model Ensemble Applicable to Federated Learning](https://mlsanity.com/paper/2009.01974/) (2020)
- [Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes](https://mlsanity.com/paper/1704.02801/) (2017)
- [FedPD: A Federated Learning Framework with Optimal Rates and Adaptivity to Non-IID Data](https://mlsanity.com/paper/2005.11418/) (2020)
- [Nonparametric Estimation of Heterogeneous Treatment Effects: From Theory to Learning Algorithms](https://mlsanity.com/paper/2101.10943/) (2021)
- [TransTab: Learning Transferable Tabular Transformers Across Tables](https://mlsanity.com/paper/2205.09328/) (2022)
- [Automated versus do-it-yourself methods for causal inference: Lessons learned from a data analysis competition](https://mlsanity.com/paper/1707.02641/) (2017)
- [XTab: Cross-table Pretraining for Tabular Transformers](https://mlsanity.com/paper/2305.06090/) (2023)
- [On Inductive Biases for Heterogeneous Treatment Effect Estimation](https://mlsanity.com/paper/2106.03765/) (2021)
- [Perfect Match: A Simple Method for Learning Representations For Counterfactual Inference With Neural Networks](https://mlsanity.com/paper/1810.00656/) (2018)
- [Some methods for heterogeneous treatment effect estimation in high-dimensions](https://mlsanity.com/paper/1707.00102/) (2017)
- [Architecture Agnostic Federated Learning for Neural Networks](https://mlsanity.com/paper/2202.07757/) (2022)
- [Combining multiple observational data sources to estimate causal effects](https://mlsanity.com/paper/1801.00802/) (2018)
- [In Search of Insights, Not Magic Bullets: Towards Demystification of the Model Selection Dilemma in Heterogeneous Treatment Effect Estimation](https://mlsanity.com/paper/2302.02923/) (2023)
- [Federated Causal Inference in Heterogeneous Observational Data](https://mlsanity.com/paper/2107.11732/) (2021)
- [Transfer Learning on Heterogeneous Feature Spaces for Treatment Effects Estimation](https://mlsanity.com/paper/2210.06183/) (2022)
- [An Adaptive Kernel Approach to Federated Learning of Heterogeneous Causal Effects](https://mlsanity.com/paper/2301.00346/) (2023)
- [Federated Adaptive Causal Estimation (FACE) of Target Treatment Effects](https://mlsanity.com/paper/2112.09313/) (2021)

## Built on by (most cited first, 2 of 2)

- [Federated Causal Inference in Healthcare: Methods, Challenges, and Applications](https://mlsanity.com/paper/2505.02238/) (2025)
- [A Survey on Federated Causal Discovery and Inference](https://mlsanity.com/paper/2606.23741/) (2026)

## Abstract (arXiv)

Machine learning methods for estimating heterogeneous treatment effects (HTE) facilitate large-scale personalized decision-making across various domains such as healthcare, policy making, education, and more. Current machine learning approaches for HTE require access to substantial amounts of data per treatment, and the high costs associated with interventions makes centrally collecting so much data for each intervention a formidable challenge. To overcome this obstacle, in this work, we propose a novel framework for collaborative learning of HTE estimators across institutions via Federated Learning. We show that even under a diversity of interventions and subject populations across clients, one can jointly learn a common feature representation, while concurrently and privately learning the specific predictive functions for outcomes under distinct interventions across institutions. Our framework and the associated algorithm are based on this insight, and leverage tabular transformers to map multiple input data to feature representations which are then used for outcome prediction via multi-task learning. We also propose a novel way of federated training of personalised transformers that can work with heterogeneous input feature spaces. Experimental results on real-world clinical trial data demonstrate the effectiveness of our method.