Title: LLM-Driven Treatment Effect Estimation Under Inference Time Text Confounding
Venue: NeurIPS 2025
Accepted: yes
Keywords: Treatment Effect Estimation

LLMs

Text Confounding

Causal Inference

CATE

Doubly Robust
Abstract: Estimating treatment effects is crucial for personalized decision-making in medicine, but this task faces unique challenges in clinical practice. At training time, models for estimating treatment effects are typically trained on well-structured medical datasets that contain detailed patient information. However, at inference time, predictions are often made using textual descriptions (e.g., descriptions with self-reported symptoms), which are incomplete representations of the original patient information. In this work, we make three contributions. (1) We show that the discrepancy between the data available during training time and inference time can lead to biased estimates of treatment effects. We formalize this issue as an \emph{inference time text confounding} problem, where confounders are fully observed during training time but only partially available through text at inference time. (2) To address this problem, we propose a novel framework for estimating treatment effects that explicitly accounts for inference time text confounding. Our framework leverages large language models (LLMs) together with a custom doubly robust learner to mitigate biases caused by the inference time text confounding. (3) Through a series of experiments, we demonstrate the effectiveness of our framework in real-world applications.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigate the CATE estimation problem with text information during inference time while the complete confounder is available during training time. To bridge this gap, the authors first train a doubly robust estimator using the training dataset. Then, the author construct text description of the confounders in the training data. Based on the generated text, the algorithm constructs a mapping from the text to the estimated treatment effect. Experiemental results on International Stroke Trial and MIMIC-III demonstrate the effectiveness of the method.\n\nSTRENGTHS AND WEAKNESSES:\nThe paper is well written and presented. The experimental results demonstrate the significant advancement of the proposed method. However, I think some weaknesses still exist as follows:\n\n1. The motivation is somehow weak. Because the text is much easier to obtain than complete measurement of confounders, it is not realistic that during training time, the text is unavailable.\n\n2. The novelty is not sufficient. The estimation method is the off-the-shelf doubly robust method. The incremental part is only the text generation and mapping function.\n\n3. The text generation based on LLM may be misaligned with the real text collection. Therefore, the rationality of the mapping function from the text to treatment effect is not convincing.\n\nQUESTIONS:\n1. The Assumption 3.2 is disjointed with the other parts of the paper. How to understand the noise variable in the text generation mechanism?\n\n2. How to obtain the text of the test set in MIMIC-III and IST?\n\n3. The proposed TCA method use doubly robust method as the estimation. However, the baselines does not use doubly robust method. For fair comparision, I suggest to supplment the TBE-Doubly Robust as the baseline.\n\nLIMITATIONS:\nThe authors present the limitations in Appendix G. As the authors claimed, the quality of the generated text and representation is important to the effectiveness of the method. And the other point is also important. Since the generated text contains less information than the confounder variable $X$, the single value of $T$ refers to multiple value of $X$. However, this fact is not considered in the method and a deterministic mapping is learned. Therefore, intergrating uncertainty quantification is necessary to this method.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a method for causal effect estimation when confounder values are observed as textual description at test/inference time. The proposed problem setup and the solution proposed are novel and are first of its kind.The proposed method has real-world use cases especially in healthcare domains. The proposed idea identifies conditional average treatment effect conditional on textual confounders using the conditional average treatment effect conditional on actual/numerical values of confounders. Experimental results presented show improvements over naive baselines.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1. The proposed problem setup of causal effect estimation under test time text confounders is novel, and has real-world significance.\n2. The proposed solution is intuitive and theoretically grounded.\n3. The experimental results show good performance compared to baselines in all experimental setups.\n\nWeaknesses:\n\n1. In the abstract the following line overstates the contributions as it is well known fact.\n\n\"We show that the discrepancy between the data available during training time and inference time can lead to biased estimates of treatment effects.\"\n\n2. A major issue with the current problem setup is that all covariates X are considered to be confounders. This holds very rarely in real-world. Many variables should not be used for adjustment as there may be covariates that are descendants of the treatment variable.\n\n3. The significance of Lemma 4.1 is unknown because it is obvious that the bias of naive estimator is positive.\n\n4. Finally, the proposed method tries to estimates $\\tau^t(t)$ instead of $\\tau^x(x)$. In my opinion, identifying $\\tau^x(x)$ at test time would be even more interesting. At least comparing the estimated values of  $\\tau^t(t)$ and $\\tau^x(x)$ in \"in-sample\" data and checking that they do not deviate much from each other would be a great addition to the paper and validates the idea of estimating $\\tau^t(t)$.\n\nQUESTIONS:\nPlease see the weaknesses section and provide your responses.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nEstimating treatment effects is an important task in many empirical sciences such as healthcare and political science. However, the information available when we seek to estimate treatment effects may not always be sufficient. Specifically, at inference time, we may only have access to text data that are noisy proxies of the actual confounders between treatment and outcome. The authors propose a novel methodology for estimating treatment effects when practitioners have access to only covariates at training time and only text data at inference time. The authors demonstrate the effectiveness of their method through semi-synthetic experiments.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\nOne clear strength of this paper is its clarity. The authors explain clearly how their problem setting differs from problem settings considered previously in the literature. They also give a clear example (Figure 1) of when covariate data may be available at training time but not at inference time. Throughout the paper, the authors also do a great job introducing their problem setup, why naive approaches fail, and the assumptions required for their novel TCA framework to be valid.\n\nThe observation that generating synthetic text data at training time to learn a function on how average treatment effects vary based on text data to aid in the estimation of treatment effects at inference time is an interesting observation. As large language models improve in quality, it is great to see methods such as the authors’ that use them without relying on large language models’ reasoning abilities, which has been shown to be limited and flawed.\n\nPrevious work that considers text data as a proxy for estimating treatment effects requires two proxies and two separate and independent pieces of text data. In the authors’ method, they only require one instance of text data. Even though their method requires a training dataset with access to the ground truth covariates, a condition not required in previous work, the authors’ method applies to a new, and realistic, scenario and is thus valuable to the community. Practitioners now have more methods to choose from depending on the specific dataset that they have.\n\nThe semi-synthetic experiments are promising and demonstrate how different prompts to large language models affect the downstream accuracy of the authors’ proposed method pipeline. Table 2 also demonstrates that potential biases that may be present in LLMs do not significantly impact the fairness of downstream estimates for the treatment effect.\n\nWeaknesses\n\nI have no major complaints regarding this paper, but I have some questions below that I hope the authors can answer and discuss.\n\nQUESTIONS:\nIn Appendix D describing implementation details for the semi-synthetic experiments, the authors state that they used OpenAI’s APIs to generate surrogate text data. This implies that the authors passed covariate data from MIMIC-III through OpenAI’s API, which violates MIMIC-III’s terms and conditions (https://physionet.org/news/post/gpt-responsible-use). Can the authors verify whether they passed protected information through OpenAI’s API? I will raise my score if the authors are able to address this potential ethics violation.\n\nWhat happens if the function $h$ for generating text data from covariates differs between training time and inference time? For instance, suppose the LLMs generate text data significantly differently from the way health practitioners write their notes. Would the authors expect downstream estimates for the treatment effect to be biased in this case? \n\nWhat happens if, at inference time, text data is affected by the treatment and/or outcome? For instance, suppose a practitioner used discharge notes as $T$ at inference time, then the text data may likely be a child of both $A$ and $Y$, and conditioning on $T$ will induce confounding between the treatment and outcome. Is this something that practitioners need to worry about when considering the authors’ method? What are steps that practitioners can take to mitigate this concern?\n\nIn datasets such as MIMIC-III, clinical text data in addition to covariate data may be available at training time. In such a scenario, would the authors’ proposed method still apply by using the available text data directly to learn $f_\\theta$? Would it still be necessary to generate surrogate text data $\\tilde{T}$ in this case? Would the variability of downstream estimates for the conditional average treatment effect be improved in this case?\n\nDoes line 198 have a typo? Instead of “Y(1) if the treated”, did the authors mean to write “Y(1) if treated”?\n\nDoes line 214 have a typo? The final sentence seems to be incomplete and does not have any ending punctuation.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper tackles the challenge (which they refer to as \"inference time text confounding\") when CATE estimation is done at inference time using text and confounders are not observed, but they are observed during training. The paper provides a framework for CATE estimation in this setting, and shows its real-world utility.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n* The paper is very clearly written and the problem setting is clear and well-formulated \n* Strong performance results compared to baselines. The robustness across varying conditions (confounder strengths and prompt strategies) enhances the credibility of results.\n* The paper is highly reproducible\n\nWeaknesses\n* The TCA method seems to be a straightforward application of existing methodologies (meta-learners and DR pseudo-outcomes)\n* While \"Inference Time Text Confounding\" is common in medical domains, it's not fully clear to me how this connects to CATE estimation. For example in Figure 1 I don't understand what CATE is. CATE estimation is useful when we are trying to decide on an intervention. But what is the *intervention*  in this case?\n* Stronger baselines could be included. For example, the framework seems conceptually very similar to the CaML framework, in that you are (1) getting pseudo-outcomes as a pre-processing step, then (2) generating embeddings (3) Predicting the pseudo-outcomes from the embeddings. How does the performance of would this compare to using CaML out of the box?  \n\n[1]https://arxiv.org/pdf/2301.12292\n\nQUESTIONS:\n(1) Can you please explain what are several real-world settings where inference-time text confounding applies, and specifically what are the interventions and outcomes we would care about that would motivate CATE estimation\n(2) How does your method compare to strong baselines (e.g., CaML, and SIN [2])\n\n[2] https://proceedings.neurips.cc/paper/2021/file/d02e9bdc27a894e882fa0c9055c99722-Supplemental.pdf\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-09"}