Technical Summary: Joint Estimation of Treatment Effects from Incomplete Health Data

Problem Statement

The article addresses the critical challenge of estimating causal treatment effects from observational Electronic Health Records (EHRs) when two compounding difficulties are present: time-varying confounding and Missing-Not-At-Random (MNAR) missing values.

Standard methods for Target Trial Emulation (TTE) often fail in real-world EHR environments for the following reasons:

1. Time-varying confounding: Treatment decisions are driven by evolving patient profiles, creating complex dependencies that existing methods often handle only under restrictive parametric assumptions or without explicit causal structure.
2. MNAR missing values: Critical biomarkers in EHRs frequently exhibit missing rates of 50% to 80%. These gaps are driven by clinical decisions (e.g., tests are only ordered if a condition is suspected), meaning the probability of missing data depends on the unobserved value itself.
3. Compounded bias: Standard imputation strategies (e.g., MICE, MissForest) typically assume Missing-At-Random (MAR) and are validated on low missing rates. When applied to high-rate MNAR data, they distort the correlation structures between biomarkers and outcomes, leading to biased causal estimates.

Existing approaches often treat imputation and causal inference as separate stages or rely on variational/adversarial approximations that introduce silent errors in counterfactual estimation.

Methodology

The authors propose a two-stage pipeline designed to jointly handle MNAR missing values and causal structures over time.

Stage 1: LLM-Driven Evolutionary Imputation

To address input data with gaps, the authors introduce an evolutionary search framework where a Large Language Model (LLM) acts as a mutation operator to generate imputation strategies.

• Mechanism: Instead of the LLM directly filling missing values, it proposes candidate Python modules (imputers) $g^{(k)}$.
• Self-supervised Proxy Score: The quality of each candidate is evaluated on a proxy holdout set (masked observed data) using a composite score $s(g)$:
  $$s(g) = \text{RMSE}(g) + \lambda_Y \Delta_Y(g) + \lambda_T \Delta_T(g)$$
  This score penalizes pointwise reconstruction error (RMSE) and, crucially, the distortion of covariate-outcome correlations ($\Delta_Y$) and covariate-treatment correlations ($\Delta_T$). This ensures that the imputed data preserve the causal structure required for downstream estimation.
• Evolutionary Loop: Starting from a seed imputer, the LLM iteratively suggests variations based on the current best candidate and a summary of the search history, accepting only those that strictly improve the proxy score.
• Backends: The method was evaluated using three LLM backends: GPT-5.4, GPT-OSS-120b, and Qwen3.5-Plus.

Stage 2: CausalFlow-T (Exact Counterfactual Inference)

The completed data $\hat{D}$ are fed into CausalFlow-T, a novel model for exact counterfactual inference.

• Architecture: It combines a Causal Masked Autoregressive Flow (CausalMAF) with a Long Short-Term Memory (LSTM) encoder.
• DAG Constraints: The flow is constrained by an expert-specified Directed Acyclic Graph (DAG). The autoregressive ordering of variables strictly follows the topological sort of the DAG. This ensures that interventions (do-operations) are propagated only through causal descendants, preventing anti-causal leakage.
• Exact Inference: Unlike Variational Autoencoders (CVAE) or GAN-based methods that optimize a lower bound (ELBO) or adversarial objectives, CausalFlow-T maximizes the exact log-likelihood.
• Counterfactual Procedure: It applies Pearl's Abduction-Action-Prediction (AAP) procedure:
  1. Abduction: Invert the flow exactly to recover the exogenous noise $z_t$ for a specific patient.
  2. Action: Intervene on the treatment variable ($A \leftarrow a'$) while keeping $z_t$ fixed.
  3. Prediction: Propagate the intervention through the DAG-constrained flow to generate the counterfactual outcome.
     This guarantees that the same patient-specific noise is used for both factual and counterfactual arms, a property variational methods cannot offer.

Main Contributions

1. CausalFlow-T: A DAG-constrained normalizing flow enabling exact, invertible counterfactual inference for longitudinal data. The authors demonstrate it is the only model satisfying five reliability criteria simultaneously (subgroup calibration, tail variance preservation, arm reconstruction, hazard ratio recovery, and stability).
2. Identification of Non-Overlapping Failure Modes: Through ablation studies, the article reveals two distinct failure modes in causal inference:
   • Separation of Confounding: Requires explicit causal structure (DAG constraints). Models without DAG constraints retain systematic biases even with low factual error.
   • Structural Propagation: Requires exact inference. Variational methods (even with correct DAGs) fail to recover true treatment effects under complex mediation paths due to approximation gaps (Posterior Collapse).
3. LLM-Driven Evolutionary Imputation: A novel imputation strategy that outperforms standard statistical methods (LOCF, MissForest) and flow-matching baselines in preserving causal correlation structures under 30%–80% MNAR missing values.
4. Joint Evaluation Framework: The article notes that pointwise reconstruction accuracy and causal fidelity are not aligned objectives. It introduces a two-tier evaluation (biomarker quality + downstream causal quality) to reveal trade-offs invisible to standard metrics.

Results

Synthetic and Semi-Synthetic Benchmarks

• CausalFlow-T Performance: On four synthetic datasets (including a complex 17-node CVD risk benchmark) and a semi-synthetic FIRE-based dataset, CausalFlow-T was the only model to pass all reliability criteria.
  • Competitors like GNN-CVAE and CVAE suffered from Posterior Collapse (predicting zero effects where real effects existed) or variance collapse.
  • TARNet and unrestricted flows failed to recover correct hazard ratios or invert treatment directions.
• Imputation Performance: The LLM-driven imputer (particularly GPT-5.4) achieved the best pooled rank across biomarker and causal metrics.
  • It maintained the lowest pointwise MAE and RMSE across 30%, 50%, and 80% missing values.
  • Crucially, it was the only strategy that preserved the recovery of the Average Treatment Effect (ATE), while statistical methods (MissForest, CausalCFM) progressively deteriorated as missing rates increased.
  • The study highlighted a trade-off: while CausalCFM sometimes preserved correlations better at lower missing rates, LLMs dominated in high-missing-rate regimes and temporal extrapolation.

Real-World Application

The pipeline was applied to a cohort of 6,114 adults with Type 2 Diabetes in Swiss primary care (FIRE database) to estimate the effect of initiating GLP-1 receptor agonists compared to SGLT-2 inhibitors on body weight after one year.

• Result: The pipeline estimated a per-protocol weight loss difference of −0.98 kg [95% CI −1.01, −0.96], favoring GLP-1RAs.
• Validation: This estimate is consistent with the SUSTAIN-8 Randomized Controlled Trial (RCT), which reported a difference of −1.06 kg, despite the real-world data exhibiting >50% MNAR missing values and mixed drug agents.

Significance and Claims

The article claims that jointly addressing causal structures and MNAR missing values is both necessary and feasible for the practical estimation of treatment effects from observational data.

• Beyond RCTs: The work demonstrates that ML-based pipelines can extend causal evidence beyond the restricted populations and conditions of RCTs to real-world clinical settings where studies are not feasible.
• Reliability over Accuracy: The authors argue that low factual prediction error (MAE) is insufficient for causal inference. Models must meet structural criteria (variance preservation, correct mediation propagation) to be reliable.
• Linking Imputation and Causality: The study notes that imputation decisions have direct, measurable impacts on causal validity, requiring joint evaluation metrics rather than treating imputation as a pure reconstruction task.
• Modest Limitations: The authors acknowledge limitations, including dependence on expert-specified DAGs (sensitivity analyses show bias increases with the distance of unique paths), the use of continuous relaxation for binary survival outcomes (affecting calibration), and the assumption of conditional exchangeability. They note that future work should investigate discrete flow architectures and automated DAG refinement.

In summary, the article offers a robust, two-stage framework that bridges the gap between incomplete, messy real-world EHR data and rigorous causal inference, validated by both synthetic ground truths and consistency with established RCT evidence.