LLM-Driven Treatment Effect Estimation Under Inference Time Text Confounding

2025-07-03

AI Summary

This paper addresses the “inference-time textual confounding” problem in personalized medicine: while training leverages structured electronic health records (EHRs) containing complete confounders, inference relies solely on incomplete, self-reported patient text—inducing estimation bias. We formally define this challenge and propose LLM-DR, a novel framework integrating large language models (LLMs) with a customized doubly robust (DR) estimator. LLM-DR employs LLMs to implicitly recover confounding information from unstructured text and applies DR estimation to jointly correct for residual bias and model misspecification. Experiments across multiple real-world clinical datasets demonstrate that LLM-DR significantly improves both accuracy and robustness of treatment effect estimation. Our approach establishes a new, interpretable, and generalizable causal modeling paradigm for text-driven personalized decision-making.

Technology Category

• Machine Learning: Causal Learning
• Natural Language Processing: (Large) Language Models
• Reasoning under Uncertainty: Causality

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions

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 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 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.

Problem

• Estimating treatment effects with incomplete inference-time text data
• Addressing bias from training-inference data discrepancy
• Leveraging LLMs to mitigate text confounding biases

Innovation

• LLM-driven framework for treatment effect estimation
• Custom doubly robust learner mitigates biases
• Addresses inference time text confounding issue