Skip to main content
Back to timeline
arXivSource publication:

Mediation analysis stops refitting per table: MediPFN learns interval error quantiles in advance and runs 59.1 times faster than bootstrap in paired timing

Synopsis

Causal mediation analysis splits a treatment effect into an indirect pathway through mediators and a direct pathway through other processes, and attaching a confidence interval to the indirect effect usually means refitting several nuisance functions on every table and resampling repeatedly. MediPFN adapts a pretrained tabular foundation model to predict those nuisance functions, then a separate head learns quantiles of the estimator's standardized error from simulated tables with known indirect effects, with symmetrization and independent calibration producing the interval. On synthetic, distribution-shift and real-covariate tables its coverage sits closer to the 95% nominal level than several comparators, and paired timing gives a median 59.1-fold speedup over bootstrap.

Source-provided article image: MediPFN: Amortized In-Context Nuisance Learning and Inference for Causal Mediation Analysis
arXiv

Interpretation

Mediation analysis can now reuse a pre-learned nuisance predictor and error distribution instead of fitting conditional means and propensity scores separately for each new table and estimating uncertainty by resampling. Conventional implementations fit the potentially nonlinear functions for treatment assignment, mediator distribution and outcome response separately in each study, and the uncertainty of the indirect-effect estimate depends on how their estimation errors propagate. On synthetic data, MediPFN reaches coverage of 94.01%, 92.97% and 91.67% in the in-distribution, source-shift and composition-shift settings and 94.14% on the first real-covariate set, closest to 95% among the displayed methods in each setting.

A frozen tabular foundation model receives a new table's observed labels only through a lightweight adaptation head, and a second head emits quantiles of the estimator's error directly. The adaptation head learns a small additive adjustment to each standardized feature, and the output stage carries separate scale and offset for classification and regression, adjusting log-odds before mapping back to a propensity score and rescaling and shifting a predicted mean. The error head takes a 1,024-dimensional pooled representation plus 14 table and estimator summaries, is trained on 10,000 mechanism-balanced synthetic tables, fine-tuned on 1,536 semi-synthetic tables built from eight OpenML covariate sources, and calibrated on 256 tables from two separate OpenML sources to set the interval expansion.

Perspective

Evaluation covers synthetic tables, source-shift and composition-shift tables, 576 stress tables, and semi-synthetic tables generated by resampling covariates from eight OpenML datasets; identification rests on the three standard causal assumptions of consistency, positivity and sequential exchangeability.

Performance under source shift, small samples and heavy-tailed outcomes shows that coverage and interval length need to be assessed together; the authors list broader simulation priors, estimation uncertainty under wider distribution shift, other causal settings such as instrumental-variable models, and sensitivity analysis for unmeasured mediator-outcome confounding as future directions.