Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
49,199 characters · 27 sections · 30 citation commands
Causal Effect Estimation with Latent Textual Treatments
\twocolumn[ \icmltitle{Causal Effect Estimation with Latent Textual Treatments}
\icmlsetsymbol{equal}{*}
\icmlaffiliation{huji}{School of Computer Science and Engineering, The Hebrew University of Jerusalem} \icmlaffiliation{stanfordecon}{Department of Economics, Stanford University} \icmlaffiliation{stanfordbusiness}{Graduate School of Business, Stanford University}
\icmlcorrespondingauthor{Omri Feldman}{[email removed]}
\icmlkeywords{Machine Learning, ICML}
\vskip 0.3in ]
\printAffiliationsAndNotice
Whether it is the aspects of political speech that are effective in creating voter support or how social media posts can be most engaging, understanding the causal effect of textual features on downstream outcomes is of great interest across domains gerber2011large, fenerty2012effect, egan2013skilled, fong2016discovery, green2019get, luong2021text, ash2023text, ellickson2023estimating. However, researchers face two particular challenges when designing experimental pipelines for such questions. First, they need to be able to modify text in systematic ways, while ensuring that documents remain readable. Second, valid causal estimation requires a separation between those intended modifications and other variations in the text feder2022causal. While large language models (LLMs) have revolutionized our ability to generate texts at scale, addressing both of the above challenges requires several key methodological enhancements.
In this paper, we propose a novel framework for applying controlled interventions on textual features and estimating their causal effect on downstream outcomes. We build on the high-dimensional representation of language in LLMs to isolate semantic features and systematically intervene on those representations, producing text that varies primarily in targeted features. To facilitate a methodical choice between various candidate features, we also introduce a scoring metric that quantifies the effectiveness of our interventions and the potential undesired confounding artifacts it may introduce. We then estimate causal effects on downstream outcomes using methods from causal machine learning.
To understand the appeal of this approach, consider the following example. martinvenugopal recently assembled a dataset of transcripts of local government meetings across various towns in California. A key feature of these meetings is a public comment period, in which members of the public may address the municipal government directly. These comments can vary widely in their “civility", for which a score is provided for every observation in the dataset. A social-science researcher may be interested in understanding which aspects of speech result in a comment being regarded as more or less civil. Since civility is a complex concept, there may be many atomic features in the text that can impact the perceived civility of the comment. Our framework provides a researcher with a tool to generate hypotheses about which atomic features impact perceived civility and, subsequently, generate quasi-counterfactual comments that vary in these features.
The researcher's goal may then become to validate these hypotheses through a randomized experiment, in which these quasi-counterfactual texts are shown to respondents who evaluate their civility. Similarly, the researcher may instead be interested in seeing if changes in concepts in a text that are correlated with civility impact some alternative downstream outcome of interest, such as respondent agreement. Assessing causality in such a setting presents unique challenges. While randomization of steered texts to respondents will ensure potential outcomes are uncorrelated with treatment assignment, steering itself is rarely “surgical", and will likely perturb correlated nuisance features. As a result, steered texts are not true counterfactuals, and we control for such unwanted variations by adjusting for dense contextual embeddings using robust causal machine learning.
In implementing our approach to causal estimation, we must address a fundamental challenge to causal identification: since the treatments (along with possible confounders) are expressed in the text, embeddings produced by a sufficiently expressive model will perfectly predict treatment itself. Using such embeddings as controls will (in the parlance of causal inference) lead to a positivity violation, with no overlap between treated and control documents and degenerate propensity scores. This leads to a central trade-off: controlling for more sophisticated text embeddings destroys variation in treatment, while weaker embedding models fail to provide sufficient control for nuisance features, leading to omitted-variable bias.
To resolve this tension, we propose a novel residualization procedure which removes treatment information from the embeddings before proceeding with treatment effect estimation. We provide theoretical results that establish causal identification and bound the resulting estimation error. We then validate this approach in practice via semi-synthetic simulations and show consistent, marked improvements over naive estimation procedures that control for full embeddings directly, in line with our theoretical guarantees (see example in \hyperref[fig:cate_estimation]{Figure \ref*{fig:cate_estimation}} ). Our approach constitutes a robust foundation for causal effect estimation with latent treatments, and can be extended naturally to other modalities.
In this section, we formalize the notation for Sparse Autoencoders and the causal inference framework used throughout.
Sparse Autoencoders are designed to decompose the dense, polysemantic activations of an LLM into a high-dimensional, sparse latent space that corresponds more closely to interpretable features bricken2023monosemanticity.
Let $a \in \mathbb{R}^d$ be the vector of activations from a specific layer of a language model. An SAE consists of an encoder $W_{enc} \in \mathbb{R}^{h \times d}$ and a decoder $W_{dec} \in \mathbb{R}^{d \times h}$, where the latent dimension $h$ is typically much larger than the hidden dimension $d$. The encoder maps the activations into a sparse feature vector $z \in \mathbb{R}^h$ and the decoder maps $z$ back to $\mathbb{R}^d$,
where $\sigma(\cdot)$ is a nonlinear activation function (e.g. ReLU, JumpReLU, TopK) and $b_{enc}$ and $b_{dec}$ are bias terms. During training, a sparsity penalty is applied to $z$ to ensure that only a small number of features are active for any given input. We denote the indices of the SAE latent space as $\phi \in \Phi$, latent feature activation as $z_\phi$, and the $\phi$-th column of the decoder matrix as $W_{dec(\phi)}$.
Our ultimate goal is to estimate the causal effect of controlled variation in text features on downstream outcomes of interest. We adopt the potential-outcomes framework from causal inference. We write $T \in \{0, 1\}$ for a binary treatment variable (e.g., the presence of a specific linguistic feature) and $Y$ for an observed downstream outcome. For a given unit, we consider potential outcomes $Y(1)$ and $Y(0)$, representing the outcome under treatment and control, respectively. The realized outcome is $Y = Y(T)$. The individual treatment effect is $Y(1) - Y(0)$, which is unobservable. We aim to estimate the Conditional Average Treatment Effect (CATE)
which is the expected effect of the treatment given a set of covariates $\tilde{X}{=}\tilde{x}$. In the context of text-as-treatment, $\tilde{X}$ typically consists of high-dimensional contextual embeddings that capture the nuisance properties of the text (e.g., length, base semantic content, style). Standard identification of $\tau(\tilde{x})$ requires two fundamental assumptions:
Throughout, we will maintain the additional assumption that units $i$ with observations $(Y_i, T_i, \tilde{X}_i)$ are drawn i.i.d from a population distribution. In particular, the SUTVA assumption of no interference is implicit: potential outcomes for a given unit are independent of the treatment status of other units, $(Y_i(0), Y_i(1)) \perp T_{j\neq i}$.
SAEs have been shown to be useful in decomposing the polysemantic representations of LLMs into monosemantic, human-interpretable features Bills2023Autointerp, bricken2023monosemanticity, cunningham2023sparseautoencodershighlyinterpretable, templeton2024scaling. Recent research has leveraged the granularity of these features for steering and hypothesis generation movva2025sparseautoencodershypothesisgeneration, peng2025usesparseautoencodersdiscover, zheng2025model. However, alongside their usefulness for interpretability, benchmark results wu2025axbenchsteeringllmssimple, huang2024ravel, kantamneni2025sparseautoencodersusefulcase suggest that SAEs and SAE-based steering perform poorly compared to other methods. Our work extends this literature by proposing a new metric based on the model's response to a steering intervention, establishing a more rigorous link between a feature's interpretability and its utility in downstream causal experiments.
Estimating causal quantities with textual data presents multiple unique challenges feder2022causal, egami2022make. Recent work addressed these challenges when the textual information is used as controls veitch2020adapting, mozer2020matching, treatment fong2016discovery, pryzant2021causal, gui2022causal, ash2023text, angelopoulos2024value, and outcome modarressi2025causal. Particularly relevant to our work is imai2025causalrepresentationlearninggenerative, which studies a similar setting under slightly different assumptions. We contribute to this literature by presenting a novel steering pipeline that enables controlled interventions and addresses inherent positivity violations in the text-as-treatment setting.
Throughout this paper, we present empirical results from three labeled text datasets.
Dataset A: Local Government Speech. A dataset of city council meetings in California produced by martinvenugopal. The dataset is comprised of statements made by members of the public in city council meetings across California. The label corresponds to a binarized measure of “civility": a label of 1 indicates a highly civil statement, while a label of 0 corresponds to a highly uncivil statement.
Dataset B: Political Ads. A dataset of US political advertisements, as presented in sobhani2024multienvironmenttopicmodels. A label of 1 indicates a Republican politician and a label of 0 indicates a Democratic politician.
Dataset C: Us Vs. Them. A dataset of reddit comments posted in response to news articles across the political spectrum, curated by huguet-cabot-etal-2021-us. The comments were scored for populist attitudes towards out-groups, which we binarize such that a label of 1 indicates a discriminatory comment towards the out-group and a label of 0 indicates a supportive comment towards the out-group.
The hypothesis generation phase identifies a set of candidate SAE features $\Phi_+ \subset \Phi$ that differentiate between semantic classes. In particular, we use a labeled dataset $\mathcal{D} = {(X_i^\mathcal{D}, Y_i^\mathcal{D})}_{i=1}^{N}$ to find the features that correlate with the target concept. To narrow the search space from the high-dimensional and sparse SAE latent space, we first apply a series of standard filtering operations that removes spurious and noisy features (see \hyperref[sec:candidatefiltering]{Appendix \ref*{sec:candidatefiltering}} for details).
To ensure features are comparable, we calculate $z$-scores of the activations of the remaining features. For each feature $\phi$, we then calculate the absolute mean difference
in normalized activations $Z_\phi$ between the two classes. We retain only the $k$ features with the highest $\Delta_\phi$, experimenting with $k \in \{16, 32, 64, 128, 256\}$. To identify which of these top-$k$ features are the most important for the classification, we train an $L_1$-regularized logistic regression model (i.e. a linear probe). We perform $K$-fold cross-validation to select the optimal regularization parameter. After identifying the best-performing model, we extract the coefficient $\beta_\phi$ for each feature and rank the features by $|\beta_\phi|$. Given that the coefficients of an $L_1$-regularized Logistic Regression are biased and sensitive to the number of features, we turn to feature persistence as a way to rank our SAE features and measure the median position of each feature across the different regression fits. This gives a higher ranking to features with a consistently high coefficient.
To transform our ranked candidates into a set of steering vectors, we map each feature index to its corresponding semantic description provided by Neuronpedia neuronpedia. Our analysis shows that the features identified by the sparse probes indeed represent the multi-faceted nature of the target concepts, as captured by our training datasets. Specifically, we see that some features relate to specific linguistic markers, and some to broader, topical concepts.
While the linear probe identifies features as predictive, we leave it to the practitioner to determine which specific aspects of the target concept are most relevant as a causal hypothesis. For the purpose of this work, we construct $\Phi_+$ by manually curating a subset of features that exhibit both high persistence and clear semantic alignment with the core aspects of our datasets. This human-in-the-loop filtering ensures that the resulting steering experiments are relevant and target interesting scientific questions.
With the candidate features $\Phi_{+}$ identified, we generate quasi-counterfactual texts by intervening on the latent representation of the specific feature $\phi$ in the SAE, where each such feature can serve as a treatment in a downstream experiment.
We prompt the LLM to act as a relevant persona and ask it to respond to or comment on an input text, without constraints on tone, sentiment, or level of agreement. From a research design perspective, this allows practitioners to implement a wide array of studies. For example, LLMs can also be prompted to rephrase, summarize, or provide assistance. Moreover, our proposed approach is not limited to using the same dataset for discovery and steering, with trivial cross-dataset transferability.
In each forward pass of the model we implement an adaptive steering mechanism that scales the intervention according to the local token representation. Let $a \in \mathbb{R}^d$ represent the residual stream activations and $\|a\|_2$ represent the $L_2$ norm of $a$. For a target feature $\phi \in \Phi_{+}$, we define the steered latent vector $z'$ such that
where $\alpha \in \mathcal{A}$ is the steering factor. Unlike global steering methods, this approach ensures that the steering operation scales with the current state of the residual stream, preserving its relative magnitude across varied contexts. This is particularly useful in our setting for comparing steering interventions across layers and models.
The modified activations are then reconstructed and passed back to the model, similarly to (ref). To maintain the original model's performance and nuances, we preserve the SAE reconstruction error $\epsilon_a = a - \hat{a}$:
We ran our experiments on three instruction-tuned LLMs: Gemma-2-9B-IT, gemmateam2024gemma2improvingopen, Llama-3.1-8B-Instruct grattafiori2024llama3herdmodels and Qwen2.5-7B-Instruct qwen2025qwen25technicalreport. For Gemma-2-9B-IT, we used 16k SAEs for layers 9, 20, and 31 lieberum2024gemma. For Llama-3.1-8B-Instruct and Qwen2.5-7B-Instruct, we used 131k SAEs for layers 7, 15 and 23 arditi2024finding. Refer to \hyperref[sec:appendixsteering]{Appendix \ref*{sec:appendixsteering}} for further details, prompts and examples.
While the steering factor $\alpha$ is a controlled hyperparameter of the intervention, the actual concept intensity of the generated text can vary across contexts, specific features, and models. Thus, to evaluate the success of our intervention, we define the ex-post concept intensity as the mean token-level cosine similarity between the model's activations and the relevant decoder column. For a sequence $\nu$ of length $M$, the raw intensity $I$ is given by:
To ensure comparisons across different features, layers, and models, we introduce the normalized intensity score $I^*$:
where $\widehat{I}$ is a linear regression fit of mean intensities $\bar{I}_{\alpha}$ on steering factors $\alpha \in \mathcal{A}$, $|\hat{\beta}_1|$ is the regression coefficient, $R_I$ is the range of $\bar{I}$, and $\text{MAE}(\bar{I}, \widehat{I})$ captures the residuals.
Intuitively, this scoring rewards features that exhibit a consistent, proportional response to steering. A high score indicates that the concept intensity scales linearly with the intervention magnitude, whereas a low score penalizes features that saturate early, plateau, or otherwise deviate from a proportional linear response. See \hyperref[fig:full_intensity_spectrum]{Figure \ref*{fig:full_intensity_spectrum}} for examples.
We also wish to identify features that do not degrade the model's output quality. We prompt Gemini 2.5 Flash-Lite comanici2025gemini25pushingfrontier and ask it to rate each steered response on a scale of 0 to 3 based on its coherence and fluency (see prompt in \hyperref[sec:apxjudgeprompt]{Appendix \ref*{sec:apxjudgeprompt}} ). To compare features, we aggregate these scores into a single coherence metric
representing the average quality, where $\bar{J}_{\alpha}$ is the mean coherence score for steering factor $\alpha$.
To select the most effective features for our experiments, we evaluate each candidate $\phi$ across the steering range $\mathcal{A}$. We combine the intensity and coherence measures into a final IC Score that ranks features that are both preferable in their activation response to steering and qualitatively stable in their generated output:
The IC Score serves as a rigorous metric for identifying well-behaved features. This is useful for using these features as treatments in causal experiments for several reasons. First, features with higher scores allow for steering over a more extreme or spread-out range of $\alpha$, enabling more pronounced treatments while ensuring the output remains semantically coherent. Second, high coherence across a wide intensity range allows us to sample from multiple points along the curve, facilitating the study of multi-valued or continuous causal effects rather than being limited to binary interventions. Third, a high IC Score enables the “blind" setting of steering factors to produce reasonable texts for downstream causal estimation, while a lower score would require a manual, per-case validation to identify the range of suitable steering factors. Fourth, a high score theoretically mitigates the risk of confounding. If $\alpha$ increases and $I^*(\phi)$ does not increase proportionally, it becomes more likely that adding the steering vector to the residual stream will cause an increase in other correlated concepts. \hyperref[tab:ic_scores_combined]{Table \ref*{tab:ic_scores_combined}} shows the mean IC Score per model and layer position.
Having produced texts that systematically vary in the intensity of a target feature, we now aim to estimate causal effects on downstream outcomes, such as the concept originally targeted in the hypothesis-generation step. Specifically, we want to estimate conditional average treatment effects (CATEs) of the form in (ref). In doing so, we have to address the challenge that steering may affect not only the target feature, but also other aspects of the text. As we demonstrate in this section, controlling for such other features poses unique challenges that we address in our formal results and practical implementation based on robust causal machine learning. All proofs can be found in \hyperref[apx:proofs]{Appendix \ref*{apx:proofs}} .
We now focus on estimating the effect of a specific target feature $\phi$, and we define the treatment $T_\phi$ of interest to be the (rescaled) measured intensity of that feature, $I(\nu,\phi)$. For simplicity, we assume that feature intensity is either high ($T_\phi{=}1$) or low ($T_\phi{=}0$). We write $X$ for a full embedding of the text, which we assume contains all relevant information about a document. Beyond the concept we intervene on, texts (as represented by $X$) also contain other important semantic features (e.g., topic, tone, etc.) separate from $T_\phi$ that may be different between texts with different target intensities, which we collectively write as $X^\perp = g^\perp(X)$.
We are interested in estimating the effect of the intensity $T_\phi$ on a downstream outcome $Y$. We assume that $T_\phi$ and $X^\perp$ together capture any aspect of the text relevant for downstream outcomes, so that the outcomes can be written as a function of $T_\phi$, $X^\perp$, and some exogenous (independent) noise $\varepsilon$: $Y = f(T_\phi,X^\perp; \varepsilon)$. In line with the notation in \hyperref[sec:TE-Intro]{Section \ref*{sec:TE-Intro}} , for every unit we consider potential outcomes
which represent pairs of (correlated) random variables that express the downstream outcomes of a hypothetical pair of texts that only differ in the target intensity (treatment), but not otherwise. The treatment effects of interest are averages over $Y(1) - Y(0)$. The realized outcome is $Y = Y(T_\phi)$. Throughout, we assume that $Y$ is scalar with $E[| Y |^2] < \infty$.
As we aim to estimate treatment effects $Y(1) - Y(0)$ (and their averages), we face two challenges: First, steering may affect parts of the text $X$ beyond the target concept, so the components $X^\perp$ may correlate with the target intensity (treatment) $T_\phi$, violating unconfoundedness. Second, since the concept intensity $I(\nu,\phi)$ is measured from the text, simply controlling for the full text embedding $X$ implies that treatment $T_\phi$ is fully explained, violating positivity.
To illustrate the first challenge, assume that we aim to estimate the overall average treatment effect $\tau = E[Y(1) - Y(0)]$, but that $X^\perp$ differs between the high intensity ($T_\phi{=}1$, treated) and low intensity ($T_\phi{=}0$, control) groups. Then, simply taking the difference in means between the two groups generally leads to bias,
due to confounding. One way in which we could overcome this challenge would be to assume that steering does not create any systematic differences across groups:
In practice, this assumption may be unrealistic: as steering creates systematic differences in the target intensity, it likely also affects related (and possibly even unrelated) textual features in systematic ways. Instead, we may want to control for possible confounders, leading to the second challenge. Indeed, note that one way of avoiding confounding would be to control for the text embedding $X$. However, since $X$ also perfectly captures the target intensity $T_\phi$, we have that $ P(T_\phi{=}1|X) \in \{0,1\}$ -- positivity is violated, and we lose identification of effects of $T_\phi$.
Above, we have argued that the estimation of treatment effects is hindered by confounding through other changes in texts, and that a naive approach to controlling fails. In theory, we could address both challenges by controlling only for changes in the aspects of the text that are separate from the treatment $T_\phi$, as captured by $X^\perp$:
Implementing an estimator for $\tau$ in practice based on this result would require separating out the component $X^\perp$ that is orthogonal to the intensity (treatment) $T_\phi$. However, we do not generally measure this component $X^\perp$ separately, and only have access to the full embedding $X$ and the implied treatment $T_\phi$. Instead, we now assume that we have a possibly imperfect vector $\tilde{X} = \tilde{g}(X)$ of control variables available that balances the desire to control for $X^\perp$ with the requirement to preserve variation in treatment $T_\phi$:
The first assumption requires sufficient variation in intensity $T_\phi$ after controlling for $\tilde{X}$ and implies that $\tilde{X}$ cannot capture too much of the full embedding $X$. The second assumption says that the remaining systematic variation beyond $T_\phi$ in the embedding is not too large after controlling for $\tilde{X}$. It would be fulfilled in the two cases from above: If steering is perfect ( \hyperref[prop:tightsteering]{Proposition \ref*{prop:tightsteering}} ), then $W_1^{(d)}(\nu_1(\tilde{X}), \nu_0(\tilde{X})) = 0$ a.s.\ for any choice of control variable. If $\tilde{X} = X^\perp$ ( \hyperref[prop:perfectcontrolling]{Proposition \ref*{prop:perfectcontrolling}} ), then likewise $W_1^{(d)}(\nu_1(\tilde{X}), \nu_0(\tilde{X})) = 0$ a.s.\ since there is no remaining variation in $X^\perp(1), X^\perp(0)$.
We now consider how well we can estimate treatment effects by controlling for the (possibly imperfect) proxy $\tilde{X}$ of $X^\perp$. We assume that our goal is to approximate conditional average treatment effects (CATEs) $\tau(\tilde{x}) = E[Y(1) - Y(0)|\tilde{X} {=} \tilde{x}]$ by estimating
from the steered data. Writing
we assume that $\mu_1, \mu_0$ are not too sensitive to small variations in $x^\perp$:
We obtain a bound for the mistake we make when we approximate $\tau(\tilde{x})$ in terms of $\tilde{\tau}(\tilde{x})$:
The same bias bound applies if we approximate the overall treatment effect $\tau = E[Y(1) - Y(0)]$ by controlling for $\tilde{X}$. Hence, we can still recover a suitable estimator based on transformations $\tilde{X}$ that capture confounding features well enough, while not perfectly predicting $T_\phi$.
Above, we have argued that an approximation for $\tau(\tilde{x})$ with bounded bias may be achieved by controlling for appropriate text features $\tilde{X}$. We implement this controlling procedure based on robust causal machine learning chernozhukov_dml using the R learner nie2021quasi as implemented in EconML econml. Writing $\tilde{\mu}(\tilde{x}) = E[Y |\tilde{X}{=}\tilde{x}]$ (outcome prediction) and $\tilde{\pi}(\tilde{x}) = E[T_\phi |\tilde{X}{=}\tilde{x}]$ (propensity score), our target estimand $\tilde{\tau}(\cdot)$ defined in (ref) solves the least-squares problem
Motivated by this identity, we predict $Y$ and $T_\phi$ from $\tilde{X}$ on our sample $(Y_i,T_{\phi i},\tilde{X}_i)_{i=1}^n$ using cross-fitting to obtain out-of-fold estimates $\hat{\mu}_{-i}(\tilde{X}_i), \hat{\pi}_{-i}(\tilde{X}_i)$ of $\tilde{\mu}(\tilde{X}_i), \tilde{\pi}(\tilde{X}_i)$. We then solve the sample analogue
over regularized linear functions to estimate the approximate CATEs $\tilde{\tau}(\cdot)$. This procedure is robust: it leads to a high-quality estimation of $\tilde{\tau}(\cdot)$ even if $\tilde{\mu}, \tilde{\pi}$ are both estimated only at a rate $o(n^{-1/4})$ nie2021quasi.
We validate our steering and causal estimation pipeline in a series of semi-synthetic simulations. We sample 300 texts from each of the real-world datasets from \hyperref[sec:Datasets]{Section \ref*{sec:Datasets}} , construct $\Phi_+$ with three features per layer, and run 25 steering intervention for each text with equally spaced steering factors $\alpha \in [0, 1.2]$. This procedure produces $7,500$ texts per feature and an overall of $607,500$ generated texts.
We use the LLM-as-judge coherence measure to remove all steered texts the judge flagged as invalid. For the remaining set, we construct our binary treatment as whether a steered text has a measured feature intensity in the top ($T_\phi{=}1)$ or bottom ($T_\phi{=}0)$ quintile of the given sample distribution of $I(\nu,\phi)$. Among these selected extremes, we remove all steered texts originating from base texts that do not produce at least one steered text in each treatment group. This allows us to ensure that there are no systematic differences between the two groups arising from differences in the corresponding sets of base texts. We then note the number of steered texts produced by each base text in each treatment group, and use the inverse of this count as the weight assigned to the corresponding steered text in CATE estimation. While we run our method on a wide selection of SAE features, we remove features with a very low IC score from our analysis. For an overview of the results for those features which do not meet this criterion, see \hyperref[apx:spec]{Appendix \ref*{apx:spec}} .
We produce embeddings $X_i$ for each document via EmbeddingGemma-300M vera2025embeddinggemmapowerfullightweighttext, all-mpnet-base-v2 song2020mpnetmaskedpermutedpretraining, or all-MiniLM-L6-v2 wang2020minilmdeepselfattentiondistillation. The resulting embedding vectors are then stacked row-wise into a matrix for the given data sample. We begin by using PCA to rotate this embedding matrix, producing the matrix $\mathbf{X}$ with $n$ rows, one per document. \footnote{When performing PCA, we keep the maximum number of dimensions (given by the minimum of the number of steered texts and the number of embedding dimensions).} This will be our default representation of text for the remainder of our simulation.
Our theoretical results in \hyperref[sec:limitedcontrol]{Section \ref*{sec:limitedcontrol}} show that we can improve causal estimation on downstream outcomes by constructing control variables $\tilde{X}$ based on embeddings $X$ that capture confounding features well enough, while also allowing for substantial residual variation in $T_\phi$. In order to implement this guidance, we construct residualized versions of the full text embeddings. Specifically, we propose two strategies for residualization of our produced embeddings (e.g., producing $\mathbf{\widetilde{X}}$ from $\mathbf{X}$).
Dimension-by-dimension: We separately residualize each dimension (column) of our embedding matrix $\mathbf{X}$: let $\mathbf{x}_j$ denote the $j$-th column of $\mathbf{X}$. We remove treatment information encoded in $\mathbf{x}_j$ by fitting a predictive model to predict $\mathbf{x}_j$ from $\mathbf{T}_\phi$ (which is the vector of treatment indicators $T_{\phi i}$), yielding a predicted value $\mathbf{\hat{x}}_j$\footnote{The model fitting and prediction is accomplished via cross-fitting, with all predictions generated out-of-sample}. We then construct our residualized version $\mathbf{\tilde{x}}_j = \mathbf{x}_j-\mathbf{\hat{x}}_j$. Proceeding iteratively across columns of $\mathbf{X}$ thus allows us to construct the embedding matrix $\mathbf{\widetilde{X}}$ with treatment-related information removed.
Dropping first principal component: We perform residualization by simply dropping the first dimension of our rotated embeddings (the first principal component), e.g. $\mathbf{\widetilde{X}}$ is simply $\mathbf{X}$ with the first column removed.
We can evaluate how much predictive power our (rotated) embeddings contain with respect to $T_\phi$ by measuring the out-of-sample classification accuracy for $\mathbf{X}$ and $\mathbf{\widetilde{X}}$. The results are shown in \hyperref[fig:resid_acc]{Figure \ref*{fig:resid_acc}} . We see that using raw $\mathbf{X}$ almost perfectly predicts $T_\phi$, while both of our residualization strategies show lower accuracy, as expected. We also observe that the first dimension of our rotated embeddings, corresponding to the first principal component, carries substantive predictive power for treatment.
We can further analyze the concentration of treatment information in the first principal component by examining the performance gap in predictive accuracy of treatment between raw and residualized embeddings, component by component. \hyperref[fig:violin_panel]{Figure \ref*{fig:violin_panel}} shows that the effects of residualization are heavily localized to the first principal component, while having a negligible effect on all other components. \hyperref[fig:corr_panel]{Figure \ref*{fig:corr_panel}} goes further and shows that there is a clear positive correlation between the IC score and the accuracy gap on the first principal component. This suggests that SAE features that are more effective for steering also result in concentration of treatment information along the primary axis of embedding variation.
Taken together, these results demonstrate that when an SAE feature is highly effective for steering, dropping the first principal component is likely a sufficient strategy for residualization. For features with lower IC scores, dimension-by-dimension residualization will likely yield better results due to decreased concentration of treatment information in the first principal component. Dimension-by-dimension residualization will be our primary specification in our simulation study; results using residualization via dropping the first principal component can be found in \hyperref[apx:spec]{Appendix \ref*{apx:spec}} .
We construct synthetic outcome variables given $\mathbf{\widetilde{X}}$ from a wide variety of functional form specifications, with details provided in \hyperref[apx:spec]{Appendix \ref*{apx:spec}} . CATE estimation is performed via robust causal machine learning following \hyperref[sec:estimation]{Section \ref*{sec:estimation}} with nuisance functions estimated via random forests breiman_rf. Residualization is similarly achieved via random forests to produce $\mathbf{\hat{x}}_j$. We simulate outcomes $y_i$ across a variety of functional forms, embedding models, LLMs, SAE features, and datasets. For all simulations, the true average treatment effect (ATE) is 5. \hyperref[fig:dgp_compare]{Figure \ref*{fig:dgp_compare}} shows a summary of the results.
These results demonstrate a substantial reduction in both bias and RMSE following the application of our residualization approach. This marked improvement validates the efficacy of our methodology in mitigating the bias inherent in naive control strategies that utilize the full embeddings. Furthermore, these findings are consistent with our theoretical results, confirming that removing the treatment information is essential for robust causal estimation in such settings.
Advances in text generation via LLMs and Sparse Autoencoders (SAEs) offer a transformative opportunity for social science research. However, the utility of these models has been limited by the difficulty of identifying optimal steering features and the bias caused by the entanglment of treatments and covariate information. We address these fundamental challenges by introducing a comprehensive pipeline for hypothesis generation from labeled data, quasi-counterfactual text generation and a novel covariate residualization methodology for causal effect estimation. We demonstrate both theoretically and empirically that we can isolate treatment effects and mitigate the bias that arise in text-as-treatment settings. Our results provide a rigorous new foundation for experimentation with latent treatments, shifting the paradigm from hand-crafted techniques to formal and controllable causal experimentation.
This paper presents work whose goal is to provide social scientists with a new methodology for experimentation and causal effect estimation in text-as-treatment settings. We believe our work extends the existing toolkit available to researchers and can advance the understanding of causal relationships in fields such as economics, political science, business, and marketing. While the development of controlled text generation and causal estimation methods carries broader societal implications for areas such as personalization and targeting, we feel there are no specific consequences that warrant further highlighting here beyond the general advancement of machine learning and its application to social science research.