LLM-Driven Treatment Effect Estimation Under Inference Time Text Confounding

Thank you for your constructive review and your helpful comments!

Response to Weakness

W1

We appreciate the reviewer’s concern, but respectfully clarify that our setting—where text is unavailable during training but present at inference—is both realistic and increasingly common, especially in the medical domain.

In many real-world healthcare applications, it is common that structured data (e.g., demographics, vitals, labs) is reliably recorded and readily available during training, particularly in randomized controlled trials or publicly released clinical datasets. In contrast, unstructured text is often sparse or inaccessible, or incomplete, inconsistently formatted, and lacks standardization. Structured tabular data are carefully recorded and quality checked in clinical datasets, making them easily available. This is typically driven by regulatory need and approval for medical devices, where only high-quality information is preferred over self-reported texts. As a result, unstructured text is either unavailable, inconsistently documented. Models are often built on well-structured datasets from electronic health records or clinical trials.

Even when some text is present during training (e.g., clinical notes), it is often incomplete, inconsistently formatted, and often differs substantially in style and content from inference-time text (e.g., from a different distribution and with language style). For example, training-time notes are authored by clinicians, while inference-time inputs in telemedicine, chatbots, or symptom-reporting apps are often patient-generated. Such free-text symptom descriptions often fail to cover the full set of confounding variables.

Our setting, where models are trained on structured data but must make inferences on text data is both practically motivated and clinically relevant. Our work is the first to formally characterize the challenge of inference-time text confounding and to propose a concrete solution that is theoretically principled and empirically validated in such real-world conditions.

W2

We would like to clarify that our method is not the off-the-shelf doubly robust method. Our TCA differs substantially from the standard DR-Learner. A simple off-the-shelf combination of our text-generation step with existing DR-learners would be biased.

Our main contribution includes

(1) Problem formalization: We are the first to formally define the inference-time text confounding problem, which is distinct from standard causal inference settings. We show that naïve applications of existing methods (including DR learners) are biased in this setup, and we derive their pointwise confounding bias.

(2) Theoretical contribution: We propose an estimation framework that connects the CATE $\tau^x(x)$ w.r.t true confounder $X$ to the text-based CATE $\tau^t(t)$, and we prove identifiability and unbiasedness of our estimator under this framework. This is not a trivial extension of existing DR methods.

(3) Modular estimation strategy: While our approach leverages the doubly robust estimation to benefit from the doubly robust properties, it could also be easily adapted to use other meta-learners in the final step.

(4) Empirical validation: We further conduct new experiments where we combine our text-generation pipeline with a standard DR learner (referred to as TBE-DR). Importantly, we made the experiment fair: we used the same setup for the LLMs, etc. We report PEHE scores on two datasets.

=> Our TCA outperforms TBE-DR by a large margin.

This baseline uses the same underlying estimation technique as ours but lacks our framework for confounding adjustment. These results highlight that the core innovation lies in our confounding adjustment framework, not in the choice of DR itself.

W3

Thank you for raising this important point. We acknowledge that LLM-generated text may not perfectly align with real-world text distributions. Our method is explicitly designed to tolerate such mismatches. To assess robustness, we conducted domain shift experiments by varying the LLM prompt strategies across datasets (Figure 3(b)). Despite this shift, our method consistently outperforms all baselines, demonstrating robustness to the distribution shift between text data.

Response to Questions

Q1

Thank you for your question. Assumption 3.2 defines the data-generating mechanism of the induced text confounder $T$, which is generated from the true confounder $X$ as $T = h(X, \epsilon)$, where the noise variable satisfies $\epsilon \perp Y(a) \mid X$. This assumption formalizes the intuition that $T$ is a noisy (i.e., stochastic) transformation of the underlying true confounder $X$, which reflects the nature of real-world text generation in clinical settings. For example:

(1) Doctor’s notes or patient self-descriptions are often incomplete, informal, and subject to variability (noise) such as different writing styles, which are, however, independent of the outcome of interest.

(2) LLM-generated text, even when conditioned on structured inputs, introduces randomness (e.g., due to sampling variability or temperature settings, or seed, etc.).

The noise variable $\epsilon$ models this stochasticity in text generation, and allows diverse formats of the text description of the true confounder.

Importantly, We do not require that $T$ preserves all relevant information from $X$. It also fits in applications where $T$ often does not capture all the necessary confounding information in $X$ for potential outcomes. This aligns with the imperfect, noisy, and incomplete nature of text data in practice. Instead, $T$ may contain only partial information about $X$. As $T$ is induced by $X$, the useful information encoded by $T$ should be included within the information contained in $X$ for predicting the outcome.

Besides, the independence condition $\epsilon \perp Y(a) \mid X$ ensures that: Once we condition on $X$, the residual information in $T$ does not provide any additional information about the potential outcomes.

This assumption is more for technical reasons to ensure identifiability, that is, identifiability of $\tau^t(t)$ by adjusting for inference time text confounding, which allows us to express the CATE under text ($\tau^t(t)$) as a conditional expectation of the true CATE ($\tau^x(x)$), i.e., $\tau^t(t)=\mathbb{E}\left[\tau^x(X) \mid T=t\right]$. as shown in Lemma 4.3 and Appendix B.2. The assumption thus allows us to use our two-stage procedure for adjusting inference-time confounding when only $T$ is observed at test time.

Q2

The text for our experiments on top of MIMIC-III and IST is generated by LLMs to simulate text with self-reported symptoms. In the datasets, there is no text available that captures such symptom descriptions. Hence, we use LLMs to generate textual confounders $\tilde{T}$ from structured clinical covariates $X$. For each patient, we construct prompts by templating the key features of $X$ into natural language constraints. The generated texts are about 150–200 tokens long. We split the data into 80% train and 20% test.

Q3

Thank you for the suggestion. We added the TBE-Doubly Robust (TBE-DR) learner as another baseline. The corresponding results are reported above under W2.

Response to Limitations

Thank you for the thoughtful comment.

Assumption 3.2 formalizes the intuition that $T$ is a noisy (i.e., stochastic) transformation outcome of the underlying true confounder $X$. We do not assume the mapping from $X$ to $T$ is deterministic, one-to-one, or invertible. Nor do we require $X$ to be reconstructed from $T$ or that $T$ preserves all confounding information. A single $T$ may correspond to multiple possible values of $X$.

It is true that the generated text $T$ often contains less information than the true confounder $X$, and this is precisely why our framework is needed—to adjust for inference-time confounding when only $T$ is available to use at inference time.

To study this systematically, we conduct experiments under varying levels of information preserved in $T$, which we refer to as high-info, medium-info, and low-info settings, based on how much information from $X$ is retained in the text.

We also appreciate the reviewer’s suggestion on incorporating uncertainty quantification. We thus measure the discrepancy between estimated $\hat{\tau}^t(t)$ and ground-truth $\tau^x(x)$. This discrepancy captures the uncertainty introduced by the loss of information in the transition from $X$ to $T$, introducing variability in $\tau^t(t)$ relative to the true effect $\tau^x(x)$. For reasons of space, we focus on TBE-DRlearner as it is the strongest baseline. We report results on IST and MIMIC, respectively. These results demonstrate that our framework is more robust to uncertainty.

Thank you for your positive review and your helpful comments! We are very happy to answer your questions and improve our paper based on your suggestions. Below, we address your comments in detail.

Response to Weakness

Response to W1

Thank you for the suggestion. We agree that the original phrasing in the abstract may have overstated what is already known. Our intention was to highlight that we formally define this issue under inference-time text confounding, which, to the best of our knowledge, has not been previously studied in the context of CATE estimation. We will revise the abstract to better reflect our contribution.

Response to W2

Thank you for the insightful question. You are right that if the text used at inference time is affected by the treatment and/or outcome (e.g., discharge notes written after given treatment decisions), it may act as a post-treatment variable. Conditioning on such variables can induce post-treatment bias, as they may lie on causal paths from treatment to outcome (e.g., mediators or colliders).

Hence, we agree that, in real-world datasets, not all observed covariates should be treated as confounders—some may indeed be post-treatment variables (e.g., mediators) and should not be used for adjustment. We will thus add a discussion with practical considerations to help practitioners identify post-treatment variables, which should not be used for adjustment.

In our paper, we explicitly define $X$ as the set of pre-treatment variables that serve as confounders, in line with standard assumptions in the CATE estimation literature. Our method assumes that only such pre-treatment variables—whether structured or textual—are used for adjustment.

In practice, it is crucial for practitioners to apply domain knowledge to identify and exclude post-treatment variables from the adjustment set. This is typically feasible in real-world clinical settings, as most records (including textual notes) are timestamped, making it possible to determine whether a variable was recorded before or after treatment assignment.

Response to W3

Thank you for the comment. While it may be intuitively expected that the naïve estimator is biased under inference-time text confounding, the goal of Lemma 4.1 is to formally formalize the pointwise bias.

Lemma 4.1 provides an explicit expression that quantifies how the omission of confounding information (when using $T$ instead of $X$) affects the estimated treatment effect. This formal result is important for rigorously establishing the limitations of naïve estimators that use text representations without proper adjustment.

We provide this intuition in Remark 4.2, which highlights that, in many real-world settings, the text variable $T$ often captures only a subset of the information in $X$, resulting in $Y(a) \not\perp A \mid T$. As a result, the naïve estimator is biased, and this bias is typically positive under common conditions. We will revise the paper to better emphasize that Lemma 4.1 is not just stating that bias exists, but rather providing a precise mathematical form of that bias.

Response to W4

Thank you for the insightful comment. We agree that identifying $\tau^x(x)$ at test time would be even more interesting. However, in our setting, confounder $X$ is not accessible at inference time, which makes identification of $\tau^x(x)$ fundamentally impossible.

We appreciate your suggestion and have incorporated it. We conducted the comparison between estimated $\hat{\tau}^t(t)$ against the ground-truth $\tau^x(x)$ in in-sample data. We report both the averaged bias $\mathbb{E}\left[\hat{\tau}^t(t) - \tau^x(x)\right]$ and the root mean squared error (RMSE) $\sqrt{\mathbb{E}\left[\left(\hat{\tau}^t(t) - \tau^x(x)\right)^2\right]}$ as metrics. The results for both the IST and MIMIC-III datasets across several baselines are below.

This comparison provides empirical evidence that our method produces consistent and accurate approximations of $\tau^x(x)$, with minimal deviation, thereby validating the effectiveness of our approach.

Action: We will add the above experiments to our revised paper.

Thank you for your positive review and your helpful comments! We are very happy to answer your questions and improve our paper as a result.

Response to Questions

Answer to Q1

Thank you for your important question.

We have access to ChatGPT Edu through our university, which is SOC 2 Type II compliant [1]. This gives us enhanced privacy controls not available in the standard API, free version, or even ChatGPT Team.

As part of this setup, we have applied for a Business Associate Agreement (BAA) [2] to gain access to Zero Data Retention (ZDR). With ZDR enabled, the store parameter for /v1/responses and v1/chat/completions will always be treated as false [3]. In this setup, OpenAI does not retain any request or response data after processing—meaning there is no storage, no logging, and, more importantly, no human review.

This configuration provides a privacy approach that is functionally equivalent to the one recommended by PhysioNet [4], where Azure OpenAI is used with human review explicitly disabled.

=> We confirm that our analysis is thus HIPAA-compliant and thus in line with the PhysioNet rules.

[1] https://openai.com/index/introducing-chatgpt-edu/

[2]https://help.openai.com/en/articles/8660679-how-can-i-get-a-business-associate-agreement-baa-with-openai-for-the-api-services

[3] https://platform.openai.com/docs/guides/your-data#zero-data-retention

[4] https://physionet.org/news/post/gpt-responsible-use

Answer to Q2

Thank you for raising this important point.

We acknowledge that LLM-generated text may not perfectly align with real-world text distributions and that a domain shift between training-time and inference-time text can occur in practice. Such shifts can affect methods’ performance.

To evaluate this scenario, we conducted robustness experiments by varying the LLM prompt strategies used to generate text across datasets (Figure 3(b)). This introduces a controlled distribution shift in the text inputs. Despite these differences, our method consistently outperforms all baselines, demonstrating its robustness to moderate levels of text distribution mismatch.

While discrepancies between the surrogate text $\tilde{T}$ and the true inference-time distribution $T$ may exist, they can potentially be mitigated through domain adaptation techniques. For instance, if a small sample of real-world text $T$ is available, one could adapt the generation process of $\tilde{T}$ to better align with it. We consider this a promising direction for future work.

For example, there are several potential e strategies for extending our approach when a few text samples from the test distribution are available: (i) In-context alignment: One strategy is to incorporate these additional samples directly into the prompt, guiding the LLM with instructions such as: "Here are example descriptions. Please formulate your descriptions in a similar style, level of detail, etc." (ii) Alignment via fine-tuning: Anothe strategy is fine-tuneing the LLM by introducing a loss term that minimizes the distributional distance between generated text and the test distribution samples—for example, using Wasserstein distance or Maximum Mean Discrepancy (MMD) as the alignment loss.

Answer to Q3

Thank you for highlighting this important point. We fully agree that if text data is generated after treatment assignment or outcome realization, it may act as a post-treatment variable (e.g., mediators), and adjusting for such variables can introduce collider bias or post-treatment bias.

In our work, the $T$ used is only the set of pre-treatment variables that serve as confounders. In practice, it is crucial for practitioners to apply domain knowledge to identify and exclude post-treatment text from the adjustment set for inference data. This is typically feasible in real-world clinical settings, as most health records are timestamped, making it possible to determine whether a variable was recorded before or after treatment assignment.

Answer to Q4

Thank you for the insightful question. Yes, if clinical text data is available at training time, our method can still be applied.

If the clinical text contains additional confounding information that is not captured in the structured covariates, we can include this text together with the covariates $X$ as an additional input to the text generation (i.e., the prompt could look like “A patient with features $X$ and the following reported self-reported symptoms[...]”) in the estimation procedure to avoid confounding bias. Further, this additional information helps to improve the accuracy of CATE estimation in downstream estimates. This flexibility resembles another strength of our method.

However, if the clinical text contains only confounding information that is also captured in $X$ or is less informative, but follows a different distribution compared to the inference-time text (e.g., different style or domain), then this text contains only noisy information that does not improve estimation. Thus, such text should not be considered, and we can apply our framework as is, using only the covariates $X$, which thus still requires generating surrogate text data.

Answer to Q5, 6 Thank you! We will fix the typos.

Thank you for your constructive review and your helpful comments! We are very happy to answer your questions and improve our paper as a result.

Response to weakness

Response to W1

Thank you for the comment. We would like to clarify that TCA is not a straightforward application of standard meta-learners or doubly robust (DR) learners.

Our contributions go beyond existing methodologies in the following ways:

(1) Problem formalization: We are the first to formally define the inference-time text confounding problem and show that naïve applications of existing meta-learners are biased. We derive pointwise confounding bias for these baselines, highlighting their limitations.

(2) Theoretical insight: We provide the identifiability result connecting the target CATE under text ($\tau^t(t)$) to the CATE under structured confounders($\tau^x(x)$), and propose a two-stage estimation procedure to adjust for inference-time confounding. This theoretical connection is non-trivial and not addressed in prior work.

(3) Practical framework: We introduce an LLM-driven pipeline to generate text and translate the theoretical result to a novel, practical framework for dealing with real-world conditions where structured data is unavailable at inference time.

In contrast, the baselines we include in our experiments are direct and straightforward applications of existing meta-learners using text as input. Even though they share the same initial text-generation step (using LLMs), they remain biased, as demonstrated both theoretically (Lemma 4.1) and empirically in our experiments (Table 1). ] A simple off-the-shelf combination of the text-generation step with existing DR-learners would be biased.

To show this empirically, we performed new experiments against an off-the-shelf combination of our text-generation step and several more meta-learner baselines, including DR-learner, X-learner, and R-learner. We call these baselines “TBE-DR”, “TBE-X”, “TBE-R”. Importantly, we made the experiment fair: we used the same setup for the LLMs, etc., so that the performance gains must be attributed to our framework. We again report PEHE scores on two datasets.

=> Our TCA shows consistent improvements by a large margin.

Our method outperforms the TBE-DR baselines despite using the same underlying estimation technique—highlighting that the core contribution lies in how we adjust for confounding, not in the DR method itself. Our method also outperforms other straightforward applications of existing meta-learner baselines. This demonstrates the effectiveness of our TCA framework and proves our method differs substantially from existing methodologies (meta-learners and DR pseudo-outcomes).

Action: We will add the new experiments to our paper.

Response to W2

Thank you for the question. In Figure 1, we illustrate a common setup from clinical datasets such as the IST dataset. The intervention in this example is a binary treatment: whether the patient receives the treatment assignment of aspirin allocation or not. The outcome is whether the patient experiences a recurrent stroke. Patient covariates, which are only observed directly during training time, include demographic variables (e.g., gender, age) and clinical measurements (e.g., presence of atrial fibrillation, blood pressure, weakness in arms/legs). In this context, the CATE is the expected difference in the probability of having a stroke recurrence if a patient with covariates $X=x$ were given aspirin versus or not.

=> The motivation is that such CATE estimates help inform personalized treatment decisions: for which subgroups of patients is certain medicine beneficial, and for which it may not help much.

The challenge we focus on in the paper is that, at inference time, when a new patient comes to ask about his/her case via telemedicine or remote healthcare consultations, or medical chatbots), the structured features may not be available (e.g., due to the limitation of diagnostic tests or clinical measurements, sensor limitations, or privacy constraints), but this patient can describe his/her feeling. CATE predictions are made using only textual descriptions with self-reported symptoms of the patient.

Our framework is designed for this setup: to address this inference-time text confounding setting by learning to estimate CATE from structured data at training time, and using text at inference time, while adjusting for confounding via the induced confounding information in the text.

Response to W3

Thank you for the suggestion. We would like to clarify that SIN [1] and CaML [2] focus on a different setting and that both are not meaningful baselines from a theoretical and empirical point.

The baselines (SIN and CaML) focus on a different causal graph: these methods focus with high-dimensional treatments (e.g., predicting treatment effects of a new drug by using the molecular structure or a text description), while our settings focus on high-dimensional confounders (e.g., the text descriptions is what we use to adjust for confounding). Hence, the way of how high-dimensional information enters the machine pipeline is also different. This fundamental difference leads to different causal graphs and objectives.

There are also some further differences: (i) Our method does not require additional information on complex treatments. (ii) We focus on mitigating inference-time confounding, whereas SIN and CaML are developed for estimating treatment effects of complex structured treatments (text or molecules). (iii) We additionally assume that text is only available at inference time, not at training time (while, for SIN and CaML, it is mandatory that text is available at training time). Especially because of reason (iii), both baselines are not directly applicable to our setting.

In addition to the theoretical explanation, we would like to provide empirical evidence. Although these methods are not directly applicable to our setting, we adapted their architectures for comparison. This means we use their architectures but integrate them into our LLM pipeline (e.g., we combine their architecture with the text generation step from our TCA framework). Further, for SIN[1], we use their method design to construct pseudo-outcomes and then regress on the text generated by LLMs to simulate missing treatment information. For CaML[2], as it is a zero-shot approach that requires additional structured information for a new treatment as input, we cannot apply it in its original form. Instead, we take its model architecture to predict CATE. The results for both IST and MIMIC-III datasets are below (reported: PEHE). The results show that our method outperforms the baselines.

Action: We will add the above comparison to our revised paper.

Reference:

[1] Kaddour, J., Zhu, Y., Liu, Q., Kusner, M.J. and Silva, R. Causal effect inference for structured treatments. Advances in Neural Information Processing Systems, 2021.

[2] Nilforoshan, H., Moor, M., Roohani, Y., Chen, Y., Šurina, A., Yasunaga, M., Oblak, S. and Leskovec, J.. Zero-shot causal learning. Advances in Neural Information Processing Systems, 2023.