# Uncertainty-Guided LLM Semantic Augmentation for Heterogeneous Treatment Effect Estimation

arXiv: https://arxiv.org/abs/2607.26599 · HTML: https://mlsanity.com/paper/2607.26599/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 27 of its 58 references that are in the map.
- First posted on arXiv 29 Jul 2026.

## Builds on (most cited first, 27 of 27)

- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [Weight Uncertainty in Neural Networks](https://mlsanity.com/paper/1505.05424/) (2015)
- [Variational Autoencoders and Nonlinear ICA: A Unifying Framework](https://mlsanity.com/paper/1907.04809/) (2019)
- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Causal Effect Inference with Deep Latent-Variable Models](https://mlsanity.com/paper/1705.08821/) (2017)
- [Causal Reasoning and Large Language Models: Opening a New Frontier for Causality](https://mlsanity.com/paper/2305.00050/) (2023)
- [Quasi-Oracle Estimation of Heterogeneous Treatment Effects](https://mlsanity.com/paper/1712.04912/) (2017)
- [Uncertainty-Based Offline Reinforcement Learning with Diversified Q-Ensemble](https://mlsanity.com/paper/2110.01548/) (2021)
- [Adapting Neural Networks for the Estimation of Treatment Effects](https://mlsanity.com/paper/1906.02120/) (2019)
- [Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes](https://mlsanity.com/paper/1704.02801/) (2017)
- [Uncertainty Weighted Actor-Critic for Offline Reinforcement Learning](https://mlsanity.com/paper/2105.08140/) (2021)
- [Pessimistic Bootstrapping for Uncertainty-Driven Offline Reinforcement Learning](https://mlsanity.com/paper/2202.11566/) (2022)
- [Neural Contextual Bandits with Deep Representation and Shallow Exploration](https://mlsanity.com/paper/2012.01780/) (2020)
- [Treatment effect estimation with disentangled latent factors](https://mlsanity.com/paper/2001.10652/) (2020)
- [Identifying Causal-Effect Inference Failure with Uncertainty-Aware Models](https://mlsanity.com/paper/2007.00163/) (2020)
- [A Critical Look at the Consistency of Causal Estimation With Deep Latent Variable Models](https://mlsanity.com/paper/2102.06648/) (2021)
- [End-To-End Causal Effect Estimation from Unstructured Natural Language Data](https://mlsanity.com/paper/2407.07018/) (2024)
- [Causal Inference with Conditional Front-Door Adjustment and Identifiable Variational Autoencoder](https://mlsanity.com/paper/2310.01937/) (2023)
- [DESCN: Deep Entire Space Cross Networks for Individual Treatment Effect Estimation](https://mlsanity.com/paper/2207.09920/) (2022)
- [β-Intact-VAE: Identifying and Estimating Causal Effects under Limited Overlap](https://mlsanity.com/paper/2110.05225/) (2021)
- [Proximal Causal Inference With Text Data](https://mlsanity.com/paper/2401.06687/) (2024)
- [Self-Distilled Disentangled Learning for Counterfactual Prediction](https://mlsanity.com/paper/2406.05855/) (2024)
- [Progressive Generalization Risk Reduction for Data-Efficient Causal Effect Estimation](https://mlsanity.com/paper/2411.11256/) (2024)
- [Improving Data-driven Heterogeneous Treatment Effect Estimation Under Structure Uncertainty](https://mlsanity.com/paper/2206.12689/) (2022)
- [Deep Disentangled Representation Network for Treatment Effect Estimation](https://mlsanity.com/paper/2507.06650/) (2025)

## Abstract (arXiv)

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.