PaperlayerView exact revision on arXiv

Source-linked AI summary

Artificial intelligence surrogates for treatment effect estimation with before-and-after data

Frances Dean, Anna Neufeld, Joshua Barrios, Geoffrey H Tison, Ahmed Alaa

arXiv:2609.27180v1stat.MLcs.LG

TL;DR

Costly or long-term clinical outcomes make treatment-effect estimation difficult, and identifying useful surrogate outcomes can also be challenging. The paper estimates effects from before-and-after measurements using AI-predicted outcomes, with a labeled-data correction when surrogacy assumptions cannot be justified. In synthetic experiments, the corrected approach has narrower confidence intervals than difference in means when less than all data are labeled, even with a poor surrogate.

  • Problem

    Clinically important outcomes can be costly to measure or require long follow-up, while surrogate biomarkers may be unavailable or difficult to identify.

    Source: §1 · p. 1

  • Method

    The framework applies an AI model trained on untreated data to treated patients’ pre- and post-treatment measurements and compares their predicted outcomes.

    Source: §1 · p. 2

  • Results

    With less than 100% of data labeled, the PPI-debiased AI surrogate has a narrower confidence interval than difference in means even with a poor surrogate.

    Source: §6.1 · p. 11

  • Takeaways & Limitations

    AI-based surrogates can still support monitoring individual response alongside confirmatory tests, even without true surrogacy.

    Source: §5 · p. 10

    Surrogacy and comparability are untestable, and without true-outcome labels the magnitude of bias from their failure cannot be quantified.

    Source: §7 · p. 15

Abstract

from arXiv · show

Estimating the causal effects of medical treatments is difficult when clinically important outcomes are costly to measure or require long follow-up. Short-term or inexpensive surrogate outcomes offer a potential alternative, but surrogate biomarkers may be unavailable or difficult to identify. Advances in artificial intelligence (AI) have enabled increasingly accurate prediction of clinical outcomes from inexpensive, high-dimensional measurements, which creates an opportunity to use AI predictions themselves as surrogates. To this end, we develop a framework for estimating treatment effects from paired measurements obtained before and after treatment for each treated individual. A pretrained AI model is applied to the before and after measurements, and our estimator compares the resulting outcome predictions. We characterize the technical assumptions under which this within-person contrast identifies the average treatment effect on the treated, even when clinical outcomes are never observed for treated individuals. When these assumptions cannot be justified, we use prediction-powered inference to correct bias using a small number of observed clinical outcomes and obtain valid inference. Synthetic and real-world cardio-oncology experiments demonstrate the validity and accuracy of the approach.

1 Introduction

Important clinical outcomes can be costly, invasive, or slow to observe, motivating the use of AI predictions as potential surrogates. The paper proposes a before-and-after framework and emphasizes that predictive accuracy alone does not establish causal validity.

1 Introduction

  • Clinical endpoints such as mortality and advanced disease progression are often costly, invasive, or long term.Earlier, more accessible measurements are therefore often used as proxies for later outcomes.Source: §1 · p. 1
  • High predictive accuracy is not enough: treatment effects must operate through mechanisms captured by the AI model.The authors note that this relationship may be absent or unclear for some models.Source: §1 · p. 2
  • Traditional surrogate outcomes have sometimes misrepresented treatment effects, motivating careful evaluation of AI-based surrogates.The passage cites examples involving bone mineral density, arrhythmias, and accelerated drug approvals.Source: §1 · p. 2
  • AI predictions are already being used in causal-effect studies, and an AI-based pathology endpoint has been FDA approved for trials.The passage identifies AIM-MASH as an example.Source: §1 · p. 2
  • The framework compares AI-predicted clinical outcomes from treated patients’ measurements before and after treatment to estimate treatment effects.The model is trained on untreated patients with measurements and clinical outcomes.Source: §1 · p. 2

2 Problem Setup

The setup pairs treated patients’ pre- and post-treatment measurements with a longitudinal, treatment-free training dataset that includes outcomes. A time-aware AI model predicts outcomes from each measurement, and the within-person prediction contrast serves as the surrogate-based treatment-effect estimate.

2.1 Mathematical Notation

  • Treated patients contribute measurements before and after treatment but no clinical outcome labels, while the training population has longitudinal measurements and outcomes.The treated observations are denoted Dobs and the treatment-free training data Dtrain.Source: §2.1 · p. 3
  • The target is the average treatment effect on the treated, although neither treated patients’ untreated potential outcomes nor their observed treated outcomes are available.The passage says this prevents estimation by covariate adjustment with a matched control group.Source: §2.1 · p. 3

2.2 Causal Effect Estimation with Before-and-after AI Surrogates

  • A time-aware model trained on untreated longitudinal data predicts a future clinical outcome from a measurement and its time horizon.For example, an ECG model can predict heart-failure risk over a specified interval.Source: §2.2 · p. 4
  • The estimator compares predictions from pre- and post-treatment measurements at the same future endpoint, treating each patient as their own control.The pre-treatment prediction approximates the counterfactual trend; the passage states that when bµ = µ, this approach identifies the treatment effect.Source: §2.2 · p. 4

3 Causal Identification with Before-and-after AI Surrogates

The paper gives conditions under which before-and-after AI surrogates identify the ATT, including standard causal assumptions, assumptions about how measurements predict outcomes over time, and assumptions about surrogacy and comparability. These conditions are strong, and identification does not follow when any are absent.

Standard Assumptions

  • Consistency requires observed measurements and outcomes to match their potential values under the treatment received, while positivity requires treated measurements to occur with positive probability in training data.The passage identifies consistency as untestable and positivity as testable through overlap of covariate supports.Source: Standard Assumptions · pp. 4–5

Assumptions on Training Data

  • Training-data assumptions require later measurements to subsume earlier predictive information and pre-treatment measurements to suffice for predicting the untreated outcome.The paper describes these as Markov covariate sufficiency and counterfactual baseline sufficiency, respectively.Source: Assumptions on Training Data · pp. 5–6

Untestable Assumptions for AI Surrogacy

  • Surrogacy requires the measurement to capture treatment-induced changes relevant to the endpoint, while comparability requires the untreated measurement–outcome relationship to transfer between datasets.These assumptions are untestable when treated clinical endpoints are unavailable, and strong prediction alone does not establish surrogacy.Source: Untestable Assumptions for AI Surrogacy · pp. 6–7

3 Causal Identification with Before-and-after AI Surrogates

4 Estimation and Inference of Treatment Effects

The paper estimates treatment effects by contrasting AI predictions for pre- and post-treatment measurements. It describes when prediction errors may cancel and proposes a labeled-data correction for settings where surrogacy fails, while warning that the uncorrected inference is invalid when its assumptions fail.

4 Estimation and Inference of Treatment Effects

  • The basic estimator averages each treated patient’s difference between predicted post-treatment and counterfactual pre-treatment outcomes.The predictions serve as surrogate outcomes, but model errors can bias the estimate.Source: §4 · p. 7
  • When average prediction bias is constant across pre- and post-treatment settings, the biases cancel under the stated assumptions, allowing valid inference.The passage also requires assumptions on the data distributions and model bias.Source: §4 · p. 8
  • Without perfect satisfaction of the assumptions, inference from the basic estimator is invalid.The passage notes that errors may shift with measurement distributions and prediction horizons, and surrogacy may fail.Source: §4 · p. 8
  • With a small labeled treated sample and untreated outcome labels, a prediction-powered correction can address model-error differences and failed surrogacy.The passage says this second estimator can estimate an unbiased average treatment effect under its assumptions; the label samples must be independent and not used to train the model.Source: §4 · pp. 8–9

5 Heterogeneous Effect Estimation

Under Assumptions 1–6, the before-and-after approach identifies conditional treatment effects for treated individuals, but unbiased individual effects are not available without matching measurements and labels. It can still support subgroup averages and, for known mechanisms, individual response monitoring.

5 Heterogeneous Effect Estimation

  • Under the same Assumptions 1–6, the approach identifies CATT(x), the average treatment effect for a treated individual conditional on the pre-treatment measurement.The section defines CATT(x) as this conditional effect and states that the approach can identify it.Source: §5 · pp. 9–10
  • With sufficient counts, low-dimensional feature groups can yield average heterogeneous effects, but the approach cannot provide unbiased individual effects.Without averaging over people with the same pre-treatment measurement alongside labels, model and surrogacy biases affect heterogeneous-effect estimates.Source: §5 · p. 10
  • AI surrogates can monitor individual responses alongside confirmatory tests or estimate effects on known mechanisms, but unknown mechanisms require labeled outcomes.The section cautions that risks involving mechanisms missed by the AI surrogate cannot be predicted for individuals without unknown bias.Source: §5 · p. 10

6 Experiments

The experiments test before-and-after AI surrogates in synthetic settings and cardio-oncology. With some labeled outcomes, prediction-powered inference (PPI) improves precision over difference-in-means for a poor surrogate and corrects surrogate bias; in real-world analyses, debiased estimates are consistent with comparator estimates over time.

6.1 Synthetic Experiments

  • Synthetic experiments compared linear and nonlinear processes across surrogate strengths, using treated and untreated samples and repeated trials to assess estimation and interval performance.The study trained on untreated trajectories and varied the labeled fraction and surrogate quality.Source: §6.1 · pp. 10–11
  • When surrogate quality degraded, the naive estimator's bias grew, while PPI corrected it; PPI interval widths decreased as surrogate quality and outcome prediction improved.The synthetic experiments varied linear and nonlinear processes and surrogate correlations, using 5% labeled data for this comparison.Source: §6.1 · p. 11, Table 1 · p. 11
  • With fewer than 100% of outcomes labeled, PPI gave narrower confidence intervals than difference-in-means for a poor surrogate and greater precision than labeled-only estimates for a great surrogate.At 100% labeled data, the authors note that the AI surrogate is unnecessary.Source: §6.1 · pp. 11–12, Table 2 · p. 12, Table 3 · p. 12

6.2 Real World Treatment Effect Estimation in Cardio-Oncology

  • The real-world ECG analysis found large distributional overlap with the training population and minimal predictive-performance differences when adding earlier or later ECG data to check assumptions.The heart-failure model's mean AUROC was 0.847 (SE 0.00276); average AUROC differences were 0.016 (SE 0.005) and 0.013 (SE 0.006) for the earlier- and later-data comparisons, respectively.Source: §6.2 · p. 13, Figure 1 · p. 13
  • In cardio-oncology cohorts, debiased AI-surrogate heart-failure effects were consistent with comparator estimates over time; the naive surrogate better recovered anthracycline than trastuzumab effects.The analysis used labels for 50% of patients and compared AI estimates with covariate-adjusted estimates using patients treated with other therapies as controls.Source: §6.2 · p. 13

7 Discussion

Before-and-after AI surrogates may speed outcome monitoring and pharmacovigilance without a control cohort, but their causal validity depends on strong assumptions. The cardio-oncology experiments find that debias-corrected AI-surrogate results are consistent with analyses using long-term outcomes and covariate adjustment.

Cardio-oncology results

  • In breast cancer patients, debias-corrected AI-surrogate results using 50% of labels are consistent with analyses using long-term outcomes and covariate adjustment.The experiments include anthracycline-treated patients (N=644) and trastuzumab-treated patients (N=195).Source: Figure 2 · p. 14

Practical implications

  • The framework needs no control cohort, a feature the authors identify as useful for monitoring treated patients and pharmacovigilance.They describe running AI models on electronic medical records to detect changes in risks for multiple outcomes after new treatments.Source: §7 · p. 14

Related frameworks

  • The approach complements Athey et al. (2019): it is suited to longitudinal data when unconfoundedness cannot be assumed, whereas their framework suits clinical trials and settings with controls.The authors also distinguish their new treatment-effect estimator from Ji et al. (2026), which augments and recalibrates existing estimators using predictions.Source: §7 · p. 14

Cardio-oncology results

  • The uncorrected AI surrogate performs better for anthracyclines than for trastuzumab, while debias correction gives valid inference in either case.The authors relate the difference to the model apparently capturing anthracyclines’ irreversible oxidative-stress mechanism better than trastuzumab’s primarily asymptomatic reversible mechanism.Source: Figure 2 · p. 14, §7 · p. 15

Limitations

  • Without labels on true outcomes, the magnitude of bias from failed assumptions cannot be quantified; surrogacy and comparability are untestable.The authors caution that strong predictive performance alone does not establish a good treatment-effect estimator.Source: §7 · p. 15

8 Additional Methods

The additional methods describe the real-world ECG data, prediction model, analyses, and study resources. The model uses ECGs and clinical covariates to predict annual heart-failure risk.

8 Additional Methods

  • The methodology’s formal proof is in the appendix, while this section presents real-world study details; experiments used Python 3.11.13 and had IRB approval 19-27430.Source: §8 · p. 15

8.1 Real ECG Data

  • The study uses 491,090 UCSF twelve-lead ECG waveforms collected from 1980 through 2021, down-sampled to 250 Hz and z-score normalized.EHR diagnosis codes and cardiologist-adjudicated MUSE diagnoses provide disease labels; the data split was 70% training, 10% validation, and 20% testing.Source: §8.1 · p. 15

8.2 AI ECG Model

  • The ECG model combines a convolutional ResNet backbone with age, race and ethnicity, sex, and cardiovascular risk factors to predict hazards at 15 years.The backbone has four stages with 3, 4, 6, and 3 residual blocks; summed hazards are exponentiated to calculate disease risk by year.Source: §8.2 · pp. 15–16

8.3 Analyses

  • Because ground-truth treatment effects are unavailable, the before-and-after analysis is compared with covariate-adjusted linear regression against other-treatment controls.The regressions adjust for demographics, cardiovascular risk factors, and other cancer treatments.Source: §8.3 · p. 16

Data Availability

Code Availability

Contributions

  • FD conceived the study and experiments; GT and JB led real-data collection, while FD, AN, and AA led statistical methods.FD ran experiments and wrote the first draft; all authors contributed to writing.Source: Contributions · p. 16

A Connection to Classical Surrogate Inference, Athey et al. (2019), and Ji et al. (2026)

This section contrasts the paper’s longitudinal AI-surrogate framework with classical surrogate inference and related methods. It distinguishes the assumptions and inferential roles of the approaches, including prediction-powered recalibration.

Classical surrogate inference

  • The Prentice surrogate defines a surrogate through testing: no relationship between X and treatment A must also be a valid test of no relationship between A and outcome Y.Source: A Connection to Classical Surrogate… · p. 20
  • The paper’s AI surrogate regresses outcomes on high-dimensional measurements, which under an analogous assumption identifies treatment-effect magnitude, unlike the one-dimensional Prentice surrogate’s direction-only estimate.The authors say Prentice surrogates do not eliminate the surrogate paradox without additional monotonicity guarantees.Source: A Connection to Classical Surrogate… · p. 20

Athey et al. (2019)

  • Athey et al. (2019) use high-dimensional surrogate indices to identify effects in contemporaneous treated and untreated groups under comparability and surrogacy assumptions.Their case-control framework requires unconfoundedness; the present framework adds individualized effects using patients as their own controls and relaxes that requirement.Source: A Connection to Classical Surrogate… · pp. 20–21

Ji et al. (2026)

  • Ji et al. (2026) extend prediction-powered inference with a recalibration step that uses labeled data to address poorly calibrated predictions.Their RePPI method is reported to improve statistically over using labeled data alone, and the authors suggest it might also improve before-and-after surrogate inference.Source: A Connection to Classical Surrogate… · pp. 21–22
  • Unlike Ji et al. (2026), which augments existing estimators and retains their identification requirements, this paper proposes a treatment-effect framework usable with or without correction.Source: A Connection to Classical Surrogate… · p. 22

Related work

Assumption comparison

  • Table 4 compares assumptions for case-control surrogate inference and the proposed approach; the caption specifically notes that Athey et al. (2019) requires unconfoundedness.Source: Table 4 · p. 22

B Identification Proof

Under Assumptions 1–6, treatment effects are identified by contrasting predicted future outcomes from treated patients’ post- and pre-treatment measurements, despite unobserved treated outcomes. Covariate-shift estimators extend the approach when treated and training populations differ, with an alternative using labeled treated outcomes when surrogacy is relaxed.

B Identification Proof

  • Identification requires consistency, positivity, covariate Markov sufficiency, counterfactual baseline sufficiency, comparability, and statistical surrogacy.Together, these assumptions support using measurements and outcome predictions from the intervention-free training distribution to identify effects for treated patients.Source: B Identification Proof · p. 23
  • Under Assumptions 1–6, ATT is the treated-population average of the difference between predicted future outcomes from post- and pre-treatment measurements.The same result identifies CATT by conditioning on the baseline measurement.Source: B Identification Proof · p. 23

C Covariate Shift Estimators

  • When pre-treatment treated data differ from training data, covariate-shift corrections are proposed, including a route that relaxes surrogacy using a small labeled treated sample.The no-treated-label estimator assumes Assumptions 1–6; the alternative uses treated labels and Assumptions 1–5.Source: C Covariate Shift Estimators · p. 25

C.1 Estimators with covariate shift assuming surrogacy

The covariate-shift estimators use weighted prediction-error corrections to address differences between treated and training populations. Their confidence intervals have asymptotic ATT coverage under stated density-ratio estimation and variance conditions, with a labeled-treated-data variant that does not require surrogacy.

C.1 Estimators with covariate shift assuming surrogacy

  • The correction weights training prediction residuals by wpre(Xt0) and changes measure from Ptrain to Pobs,post, yielding desired intervention-free expectations when pre- and post-distribution intervals match.This requires the training data to be intervention-free.Source: §C.1 · p. 26
  • Both covariate-shift prediction-powered confidence intervals have asymptotic coverage of ATT at level 1−α under the stated density-ratio and variance conditions.The surrogacy-based version requires the density-ratio estimation error to be o_p(N^-1/2); the version without surrogacy requires the same rate and finite true variance with consistent variance estimators.Source: §C.1 · p. 27, §C.2 · p. 29
  • The surrogacy-based estimator does not correct bias in the covariate-shift model.The density-ratio estimation error must satisfy the stated convergence condition; bias in that model remains a limitation.Source: §C.1 · p. 27

C.2 Estimators with covariate shift in the absence of surrogacy

  • Without surrogacy, a small labeled treated sample corrects post-treatment prediction error while covariate shift corrects pre-treatment error.This estimator assumes Assumptions 1–5 and uses labels for treated outcomes.Source: §C.2 · p. 28

D.1 Main text results

The section establishes inference for before-and-after AI-surrogate estimators, including a covariate-shift correction that relaxes surrogacy, and outlines a two-stage approach to heterogeneous effects.

D.1 Main text results

  • Under Assumptions 1–6 and constant model bias across time and samples, the estimator’s confidence interval has asymptotic coverage of the ATT.The coverage result also assumes finite true variance and consistent variance estimation.Source: §D.1 · p. 29

D.2 Covariate shift estimators

  • Proposition 8 presents a covariate-shift estimator without surrogacy; the reported PPI−L quantity equals ATT + op(N−1/2), and both-ends training adjustment splits Dtrain-like data into independent groups of size N/2.Source: §D.2 · pp. 33–34

E Two-Stage Estimation Approach for Heterogeneous Effects

  • The two-stage approach trains a surrogate model, then learns individual CATT predictions from baseline measurements for use in similar populations.It requires surrogacy assumptions, and its real-world effect-estimation accuracy cannot be directly evaluated.Source: E Two-Stage Estimation Approach for… · pp. 35–36
  • With strong assumptions, the approach can model heterogeneous treatment effects for new individuals without observing anyone’s true outcomes.The paper limits this conclusion to settings where the assumptions hold.Source: E Two-Stage Estimation Approach for… · p. 36

F.1 Synthetic Experiments

Synthetic experiments examine confounding and surrogate quality, while cardio-oncology experiments check sequential assumptions, covariate shift, and the quantity of treated labels used for correction.

F.1 Synthetic Experiments

  • The synthetic confounding experiments add a variable affecting treatment probability and outcome, with tables evaluating coverage, bias, and interval width across surrogate settings.Table 5 uses 5% labeled data; Tables 6 and 7 vary labeled-data fraction for fixed poor and great surrogates, respectively.Source: §F.1 · p. 37, Table 5 · p. 37, Table 6 · p. 37

F.2 Cardio-Oncology Experiments

  • Across 15 years, the average absolute AUROC differences were 0.016 (0.005) for adding a prior ECG and 0.013 (0.006) when evaluating the later ECG from the earlier time.These checks support the sequential assumptions in the intervention-free data.Source: §F.2 · p. 38, Figure 3 · p. 38
  • Covariate-shift correction gives nearly the same cardio-oncology results as the uncorrected analysis, with classifier AUROCs of 0.557 for anthracycline and 0.517 for trastuzumab.The authors report that the treatment data are not well distinguished from the test set.Source: Figure 4 · p. 39, §F.2 · p. 39
  • The cardio-oncology analyses vary treated-label quantities around the main-text 50% follow-up-label correction, accounting for fewer labels in later years due to loss to follow-up.Table 8 considers 25–75% of patients and applies real follow-up masks.Source: §F.2 · p. 39, Table 8 · p. 40, Figure 5 · p. 40
Loading 2609.27180v1…