Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Text Rationalization for Robust Causal Effect Estimation

Lijinghua Zhang, Hengrui Cai

TL;DR

This work tackles the challenge of estimating causal effects from text data by addressing positivity violations that arise when high-dimensional textual confounders inflate the covariate space. It introduces Confounding-Aware Token Rationalization (CATR), which learns a sparse subset of tokens via a differentiable selector and HSIC-based residual-dependence penalties to preserve confounding information while improving overlap. The authors provide nonasymptotic error bounds for neural nuisance estimators and establish sqrt(n)-consistency and asymptotic normality for the AIPW estimator under the framework, and validate the approach on semi-synthetic data and a real MIMIC-III sepsis cohort, showing improved stability, overlap, and interpretability. The method supports multimodal extensions to include structured covariates and yields interpretable token selections that highlight clinically relevant confounders. Overall, CATR offers a robust and interpretable toolkit for text-based causal adjustment with practical impact in healthcare analytics and beyond.

Abstract

Recent advances in natural language processing have enabled the increasing use of text data in causal inference, particularly for adjusting confounding factors in treatment effect estimation. Although high-dimensional text can encode rich contextual information, it also poses unique challenges for causal identification and estimation. In particular, the positivity assumption, which requires sufficient treatment overlap across confounder values, is often violated at the observational level, when massive text is represented in feature spaces. Redundant or spurious textual features inflate dimensionality, producing extreme propensity scores, unstable weights, and inflated variance in effect estimates. We address these challenges with Confounding-Aware Token Rationalization (CATR), a framework that selects a sparse necessary subset of tokens using a residual-independence diagnostic designed to preserve confounding information sufficient for unconfoundedness. By discarding irrelevant texts while retaining key signals, CATR mitigates observational-level positivity violations and stabilizes downstream causal effect estimators. Experiments on synthetic data and a real-world study using the MIMIC-III database demonstrate that CATR yields more accurate, stable, and interpretable causal effect estimates than existing baselines.

Text Rationalization for Robust Causal Effect Estimation

TL;DR

This work tackles the challenge of estimating causal effects from text data by addressing positivity violations that arise when high-dimensional textual confounders inflate the covariate space. It introduces Confounding-Aware Token Rationalization (CATR), which learns a sparse subset of tokens via a differentiable selector and HSIC-based residual-dependence penalties to preserve confounding information while improving overlap. The authors provide nonasymptotic error bounds for neural nuisance estimators and establish sqrt(n)-consistency and asymptotic normality for the AIPW estimator under the framework, and validate the approach on semi-synthetic data and a real MIMIC-III sepsis cohort, showing improved stability, overlap, and interpretability. The method supports multimodal extensions to include structured covariates and yields interpretable token selections that highlight clinically relevant confounders. Overall, CATR offers a robust and interpretable toolkit for text-based causal adjustment with practical impact in healthcare analytics and beyond.

Abstract

Recent advances in natural language processing have enabled the increasing use of text data in causal inference, particularly for adjusting confounding factors in treatment effect estimation. Although high-dimensional text can encode rich contextual information, it also poses unique challenges for causal identification and estimation. In particular, the positivity assumption, which requires sufficient treatment overlap across confounder values, is often violated at the observational level, when massive text is represented in feature spaces. Redundant or spurious textual features inflate dimensionality, producing extreme propensity scores, unstable weights, and inflated variance in effect estimates. We address these challenges with Confounding-Aware Token Rationalization (CATR), a framework that selects a sparse necessary subset of tokens using a residual-independence diagnostic designed to preserve confounding information sufficient for unconfoundedness. By discarding irrelevant texts while retaining key signals, CATR mitigates observational-level positivity violations and stabilizes downstream causal effect estimators. Experiments on synthetic data and a real-world study using the MIMIC-III database demonstrate that CATR yields more accurate, stable, and interpretable causal effect estimates than existing baselines.

Paper Structure

This paper contains 16 sections, 4 theorems, 16 equations, 5 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Contributions
  4. Notation and Problem Formulation
  5. HSIC Measurement for Subset Sufficiency
  6. Methodology
  7. Learning Architecture
  8. Objective Function
  9. Extension: Multimodal Data
  10. Theoretical Analysis
  11. Nonasymptotic Error Bound of Neural Network Approximation
  12. Statistical Inference
  13. Empirical Experiments
  14. Synthetic Experiment
  15. Real Data Analysis
  16. ...and 1 more sections

Key Result

Proposition 3.1

Let $\tilde{\boldsymbol{W}}\subseteq\boldsymbol{W}$ be any candidate subset. Define $g(\tilde{\boldsymbol{W}}) = \mathbb{E}[T\mid\tilde{\boldsymbol{W}}]$ and $Q_t(\tilde{\boldsymbol{W}}) = \mathbb{E}[Y\mid T=t,\,\tilde{\boldsymbol{W}}]$, as well as residuals $r_T = T - g(\tilde{\boldsymbol{W}})$ and Contrapositively, if the residual HSIC is positive, then $\tilde{\boldsymbol{W}}$ is not sufficient

Figures (5)

  • Figure 1: Illustration of the observational-level positivity violation induced by high-dimensional proxies. Fig A shows reasonable overlap of confounder between treated and control groups in the latent confounder space. Fig B shows stronger separation on some direction after mapping to high-dimensional sparse proxies $X$, as seen from projections onto the first two principal components with marginal KDEs. Fig C shows the estimated propensities concentrates more mass near the extremes (vertical dashed lines indicate clipping thresholds at 0.03 and 0.97), highlighting inflated variance in weighting. Fig D shows that as proxy dimension increases, the clipping fraction rises rapidly, demonstrating how redundant proxies dilute effective overlap and exacerbate observational-level positivity violations.
  • Figure 2: Causal diagram illustrating text-based confounding: $\boldsymbol{W}^c$ (confounding tokens) affects both treatment $T$ and outcome $Y$, while $\boldsymbol{W}^r$ denotes irrelevant or redundant tokens.
  • Figure 3: The Confounding-Aware Token Rationalization (CATR) framework (with multimodal extension). The main pipeline (solid blocks) implements selective rationalization for text-only inputs. The semi-transparent dashed components illustrate the extension to tabular covariates, showing that our approach naturally adapts to multi-modal settings.
  • Figure 4: Comparison of importance score distributions with and without HSIC regularization.
  • Figure 5: Fig. A shows top-$k$ selected tokens that are medically informative. Fig. B visualizes the selective rationalization output with color-encoded selection importance scores.

Theorems & Definitions (7)

  • Definition 2.1: Observational-level positivity violation
  • Proposition 3.1
  • Remark
  • Remark
  • Theorem 5.1
  • Theorem 5.2
  • Theorem 5.3