PaperlayerView exact revision on arXiv

Source-linked AI summary

Exploring Sparse Autoencoders in Text-Based Causal Confounding Adjustment

Mian Zhong, Katherine A. Keith, Anjalie Field

arXiv:2609.01322v1cs.CLcs.LG

TL;DR

Text-based causal adjustment must preserve confounding information without creating high-dimensional, high-variance adjustment problems. The paper proposes an SAE pipeline that selects minimal treatment-related features through conditional independence tests. SAE approaches perform best in the reported binary and multi-class 20NG evaluations, while multi-label EURLEX results are mixed and expose common-support and convergence concerns.

  • Problem

    Text representations for observational causal adjustment must retain sufficient confounding information while remaining sparse enough for finite-sample overlap and low-variance estimation.

  • Method

    The pipeline trains sparse autoencoders on text embeddings, selects a minimal treatment-informative feature subset using conditional independence tests, and applies CEM or DoubleML.

  • Results

    SAE-based approaches achieve the lowest bias, lowest RMSE, and highest coverage across the binary and multi-class 20NG settings, while multi-label EURLEX results are more mixed.

  • Takeaways & Limitations

    Interpretable SAE features provide diagnostics supporting falsification, while the mixed multi-label results indicate that more complex text-confounding settings warrant caution and further investigation.

  • Takeaways & Limitations

    In the multi-label EURLEX setting, matching retains less than 1% of treated units, while DoubleML can suffer from outcome-estimator nonconvergence and lack of common support.

Abstract

from arXiv · show

In many settings, studying causal questions based on text data requires adjusting for confounding information within texts. Yet there is a tradeoff in constructing text representations for adjustment: they must be sufficiently large and/or dense to preserve the confounding variables necessary for unbiased effect estimation, but sufficiently small and/or sparse to satisfy finite-sample overlap and yield low-variance estimates. To address this tradeoff, we turn to sparse autoencoders (SAEs), and propose a novel causal adjustment pipeline that iteratively selects a minimal set of SAE features via conditional independence tests. We find that SAE representations achieve better adjustments (lower bias and and higher coverage) than alternative representations in standard semi-synthetic evaluations with binary confounders, and their interpretability offers opportunities for falsification. We also introduce a more realistic semi-synthetic evaluation that uses multi-label data as the unobserved confounders and find off-the-shelf adjustment methods require increased investigation for these more complex settings. Code: https://github.com/mianzg/sae-text-confounder

1 Introduction

Text can contain confounding information that must be adjusted for in observational causal studies, creating a tradeoff between representation sufficiency and finite-sample tractability. The paper proposes an SAE-based pipeline that selects minimal treatment-relevant features and reports stronger standard-evaluation performance, while finding mixed results in a more realistic multi-label setting.

  • Motivation: Text often encodes confounders that influence both treatment and outcome, so observational analyses require adjustment for text-based confounding.For example, news topic may affect both device choice and willingness to share.
  • Motivation: High-dimensional text representations can reduce interpretability, while richer or denser adjustment sets weaken overlap and increase estimator variance.This creates a practical mismatch between theoretically attractive estimators and applied requirements for interpretable adjustment.
  • Approach: The proposed pipeline trains an SAE, uses Lasso logistic regression and conditional independence tests to select a minimal treatment-related feature subset, then applies downstream causal estimators.The selected representations are intended to support low-variance estimation and interpretable diagnostics for falsification.
  • Results: In binary-confounder evaluations, SAE representations recover simulated treatment effects more accurately than alternative representations.The evaluation uses semi-synthetic settings in which the true effect is available for comparison.
  • Results: The new multi-label oracle-confounding evaluation produces mixed results, so the paper urges caution with black-box adjustment methods in more realistic text-as-confounding settings.This evaluation is presented as moving closer to realistic settings than standard binary-confounder experiments.

2 Related Work

Prior text-causal work uses classifiers, embeddings, topic models, low-rank language-model training, and proximal inference, but interpretability remains a recurring motivation. The paper builds on this literature by applying interpretable sparse autoencoder representations to confounding adjustment.

  • Applications: Text-based causal inference spans domains including linguistics, social sciences, and public health.
  • Causal adjustment with text: Existing adjustment approaches include propensity matching or regression with extracted text features, alongside text-specific classifiers, embeddings with DoubleML, low-rank language-model training, and zero-shot LLM methods.
  • Interpretable representations: Roberts et al. use topic-model representations for matching because human-readable representations facilitate oversight of text-based analyses.The paper follows this interpretability motivation while using SAE features.
  • Sparse autoencoders: Sparse autoencoders compress data into hidden representations while enforcing sparsity, such as by restricting K-Sparse models to at most k nonzero dimensions.SAEs have become prominent for interpreting large language models.

3 Preliminaries on Confounding Adjustment

The paper studies causal adjustment when text encodes unobserved confounders and must be represented for adjustment. Its central preliminary tension is that representations must preserve confounding information while remaining sparse enough for overlap and manageable variance.

  • Problem setup: The setting assumes texts D encode unobserved confounders U affecting treatment T and outcome Y, with a representation X constructed from the texts.
  • Identification assumptions: Causal identification relies on no unmeasured confounding, overlap, consistency, homogeneous effects, pre-treatment text, and causally sufficient representations.These assumptions cannot be verified from observational data alone.
  • Sufficiency–sparsity tradeoff: Causal sufficiency favors richer representations, whereas increasing representation dimensionality or density makes strict overlap harder and increases estimate variance.If text fails to carry confounding information, no representation can make adjustment sufficient.
  • Downstream estimators: The downstream estimators compared in the paper are DoubleML and matching, with matching seeking balanced treatment and control covariate distributions in representation space.DoubleML combines machine-learning models for treatment and outcome processes but may obscure finite-data common-support problems.

4 Pipeline: Causal Adjustment with SAEs

The pipeline converts text embeddings into sparse SAE representations, selects a minimal subset informative about treatment through Lasso and conditional-independence testing, and plugs that subset into CEM or DoubleML. It also supplies diagnostics for representation sufficiency and qualitative face-validity checks.

  • Step 1: Infer text representations: The first step encodes texts with pretrained embeddings and trains a top-K SAE to obtain sparse representations.The top-K construction keeps K active neurons out of M neurons.
  • Step 2: Select features: The pipeline seeks a minimal adjustment set because fewer variables can lower estimate variance and make strict overlap easier to satisfy in finite samples.It selects SAE features most informative about treatment rather than adjusting on the full representation.
  • Step 3: Test sufficiency: A minimal sufficient subset makes treatment conditionally independent of the remaining SAE features, defining the target of feature selection.
  • Step 3: Test sufficiency: Likelihood-ratio tests compare logistic treatment models using the selected subset and the full SAE set, iterating until the test no longer rejects sufficiency.The final subset may still equal the full SAE feature set.
  • Step 4: Causal estimation: The selected features are used with CEM or DoubleML to estimate the ATT.CEM retains bins containing both treatment and control units, while DoubleML estimates propensity and outcome nuisance functions.
  • Falsification diagnostics: If selection retains every SAE dimension, the diagnostic suggests the SAE dimension M may be too small, although a smaller selected subset does not establish sufficiency.
  • Falsification diagnostics: Interpreting active SAE dimensions lets researchers inspect whether the text aspects identified as confounders have face validity.

5 Semi-Synthetic Experiments

The experiments use semi-synthetic data-generating processes with known treatment effects, progressing from binary and multiclass confounders to a new multi-label confounder setting. They compare SAE-based adjustment with TIRM and embedding-based DoubleML using treatment-effect, overlap, and balance metrics.

  • Data-generating processes: Semi-synthetic simulations combine real text and confounders with synthetically generated treatment and outcome variables, while hiding the true confounders from adjustment methods.The true treatment effect is known, enabling evaluation of estimated effects.
  • Data: 20NewsGroups supplies binary and multiclass confounders based on merged newsgroup labels, including computer, religion, and other.The binary setting uses computer and religion; the multiclass setting adds other as a third class.
  • Data: EURLEX provides a multi-label confounding evaluation using the five most prevalent selected labels and 62,007 samples.Labels outside the selected five are removed, producing the paper’s first semi-synthetic simulation built on multi-labeled confounders.
  • Baseline Pipelines: The comparison includes SAE representations, TIRM topic-model representations, and pre-trained embeddings used with DoubleML, alongside CEM and DoubleML estimators.SAE and STM representations are matched in dimensionality for comparison, while pre-trained embeddings also initialize SAE training.
  • Baseline Pipelines: Table 2 compares selected and unselected SAE representations with TIRM, embeddings, and unadjusted estimates across 100 20NG simulations.The table reports CEM treated retention and indicates whether lower or higher values are preferable for each metric.
  • Metrics: Evaluation reports mean bias, RMSE, confidence-interval coverage, effective sample size, and absolute standardized mean difference across simulations.Higher effective sample size indicates better overlap, while balance is measured on the true confounder.

6 Results

Across semi-synthetic evaluations, SAE-based pipelines generally improve treatment-effect estimation and matching diagnostics, while multi-label settings expose important limitations. Feature selection improves overlap and balance, but its behavior depends on confounding complexity and SAE hyperparameters.

  • Binary and multi-class (20NG): Across binary and multi-class 20NewsGroups settings, SAE approaches achieve the lowest bias and RMSE and the highest coverage.For the multi-class setting, SAE+CEM bias is 0.0205 and SAEselect+CEM bias is 0.0267, versus 0.1572 unadjusted.
  • Binary and multi-class (20NG): SAE pipelines outperform TIRM and Embed+DoubleML on multi-class bias, with SAE bias ranging from 0.0205 to 0.0579.The comparison covers both CEM and DoubleML, with and without feature selection.
  • Binary and multi-class (20NG): Selection retains far more treated and control units after CEM matching while maintaining similar bias metrics.With DoubleML, selection offers little additional ATT benefit, possibly because nuisance functions learn feature weighting implicitly.
  • Multi-label (EURLEX): In EURLEX multi-label evaluations, SAE+CEM has low bias and high coverage but retains less than 1% of treated units, whereas TIRM retains approximately 21%.TIRM achieves similar bias metrics but lower coverage, illustrating a tradeoff between estimation quality and retained sample size.
  • Multi-label (EURLEX): DoubleML reduces multi-label bias from 0.2863 unadjusted to 0.1461 with STM and 0.1324 with SAE, but estimates remain far from zero.The outcome nuisance estimator fails to converge, and many confounding buckets lack common support in finite samples.
  • Selection analysis: Feature selection substantially increases effective sample size and generally maintains better confounder balance, except for U = religion.For religion, post-selection |SMD| remains around 0.2, below the unadjusted value of 0.55.
  • Hyperparameter analysis: With M = 32, nearly all SAE neurons are selected for multi-class and multi-label confounding, indicating insufficient expressiveness for these label sets.This provides an empirical falsification check: selecting nearly all neurons suggests a larger SAE is needed.
  • Hyperparameter analysis: Greater sparsity generally improves RMSE and data retention, whereas lower sparsity improves coverage in the multi-label setting.The results indicate that K should be chosen according to the metrics practitioners prioritize.

K Bias RMSE

In the multi-label setting, SAE sparsity creates a trade-off: greater sparsity generally improves RMSE and retains more data, while less sparsity improves coverage. Selected SAE features are also inspected for interpretability and relevance to the true confounder.

  • K Bias RMSE: For SAEselect with M = 128, smaller K means greater sparsity because only K neurons are retained.
  • Interpretability: Selected SAE neurons are qualitatively inspected with an LLM to assess whether their learned features are interpretable.The inspection uses neurons selected at least 40 times across 100 simulations.
  • Interpretability: The frequently selected neurons are highly relevant to the true confounder in the authors’ data.This provides face-proximal support for using SAE features as interpretable adjustment variables.

7 Conclusion and Outlook

The paper develops a statistically informed SAE-subsetting pipeline for text confounding and extends evaluation to more complex causal simulations. It finds potential for SAE-based causal analysis while identifying multi-label performance and pipeline tuning as open areas for investigation.

  • 7 Conclusion and Outlook: The paper subsets SAE representations with statistically informed selection to adjust for text confounders in more complex causal simulations.
  • 7 Conclusion and Outlook: SAE features may encode confounding information beyond topics, including stylistic and semantic variation.The authors identify more diverse text confounders, such as syntactic variation, as a future direction.
  • 7 Conclusion and Outlook: The authors call for deeper investigation of multi-label performance, jointly learned representations and estimators, and systematic pipeline hyperparameter optimization.

Limitations

The evaluation is bounded by causal identification assumptions and by simulation and modeling choices that may affect results. The authors also defer comparisons across SAE architectures to future work.

  • Limitations: Results from semi-synthetic simulations may vary under different confounder labels, confounding strengths, treatment strengths, representation dimensions, or other parameter settings.The authors mitigate this concern by motivating decisions, following prior best practices, and reporting multiple settings and datasets.
  • Limitations: Using shared hyperparameter settings makes pipeline comparisons fair and principled, but optimal settings may differ across representations.The paper notes that CEM binning rules could differ for SAE and STM representations.
  • Limitations: The study does not compare other SAE architectures, which are left for future investigation.The authors characterize these comparisons as refinements rather than prerequisites for using the proposed adjustment approach.
  • Identification assumptions: Causal estimates rely on untestable assumptions including no unmeasured confounding, overlap, consistency, and homogeneous effects.

B Data Statistics

This section documents preprocessing and descriptive statistics for the 20NG and EURLEX semi-synthetic evaluations, then defines the treatment-effect, overlap, balance, and weighting metrics used in analysis. It also records implementation choices for SAE, baseline estimators, and matching.

  • 20NG: The 20NG dataset’s original labels are merged into 10 experimental labels, including computer, politics, religion, sport, cryptography, and electronics.
  • Treatment assignment: Treatment assignment distributions are visualized for 20NG computer, religion, and multi-class confounding scenarios to assess overlap.
  • 20NG: Figure 5 reports 20NG label prevalence as the fraction of total samples.
  • EURLEX: EURLEX uses the five most common coarse-level labels as a vector-valued multi-label confounder: trade, agri-foodstuffs, geography, agriculture, forestry and fisheries, and EUROPEAN UNION.
  • Treatment assignment: EURLEX treatment assignment satisfies overlap overall, but certain regions have poor overlap.
  • Metrics: ATT measures the average treatment effect among treated units, while coverage is the percentage of confidence intervals containing the true treatment effect.
  • Metrics: ESS assesses covariate overlap from balancing weights, and SMD measures balance on the true confounder U.
  • Matching: CEM bins covariates, drops bins containing only treated units, and assigns weights to matched observations.

E Results: Complementary Figures and Tables

The complementary results report adjustment performance across additional confounding settings and clarify how matching retains comparable treated and control samples. They also vary the CEM cutpoint specification as a robustness comparison.

  • The unadjusted treatment effect is the difference between mean outcomes in treated and control samples.
  • T-ret.% measures the proportion of treated samples retained after matching, while control retention measures the comparable control-supported population.
  • Table 4 reports the percentage of control samples retained across confounding settings and pipelines using the CEM estimator.
  • The matching estimator is evaluated with a 50% quantile cutpoint on positive values in Tables 2 and 3, and an 80% quantile cutpoint in Table 5.
  • The complementary figures show adjustment performance for binary computer, binary religion, and multi-labeled confounding settings.

F Interpretability of SAE features: Complementary Figures and Tables

The complementary materials examine SAE-feature interpretability across binary, multiclass, and multi-label confounding settings. They use repeated simulations and feature-activation interpretations to identify recurring neurons and characterize their semantic content.

  • Interpretability figures: Top-five neuron figures summarize the most frequent SAE features across 100 simulations for computer, religion, multiclass, and EURLEX multi-label confounding.The figures cover 20NG computer, 20NG religion, 20NG multiclass, and EURLEX multi-label settings.
  • Complementary adjustment figures: Figure 12 reports simulation results across 10 of 100 randomly selected seeds, with horizontal lines showing the average estimate across those simulations.For SAEselection+CEM, the number of selected features appears above the confidence-interval bars.
  • Neuron interpretations: EURLEX multi-label interpretations include veterinary and animal-health legal acts, including directives and disease-related materials.
  • Neuron interpretations: The interpreted 20NG computer neurons include X11 graphics, X Window tools, bitmap formats, PC hardware buses, and classified listings.
  • Complementary tables: The complementary tables include a continued Table 6 and a Table 5 comparison using CEM cutpoints at 0 and the 80% quantile of positive values.
  • Neuron interpretations: The 20NG religion interpretations include Bible references and morality or ethics discussions.
  • Neuron interpretations: The multiclass interpretations include computer-related and religion-related features, including X11 topics, graphics formats, Bible references, and morality discussions.
Loading 2609.01322v1…