Research map / Probabilistic and scientific modeling / Treatment effects
Towards Causal Foundation Model: on Duality between Optimal Balancing and Attention
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes CInA, a transformer whose last self-attention layer is shown to recover optimal covariate balancing weights via a primal-dual (SVM) connection, enabling zero-shot average treatment effect estimation on unseen datasets.
Problem
Foundation models struggle with causal inference tasks such as treatment effect estimation, which need intricate reasoning and high numerical precision, and per-dataset estimators must be re-fit on every new dataset.
Approach
Shows that optimal adversarial covariate balancing is the dual of an SVM whose primal solution is a support-vector expansion equivalent to self-attention, so a transformer trained with a penalized hinge loss on multiple unlabeled datasets yields balancing weights (and ATEs) zero-shot by reading the last attention layer's values and kernel normalizer.
Key results
- On real-world datasets CInA achieves the best average ranking of 1.83 versus 3 for the second-best method DML, and CInA (ZS) returns the best result on ACIC.
- On ER-5000, CInA (ZS) reaches DML-level MAE (11.50 vs 11.13) while being about 10^2 times faster at inference; supervised CInA (ZS-S) reaches 2.66 MAE, outperforming all per-dataset baselines.
- On 10k ACIC 2018, CInA (ZS-S) achieves ATE MAE 1.361 versus 8.572 for DML and 13.07 for the naive estimator.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use CInA instead of RieszNet | CInA on Simulation-A · Simulation-A | ATE MAE: 0.045 ± 0.01 -> 0.126 ± 0.02 | Hurt |
| Use zero-shot CInA instead of SVM | CInA (ZS) on ER-5000 · ER-5000 | ATE MAE: 11.09 ± 3.13 -> 11.50 ± 1.85 | Hurt |
| Add supervised ATE training signals (ablation) | CInA (ZS) on ER-5000 · ER-5000 | ATE MAE: 11.50 ± 1.85 -> 2.66 ± 0.33 | Helped |
| Use CInA instead of RieszNet | CInA on IHDP · IHDP | ATE MAE: 0.110 ± 0.01 -> 0.114 ± 0.01 | Mixed |
| Use CInA (ZS) instead of DML | CInA on 10k ACIC 2018 · 10k ACIC 2018 | ATE MAE: 8.572 ± 8.96 -> 1.460 ± 0.48 | Helped |
| Add supervised ATE training signals (ablation) | CInA (ZS) on 10k ACIC 2018 · 10k ACIC 2018 | ATE MAE: 1.460 ± 0.48 -> 1.361 ± 0.42 | Helped |
| Pretrain on ER-5000 instead of ACIC datasets (ablation) | CInA (ZS) on 10k ACIC 2018 · 10k ACIC 2018 | ATE MAE: 1.460 ± 0.48 -> 1.718 ± 0.74 | Hurt |
| Pretrain on ER-5000 and add supervised ATE training signals (ablation) | CInA (ZS) on 10k ACIC 2018 · 10k ACIC 2018 | ATE MAE: 1.361 ± 0.42 -> 1.702 ± 0.74 | Hurt |
Surprising findings
- Unsupervised CInA (ZS) can match DML on ER-5000 despite never seeing outcomes during training.
- SVM baseline performs well in simulation but poorly on IHDP-resampled and ACIC, where hyperparameter selection on validation datasets does not transfer.
Limitations
- Efficiency: attention cost is O(N^2 D) in the number of units, so large datasets need efficient transformers.
- Performance is limited by the availability and diversity of training dataset resources; real-world multi-dataset training is small.
- Theory relies on regularity conditions (linear independence of features) and a unit-ball RKHS assumption; the work is restricted to average treatment effect estimation.
Experimental setup
Datasets: Simulation A synthetic ER-5000 synthetic Twins IHDP IHDP-resampled ACIC 2018 LaLonde CPS LaLonde PSID
Benchmarks: Simulation A ER-5000 IHDP IHDP-resampled Twins LaLonde CPS LaLonde PSID ACIC 2018
Metrics: mean absolute error of ATE run-time
Baselines: Naive IPW Self-normalized IPW DML SVM DragonNet RieszNet
Builds on: Tarr & Imai (2021) dual SVM balancing Nguyen et al. (2022) attention as kernel Kallus adversarial covariate balancing double machine learning
Closest papers by meaning
- Prediction-Powered Causal Inferences (NeurIPS 2025 · similarity 0.87)
- Causal Inference with Conditional Front-Door Adjustment and Identifiable Variational Autoencoder (ICLR 2024 · similarity 0.87)
- Smoke and Mirrors in Causal Downstream Tasks (NeurIPS 2024 · similarity 0.87)
- Causality can systematically address the monsters under the bench(marks) (NeurIPS 2026 · similarity 0.87)
- Neural Causal Abstractions (AAAI 2024 · similarity 0.87)
- Can Large Language Models Infer Causation from Correlation? (ICLR 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.