Uncertainty-Guided LLM Semantic Augmentation for Heterogeneous Treatment Effect Estimation

2026/07/29 by Jialu Xu, Mengkun Liang, Guannan Liu +2
Computer Science · #cs.LG

paper · pdf

17 pages, 12 figures, 5 tables

arxiv created 2026/07/29 · arxiv updated 2026/07/30

Abstract

Estimating heterogeneous treatment effects is central to targeted interventions, such as personalized promotions and precision medicine. We focus on the conditional average treatment effect (CATE), a standard estimand for characterizing such heterogeneity. Even under standard identification conditions, finite-sample CATE estimation requires learning the nuisance structure for covariate adjustment and treatment-effect heterogeneity, often together with an effective representation of X. Raw numerical and categorical encodings can leave semantic relations and higher-order interactions implicit, making this joint task locally unstable. A motivating study further shows that this instability appears through partially separable assignment- and heterogeneity-side channels. Building on this observation, we propose CURL (Causal Uncertainty-guided Representation Learning), a plug-in adapter that uses estimator uncertainty to allocate pretrained semantic capacity to locally unstable units. CURL queries a frozen LLM through two role-conditioned prompts, constructs assignment- and heterogeneity-oriented representations from the observed covariates, and routes them through separated pathways. On four benchmarks, CURL improves ten host learners in most settings, while ablation, refinement-dynamics, route-reassignment, and probe analyses support the intended design and roles of the two channels.

Citations

• Deep Disentangled Representation Network for Treatment Effect Estimation
• Progressive Generalization Risk Reduction for Data-Efficient Causal Effect Estimation
• End-To-End Causal Effect Estimation from Unstructured Natural Language Data
• Self-Distilled Disentangled Learning for Counterfactual Prediction
• Proximal Causal Inference With Text Data
• Causal Inference with Conditional Front-Door Adjustment and Identifiable Variational Autoencoder
• Causal Reasoning and Large Language Models: Opening a New Frontier for Causality
• Improving Data-driven Heterogeneous Treatment Effect Estimation Under Structure Uncertainty
• Pessimistic Bootstrapping for Uncertainty-Driven Offline Reinforcement Learning
• Uncertainty-Based Offline Reinforcement Learning with Diversified Q-Ensemble
• Uncertainty Weighted Actor-Critic for Offline Reinforcement Learning
• Neural Contextual Bandits with Deep Representation and Shallow Exploration
• Identifying Causal-Effect Inference Failure with Uncertainty-Aware Models
• Treatment effect estimation with disentangled latent factors
• Adapting Neural Networks for the Estimation of Treatment Effects
• Quasi-Oracle Estimation of Heterogeneous Treatment Effects
• Quasi-oracle estimation of heterogeneous treatment effects
• Causal Effect Inference with Deep Latent-Variable Models
• Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes
• What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?
• Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles
• Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning
• Weight Uncertainty in Neural Networks
• A Practical Bayesian Framework for Backpropagation Networks
• The central role of the propensity score in observational studies for causal effects
• Large Language Models and Causal Inference in Collaboration: A Survey
• β-Intact-VAE: Identifying and Estimating Causal Effects under Limited Overlap

Related

• Transfer Learning on Heterogeneous Feature Spaces for Treatment Effects Estimation
• Improving Data-driven Heterogeneous Treatment Effect Estimation Under Structure Uncertainty
• Variable importance measures for heterogeneous treatment effects
• Causal Rule Forest: Toward Interpretable and Precise Treatment Effect Estimation
• Statistical Learning for Heterogeneous Treatment Effects: Pretraining, Prognosis, and Prediction