CATR: Confounding-Aware Token Rationalization

• The paper introduces a novel framework that uses HSIC-based diagnostics to select a sparse subset of tokens from high-dimensional text, ensuring unconfoundedness.
• It employs a relaxed-Bernoulli mask and a multi-term loss to balance predictive utility, sparsity, and preservation of confounding signals.
• Empirical evaluations demonstrate that CATR reduces bias and variance while increasing effective sample sizes and providing stable treatment effect estimates.

Confounding-Aware Token Rationalization (CATR) is a methodological framework designed to address the challenges posed by high-dimensional text covariates in causal effect estimation from observational data. By selectively identifying a sparse subset of tokens carrying the requisite confounding signal, CATR mitigates positivity violations and instability in inverse-probability weighted estimators that often arise when full text representations are used. CATR employs a residual-independence diagnostic based on the Hilbert–Schmidt Independence Criterion (HSIC) to ensure the selected token set suffices for unconfoundedness, optimizing a multi-term loss that balances predictive utility, parsimony, and preservation of confounding structure (Zhang et al., 5 Dec 2025).

1. Problem Setting and Motivation

Consider an observational paper where each unit $i$ is characterized by:

• A high-dimensional text sequence $W_i = (w_{i1}, \ldots, w_{iL_i})$,
• A binary treatment $T_i \in \{0,1\}$,
• An outcome $Y_i \in \{0,1\}$ or $\mathbb{R}$.

Under the potential-outcomes framework, each unit possesses counterfactuals $Y_i(0), Y_i(1)$. The scientific estimand is the average treatment effect (ATE):

$\tau = \mathbb{E}[Y(1) - Y(0)].$

Identification requires:

• Unconfoundedness: $\{Y(0), Y(1)\} \perp T \mid W^c$,
• Positivity: $$0 < p_{\min} \leq \Pr(T=1 \mid W^c = w^c) \leq p_{\max} < 1$$ for all $w^c$,
• Consistency: $Y = Y(T)$.

Only a typically unknown subset $W^c \subset W$ contains the true confounding information, while redundant $W^r = W \setminus W^c$ tokens are unrelated. Conditioning on the entire $W$ is standard practice but, in high-dimensional settings, leads to "observational-level positivity violations", where the estimated propensity $\widehat{g}(w)$ falls outside the desired $[\epsilon, 1-\epsilon]$ interval even if $g(W^c)$ satisfies overlap.

Positivity violations manifest as extreme propensity scores, large inverse-probability weights, high estimator variance, and inflated frequency of propensity score clipping. Toy examples (see Figure 1 in (Zhang et al., 5 Dec 2025)) demonstrate the undesirable concentration of $\widehat{g}(X)$ at the boundaries and an increased clipping fraction as proxy dimension rises.

2. Core Methodology: CATR Framework

2.1 Token Selection and Predictive Model

CATR employs a selector network $S_\theta$ assigning each token $w_{ij}$ a score $a_{ij} \in [0,1]$. A relaxed-Bernoulli distribution (e.g., Gumbel-Softmax) provides a continuous mask $s_{ij} \in [0,1]$, forming the rationalized subsequence $\tilde w_i = s_i \odot w_i$.

A shared predictor network $h_\phi$ receives $\tilde w_i$ and outputs:

• $$\widehat{g}(\tilde w_i) \approx \Pr(T=1 \mid \tilde w_i)$$,
• $$\widehat{Q}_t(\tilde w_i) \approx \mathbb{E}[Y \mid T=t, \tilde w_i]$$.

2.2 Residual-Independence Diagnostic

To verify sufficiency, CATR introduces a residual-independence diagnostic. For candidate subset $S$, define batch residuals:

$$r_{T} = T - \widehat{g}(\tilde W), \qquad r_{Y} = Y - \widehat{Q}_{T}(\tilde W).$$

The empirical Hilbert–Schmidt Independence Criterion (HSIC):

$$\widehat{\mathrm{HSIC}}(r_T, r_Y) = \frac{1}{(n-1)^2} \operatorname{tr}(KHKH),$$

where $K_{ij} = k_T(r_{T,i}, r_{T,j})$, $L_{ij} = k_Y(r_{Y,i}, r_{Y,j})$, and $H = I_n - \tfrac{1}{n}\mathbf{1}\mathbf{1}^\top$.

Proposition 1 demonstrates:

$$(Y(0), Y(1)) \perp T \mid \tilde W \implies \mathrm{HSIC}(r_T, r_Y) = 0.$$

Therefore, nonzero observed HSIC indicates the selection is insufficient to block confounding paths.

2.3 Optimization Objective

The full optimization problem,

$$\min_{\theta, \phi} \frac{1}{n}\sum_{i=1}^n \left[ \mathcal{L}_{\mathrm{sup},i} + \mu \mathcal{R}_{\mathrm{sparse}}(s_i) + \gamma \widehat{\mathrm{HSIC}}(r_{T,i}, r_{Y,i}) \right],$$

comprises:

• $\mathcal{L}_{\mathrm{sup}}$: Sum of cross-entropy losses for $\widehat{g}$ and $\widehat{Q}_t$,
• $\mathcal{R}_{\mathrm{sparse}}$: Sparsity penalty (entropy/KL),
• $\widehat{\mathrm{HSIC}}$: Residual-independence metric.

Hyperparameters $\mu, \gamma$ control sparsity and independence trade-offs.

3. Algorithmic Implementation

Algorithm 1 in (Zhang et al., 5 Dec 2025) proceeds via stochastic gradient descent in minibatches of size $m$:

1. Sample minibatch $(w_i, t_i, y_i)$.
2. Compute selection scores $a_i = S_\theta(w_i)$.
3. Sample relaxed mask $s_i \sim \mathrm{RelaxedBernoulli}(\eta, a_i)$.
4. Form rationalized text $\tilde w_i = s_i \odot w_i$.
5. Predict $\widehat{g}(\tilde w_i)$ and $\widehat{Q}_t(\tilde w_i)$.
6. Compute residuals $r_{T,i}, r_{Y,i}$ and batch HSIC.
7. Accumulate losses and backpropagate.

Primary hyperparameters:

• Relaxed-Bernoulli temperature $\eta$,
• Token budget prior $\pi$ (if KL penalty used),
• Sparsity and independence weights $\mu$, $\gamma$.

Batchwise HSIC is $\mathcal{O}(m^2)$ per batch but tractable for $m \in [32, 128]$. Computational complexity is dominated by HSIC and the encoder.

4. Theoretical Properties

Theoretical guarantees under regularity assumptions (bounded embeddings, Sobolev smoothness, bounded penalties) include:

• Fast nonasymptotic rates: For $f \in \{g, Q_t\}$,

$$\|\widehat{f} - f_*\|_{L_2(P)}^2 = O_p\left(n^{-\frac{\beta}{\beta + d}} \log^8 n\right).$$

• Product-rate condition: Nuisance estimators $\widehat{g}$, $\widehat{Q}_t$ permit doubly-robust inference.
• Consistency and efficiency: IPW and AIPW are consistent under empirical overlap; AIPW is $\sqrt{n}$-consistent and asymptotically normal if propensities are bounded away from 0 and 1:

$$\sqrt{n}(\widehat{\tau}_{\rm AIPW} - \tau) \rightsquigarrow \mathcal{N}(0, \operatorname{Var}(\psi_1 - \psi_0)).$$

5. Empirical Evaluation

5.1 Semi-Synthetic MIMIC-III Experiment

• Design: Physician notes with confounder signal detected via infection-keyword indicators; nonlinear transformations define $g(W)$ and $Q_T(W)$.
• Method comparison: TARNet, CFRNet, DragonNet, CausalBERT, CATR.
• Estimators: OR, IPW, AIPW.
• Metrics: Absolute bias, empirical SD, bootstrap SE, CI coverage, effective sample size (ESS) ratio, clipping fraction.

Results: CATR achieves the lowest bias and variance, highest CI coverage, better ESS ratio (≈0.31 vs. 0.27 baseline), and fewer extreme propensities (≈10.7% vs. 29–31%). Ablation shows removing HSIC or sparsity deteriorates all performance measures and increases clipping.

Qualitative findings: Without HSIC regularization, top tokens selected are largely spurious. With HSIC, all salient infection confounders are recovered.

5.2 Real-World MIMIC-III Study

• Cohort: Septic ICU patients,
• Treatment: IV fluid bolus,
• Outcome: ICU readmission,
• Covariates: 42 structured plus unstructured text,
• Adjustments: structured-only, multimodal, multimodal+CATR.

Results: Structured-only yields low ESS (0.21) and 8% clipping. Adding text (no CATR) slightly improves ESS (0.28) but with 20% clipping and unstable ATE. Multimodal+CATR achieves ESS ≈0.85, 0% clipped propensities, and stable ATE (≈–0.026, SE ≈0.041). Token selections highlight clinically meaningful terms such as “sepsis” and “infection”.

6. Limitations and Future Directions

Key limitations:

• HSIC provides a soft test; zero HSIC is necessary but not sufficient for unconfoundedness.
• Mini-batch HSIC estimates are noisy with small batch sizes.
• Hyperparameter selection (especially $\mu, \gamma$, token allowance) demands careful validation and may be computationally intensive.
• Token selection quality is contingent on the information encoded by the fixed pretrained embedding.

Ongoing and future research directions include:

• Integrating counterfactual necessity and sufficiency diagnostics.
• Combining CATR with post-hoc calibration or covariate balancing to stabilize IPW.
• Extending token selection to multimodal (text and structured) covariates.
• Exploring alternative or differentiable independence criteria.

7. Context and Significance

CATR addresses the unique methodological problem posed by high-dimensional text when used as covariates in causal inference. By selecting a minimal sufficient subset of tokens, CATR avoids the fragility of propensity estimation in large text spaces while preserving requisite confounding information relevant for effect identification. Empirically, CATR demonstrates improved estimator stability, reduced bias/variance, enhanced effective sample size, and improved interpretability over baseline and ablated approaches. The framework thus advances the robust integration of unstructured text into causal effect estimation pipelines, especially in settings where naive adjustment for entire documents would undermine positivity and statistical efficiency.

(Zhang et al., 5 Dec 2025)