Research map / Probabilistic and scientific modeling / Treatment effects
LLM-Driven Treatment Effect Estimation Under Inference Time Text Confounding
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 TCA, a framework that estimates treatment effects under inference-time text confounding by training on LLM-generated text surrogates of structured confounders with a doubly robust pseudo-outcome regression.
Problem
Treatment effect estimation is trained on structured confounders but at inference time only partial confounder information is available via text descriptions, causing biased CATE estimates.
Approach
Identify the text-conditioned CATE as E[tau^x(X) | T=t] and learn it by regressing doubly robust pseudo-outcomes (built from structured confounders at training) on LLM-generated text surrogates of confounders, so that at inference only text is needed.
Key results
- TCA achieves PEHE 0.116 on IST and 0.179 on MIMIC-III versus best baseline TBE-CFR at 0.192 and 0.274.
- With GPT-3.5 Turbo generating text, TCA reaches 0.141 (IST) and 0.205 (MIMIC-III), still beating all baselines.
- TBE-DR baseline (same DR learner on text) reaches 0.187 and 0.266, far behind TCA, indicating the confounding adjustment rather than DR drives gains.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use TCA's confounding-adjustment framework instead of direct text-based CATE estimation | CATE estimation · IST | PEHE: 0.192 -> 0.116 | Helped |
| Use TCA's confounding-adjustment framework instead of direct text-based CATE estimation | CATE estimation · MIMIC-III | PEHE: 0.274 -> 0.179 | Helped |
| Use TCA rather than TBE-CFR | CATE estimation · IST and MIMIC-III demographic subgroups | PEHE: IST: 0.191 -> 0.115; 0.193 -> 0.116; 0.195 -> 0.117; 0.191 -> 0.115; MIMIC-III: 0.274 -> 0.179; 0.275 -> 0.180; 0.275 -> | Helped |
| Use TCA rather than TBE-CFR with GPT-3.5 Turbo text generation | CATE estimation · IST | PEHE: 0.216 -> 0.141 | Helped |
| Use TCA rather than TBE-CFR with GPT-3.5 Turbo text generation | CATE estimation · MIMIC-III | PEHE: 0.299 -> 0.205 | Helped |
| Add TCA's confounding-adjustment framework to text-based DR estimation | TBE-DR · IST | PEHE: 0.187 -> 0.116 | Helped |
| Add TCA's confounding-adjustment framework to text-based DR estimation | TBE-DR · MIMIC-III | PEHE: 0.266 -> 0.179 | Helped |
| Vary the prompt strategy among factual, narrative, and symptom-focused prompts (ablation) | TCA · IST and MIMIC-III | PEHE | No clear effect |
| Increase the strength of true confounding (ablation) | TCA · IST and MIMIC-III | PEHE | No clear effect |
Surprising findings
- Naive text-based DR learner (TBE-DR) performs much worse than TCA, showing an off-the-shelf combination with text generation is biased.
Limitations
- Relies on semi-synthetic outcomes, since counterfactuals are never observed in real data.
- Assumes LLM-generated surrogate text approximates the true inference-time text distribution; distribution shift is not validated against real patient text.
- Potential bias from LLM-generated narratives, only checked across gender and age subgroups.
Experimental setup
Datasets: International Stroke Trial (IST) MIMIC-III
Benchmarks: International Stroke Trial (IST) semi-synthetic MIMIC-III semi-synthetic
Metrics: PEHE
Baselines: TBE-S-learner TBE-T-learner TBE-TARNet TBE-CFR TBE-DR
Builds on: DR-learner CATENets CFRNet TARNet
Closest papers by meaning
- End-To-End Causal Effect Estimation from Unstructured Natural Language Data (NeurIPS 2024 · similarity 0.88)
- A Non-parametric Direct Learning Approach to Heterogeneous Treatment Effect Estimation under Unmeasured Confounding (NeurIPS 2024 · similarity 0.87)
- KG-TREAT: Pre-training for Treatment Effect Estimation by Synergizing Patient Data with Knowledge Graphs (AAAI 2024 · similarity 0.86)
- Improving the Generation and Evaluation of Synthetic Data for Downstream Medical Causal Inference (NeurIPS 2025 · similarity 0.86)
- Learning Representations of Instruments for Partial Identification of Treatment Effects (ICML 2025 · similarity 0.86)
- Meta-Learners for Partially-Identified Treatment Effects Across Multiple Environments (ICML 2024 · similarity 0.85)
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.