Causal pretraining reduces error in estimating treatment effects
Learning to Fluctuate: Statistical Foundations for Causal Tabular Pretraining
Machine Learning
Summary
Estimating the effect of treatments or actions from data can be tricky, especially when patterns change or data is limited. The authors propose a new way to train models using labeled synthetic data that includes both average effects and their fluctuations, helping the model learn more accurate causal estimates. Their method improves accuracy substantially compared to previous approaches and can better handle shifts in effects across different datasets. They also offer mathematical guarantees and tests that show their approach captures sampling behavior well.
What this means in practice
- •For healthcare data teams: Estimate patient treatment effects more accurately when data varies or is limited by using fluctuation-supervised pretrained models.
- •For financial risk analysts: Improve causal effect estimates for interventions under changing market conditions with pretrained models informed by synthetic data fluctuations.
Authors
Zhiheng Zhang
Abstract
Causal tabular foundation models amortize effect estimation across synthetic mechanisms, but latent-effect supervision rewards posterior shrinkage instead of directly encoding the repeated-sample response needed in a fixed deployment population. We introduce fluctuation-supervised pretraining (FSP): each synthetic table is labeled by its average treatment effect plus its efficient influence-function fluctuation, while deployment remains a single frozen forward pass. Along the path $T_{λ,P}=θ(P)+λP_nψ_P$, we prove an endpoint transition: every fixed $λ<1$ retains label ambiguity of order $(1-λ)^2/n$, whereas full fluctuation makes the Gaussian label observable and reduces optimal finite-stratum causal label-prediction risk to order $n^{-2}$. One finite-pretraining bound combines label, network, episode-sampling, and optimization errors; its resulting sampling defect controls fixed-mechanism bias, mean squared error, variance, Gaussian approximation, and, with variance-head accuracy, studentized coverage. Complementary lower bounds separate the local $n^{-1}$ ATE risk that deployment observations cannot erase from the $\log N/M$ excess risk of a generic finite-dictionary episode-learning problem. Experiments trace the learned sampling response. With a raw-row/column backbone, FSP reduces large-effect-shift RMSE by 69.8% relative to latent supervision and by 39.5% relative to a released CausalPFN checkpoint on matched tables. Continuous-covariate experiments, known-effect semisynthesis and two randomized-study evaluations separate sampling-law fidelity from point-risk shrinkage and expose weak-overlap errors in both learned heads.