Federated AI Estimates Treatment Effects Without Sharing Patient Data

Estimating whether a drug, a surgery, or a policy actually works usually means pooling patient records from many hospitals into one central database. That pooling collides with privacy law and institutional caution, because the covariates needed to remove bias in observational data are exactly the kind of sensitive information organizations are reluctant to hand over. A team of researchers now reports a framework, called VFCI, that lets multiple institutions jointly train causal inference models over vertically partitioned data — the same patients, but different feature sets held by different parties — without any raw features or labels ever leaving their holders. The work, published in Data Science and Engineering, targets a setting the authors describe as rarely studied: causal inference under the vertical federated learning regime.

The problem the framework addresses is a fundamental one in causal inference. Observational data carry selection bias: some features, known as confounders, influence both the treatment a person receives and the outcome being measured. In the paper’s running example, a doctor assessing whether surgery or medication better aids patient recovery may systematically assign surgery to younger patients and medication to older ones, so age confounds the comparison. Adjustment variables that affect only the outcome, such as exercise, must also be accounted for. Estimating an individualized treatment effect — the expected difference between the outcome under treatment and under no treatment for a given instance — requires conditioning on all of these variables. In the vertical setting, those variables are scattered across parties: a pharmaceutical company evaluating a new drug for diabetic patients might hold genetic information, while several institutions hold different electronic health records of the same patients, and one designated label owner holds the treatment assignments and outcomes.

VFCI builds on split learning, a distributed paradigm in which each party holds a block of the global network rather than a copy of all data. Each client manages a bottom model that reads its local features and produces forward activations, typically hidden embedding vectors, which are the only things sent to the server. The server holds a middle model, aggregates the activations from all clients, and continues the forward computation. The label owner holds a top model, computes the loss using its private treatment and outcome labels, and sends intermediate gradients back to the server, which backpropagates and returns per-client gradients so the bottom models can update. Raw data never leave the clients at any point in this loop. Communication runs over remote procedure calls secured with Transport Layer Security, and the authors note the framework is not tied to TLS: they also provide an implementation using homomorphic encryption, in which a server-generated key pair lets clients and server encrypt and decrypt exchanged messages.

A key insight of the paper is that conventional split learning is not sufficient on its own. In a standard two-phase split neural network, the server hosts the label-dependent heads of the model, which means it observes logits and last-layer gradients that are directly label-informative. Prior label-inference attacks have exploited exactly this surface. VFCI instead adopts a three-phase split: the final layers of the top model, along with loss evaluation, are assigned to the label owner. The server therefore receives only intermediate gradients and never sees logits or last-layer gradients, which the authors identify as the tensors most often highly label-informative. The full training procedure unfolds in nine steps per epoch, from forward propagation of each bottom model through transmission of activations, forward propagation of the middle and top models, loss and gradient computation at the label owner, gradient transmission, and backpropagation through the middle and bottom models.

Because the training performs exact end-to-end backpropagation across the split points, the authors state that under the same model, data, optimizer, and hyperparameters, VFCI optimizes the same objective and produces the same parameter gradients as centralized training, changing only the wall-clock time through communication costs. This equivalence matters for credibility: the framework is not trading estimation quality for privacy in its core design. The workflow also generalizes. Treating the general training routine as an oracle, the authors show how to accommodate two-stage causal models alongside end-to-end ones. In a two-stage method such as DCN-PD, a propensity network is first trained to predict the probability of treatment assignment, and an outcome network is then trained with a propensity-dropout regularization scheme that penalizes subjects with extreme propensity scores more heavily, reducing selection bias. Under VFCI, all parties train two local models each — one for propensity, one for outcome — and the two stages run alternating instantiations of the general workflow, with the label owner transmitting propensity scores to the server between stages.

The security analysis covers a range of threat scenarios: a curious server, curious clients, server-client collusion, client collusion, and external attackers. External eavesdropping and link tampering are mitigated by secure channels such as TLS, while homomorphic encryption can protect intermediate messages in ciphertext form. The authors are explicit about the protection boundary. Raw features remain local, and the server and other clients observe only intermediate activations and gradients, which reduces direct exposure — but they caution that intermediate messages may still carry information and can be targeted by reconstruction attacks known from the split-learning literature. Likewise, they do not claim perfect label privacy, since intermediate messages used for backpropagation may still leak label-related information. For stronger protection, they note that standard defenses such as activation or gradient perturbation, gradient clipping, and secure computation on intermediate tensors can be integrated without changing the protocol boundary.

Empirically, the team instantiated two causal models in the framework: DragonNet, an end-to-end architecture with a shared representation layer and separate outcome heads that also predicts propensity scores, and DCN-PD, the two-stage model. Experiments ran on two pre-established causal benchmarking datasets. The IHDP dataset is semi-synthetic, with 747 instances — 139 treated and 608 control — and 25 features describing attributes of children and their mothers, studying the effect of specialists’ home visits on future cognitive scores. The ACIC 2018 datasets derive from linked birth and infant death data based on real-world medical measurements; the researchers randomly picked three datasets of 10,000 instances with 177 features. Both were partitioned vertically into two and four feature partitions, then split 80-20 into train and test sets. Each client held one partition, and all parties ran on AWS g4dn.xlarge EC2 instances, with Adam at a learning rate of 0.0001, batch size 32, and 50 training epochs.

Against three baselines — Local, where each client trains independently on its own features; Central, where all features and labels are pooled as an upper bound; and a two-phase SplitNN ablation — VFCI came close to the ideal. Across both models and both partition scenarios, the reported metrics — mean squared error of estimated outcomes, precision in estimation of heterogeneous effect, and absolute error in average treatment effect — were close to those of the Central methods, while VFCI achieved much smaller errors than Local, which can only use each client’s own data. Compared with the two-phase SplitNN ablation, VFCI reached comparable causal estimation accuracy under the same model capacity and training settings, while enabling the stricter privacy boundary. The authors also observed that as the number of participating clients increased from one to eight, all three metrics decreased, meaning the model became more accurate — a result they present as verification of the value of federated learning, since introducing more data sources benefits the trained model.

The label-inference experiments quantify what the three-phase design buys. Under an honest-but-curious server on IHDP, the researchers considered two attacker strategies: a logits-based attack that trains a lightweight classifier on server-side logits, applicable to two-phase SplitNN, and a gradient-based attack that trains on intermediate activations and backpropagation gradients, observable in both protocols. Reporting a composite attack accuracy averaged over the two, three-phase VFCI reduced it from 0.9716 under two-phase SplitNN to 0.7506, with the corresponding attack advantage — twice the accuracy minus one — dropping from 0.9432 to 0.5012. The improvement is attributed primarily to keeping label-dependent heads and loss evaluation at the label owner, so the server no longer observes the most label-informative tensors.

Efficiency results temper the privacy gains with practical costs. Per-epoch runtime increases with more clients due to higher communication volume and synchronization overhead. On IHDP with DragonNet, TLS incurred a 1.31-fold slowdown as clients grew from two to eight, mainly from communication overhead. Homomorphic encryption, which keeps intermediate activations and gradients confidential from the compute host, was consistently slower than TLS, with the gap widening as clients multiplied because of repeated encryption and decryption. The authors therefore recommend TLS as the default for efficiency and homomorphic encryption only when stronger protection is required, and they note that the three-phase split is orthogonal to standard remedies for synchronization bottlenecks, such as partial participation and bounded-staleness or asynchronous updates. The study’s scope carries the usual caveats of benchmark work: the datasets are semi-synthetic, the threat model assumes honest-but-curious participants who follow the protocol, and the authors themselves flag residual leakage channels in intermediate messages. Within that scope, the framework offers a concrete answer to a question that has barely been addressed — how to estimate treatment effects from multi-source, vertically distributed data while keeping the raw evidence of cause and effect where it belongs.

Subject of Research: A federated learning framework for estimating treatment effects from vertically partitioned data while preserving feature and label privacy.

Article Title: VFCI: A Federated Framework for Vertical Causal Inference

Article References: Zhou, X., Huang, H., Xu, Q., Hu, C., Chen, X., Wang, H., & Jiang, J. (2026). VFCI: A Federated Framework for Vertical Causal Inference. Data Science and Engineering. https://doi.org/10.1007/s41019-026-00357-7

Image Credits: AI Generated

DOI: 10.1007/s41019-026-00357-7

Keywords: federated learning, causal inference, treatment effect estimation, vertical federated learning, split learning, data privacy, label inference attack, DragonNet, DCN-PD, propensity score, homomorphic encryption, observational data