Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors

We are grateful that the reviewer found our perspective on correctness prediction novel, and we appreciate the opportunity to discuss the practical value of our work!

---

Q1. Generalizability to complex real-world tasks

We believe it's helpful to take a step back to examine the source of generalizability concerns and its relation to the broader goal of this work. Specifically, our framework comprises two key components: (1) finding causal mechanisms in LLMs, and (2) leveraging these mechanisms to predict model behaviors, particularly in OOD settings.

Generalizability questions fall in the design space of problem (1), e.g., how to identify the high-level causal model without a strong human prior, how to scale to more complex tasks, etc. These are indeed important problems that mechanistic interpretability community has been focused on. However, even if we could perfectly identify the mechanisms underlying every GPT-4 prediction, it remains an open question whether such mechanisms are actually useful for predicting model behaviors (Mu et al. 2021; Wang et al. 2023; Sharkey et al. 2025).

Therefore, the focus of this work is to examine this under-explored but crucial step (2). As a result, we choose to build upon existing tools developed by the interpretability community for problem (1). This means we naturally inherit the limitations of these tools.

With this context in mind, we would like to further address limitations of these interpretability tools we used!

1.1 Reliance on human priors to identify causal mechanisms

This is indeed a recognized limitation of concept localization methods. We discussed solutions in our response to reviewer uRDA Q3.

1.2. Scale to real-world complex tasks

We totally agree that there are many real-world tasks whose complex causal structures are not (yet) fully understood. However, our method remains applicable when the causal structure is partially known, and a surprising number of real-world applications fall into this category.

In fact, our experiments already include such tasks, e.g., MMLU contains 57 tasks to evaluate production-grade LLMs. While we may not know the exact causal mechanisms behind answering abstract algebra or international law questions, we can still leverage the shared multiple-choice format as high-level structure for correctness prediction.

Below, we highlight two additional practical domains where our methods are directly applicable:

• LLM evaluation. Constructing benchmarks typically requires expensive human annotations. A reliable correctness estimator allows benchmark creators to prioritize uncertain or error-prone examples. Moreover, systematic evaluation often requires structured inputs, e.g., SQuAD adversarial (Jia et al., 2017), GSM-symbolic (Mirzadeh et al., 2024), which inherently reflect high-level task structures that our framework can readily exploit.
• Verifiable text generation. In high-stakes domains like medical records generation, ensuring factual and referential accuracy. One strategy is to have LLMs generate symbolic references to structured source data that are easy to verify (Hennigen 2024). These symbolic references act as templated prompts, which can be converted to high-level models. Our methods allow estimating the correctness of texts generated using these decoding algorithms.

Q2. Computation cost of counterfactual simulation

We have provided the computation cost in Table 1. Counterfactual simulation is K times more expensive than greedy decoding, where K is the number of counterfactual samples. Empirically, a relatively small set of counterfactual samples, e.g., K=16, is sufficient to achieve high AUC-ROC.

The cost can be further reduced by caching representations before the intervention site. For $Q$ queries from the same task, assume the intervention site is at the middle layer, the average cost per query is $(K+Q+QK)/2Q$, i.e., about $K/2$ times more expensive than greedy decoding when $Q$ is large.

Q3. Which components contribute most to correctness prediction

Thanks for the suggestion! It indeed aligns with our discussion in Section 6.4 on using correctness estimation as an evaluation task for interpretability methods.

We have experimented with attention heads vs residual streams using GPT-2 on IOI task. The original IOI paper localizes the position variable in the outputs of S-Inhibition heads. These sets of heads produce an AUC-ROC around 0.8, while random heads have an AUC-ROC of 0.5. They underperform residual subspaces, likely because the manually identified heads do not fully capture the causal variable.

Q4. Risk of overfitting causal subspaces

As with any supervised method, causal subspaces or probes may overfit to the training distribution. This is exactly why we see a slight drop from ID (Table 2) to OOD settings (Table 3) across all tasks/methods. However, methods using causal features are more robust, i.e. consistently show smaller ID–OOD gaps.

We very much appreciate the reviewer’s suggestion to clarify what is distinctively causal about our methods, and also how what we are doing relates to important neighboring areas of search, like causal discovery and causal representation learning (CRL).

We see this as an opportunity to further clarify the nature of our contribution using the additional space we would get in our next version.

---

Q1. Clarification on the definition of “causal”

Our paper is grounded in the literature on mechanistic interpretability, and more specifically in techniques that perform interventions on networks to understand how they represent data and make predictions. This is in many respects easier than causal discovery or CRL, since the network is a closed deterministic system.

Nonetheless, the field does not yet have clarity on how networks make predictions, and approaches like ours are seeking to address that limitation. (For additional details on causal abstraction, the family of techniques we draw on, please see our response to reviewer uRDA Q2.4).

As we agree that using some of these terms which have technical meanings in important literature (especially “causal representation”) may be confusing to many readers, we intend to reword our summaries of the paper’s main contributions. In particular, we will replace the mention of “causal representation” in the text “assessing stability of causal representations” to “assessing the ability to simulate counterfactual behaviors under causal interventions”. This strikes us as a chance to clarify the diverse ways that causal methods are contributing to AI right now.

We thank the reviewer for the thoughtful feedback! We are especially glad that the reviewer recognizes our work as the first to address correctness prediction in LLMs from a causal perspective.

---

---

Q1. Comparison with strong baselines in correctness estimation

A major concern is that

there is no comparison to any other approach for correctness prediction.

We respectfully disagree with this characterization, as we have compared with two of the strongest baselines in the literature: confidence scores and correctness probing. We do realize that these methods are not explicitly labeled as “baselines” in our result tables, which may have caused some confusion. We will revise the manuscript to make the baseline labels more prominent.

To recap our baselines, we have reviewed existing methods (L160) in Section 4.1 (Confidence Scores with Temperature Scaling) and Section 4.2 (Correctness Probing). These baselines were chosen because they are widely adopted and representative of state-of-the-art techniques in correctness estimation (see Section 2.1 for the line of work, and also Kadavath et al., 2022; OpenAI 2023; Chen et al 2024; Orgad et al 2025 for the most relevant SOTA). We have reported their performance in Table 2 and Table 3 in Section 5.2.

---

Q2. Presentation: Clarification on background variables

2.1 Where does the background variable fit into the causal graph

Please see our responses to reviewer uRDA 2.1 and 2.2

2.2 Are causal and background variables understandable for humans

This is a valuable question from the explainability perspective–namely, how to explain complex models to humans. In our experiment, all causal variables correspond to human-interpretable concepts, as these high-level models were manually specified by interpretability researchers. However, we want to emphasize that our approach to correctness prediction does not require the causal variables to be human-interpretable. The causal role of a variable in predicting model behavior is sufficient, even if it is not directly intelligible.

---

Q3. Details on dataset construction

• For details on split generation, please see our response to reviewer w7xm Q1.
• For label balancing, we perform stratified sampling: we first partition all prompts into two groups based on the correctness of the target model predictions. We then randomly sample 1024/512/512 examples from each group.

---

Q4. Clarification: when does the causal relation $X_T → Y$ hold and when does it not hold

As discussed in our response to reviewer uRDA 2.1, the causal relation $X_T → Y$ is an abstraction of low-level neural networks.

This abstraction is faithful when $X_T$​ mediates the model’s behavior. However, neural networks often fail to generalize beyond their training distributions, where the model’s actual behaviors might diverge from what the high-level causal model $H$ predicts. When this occurs, we say the high-level causal relation $X_T → Y$ no longer holds, i.e., it is no longer a faithful abstraction of the low-level neural network model. These cases correspond precisely to our out-of-distribution settings.

---

Q5. Clarification: Why does the counterfactual dataset consist only of correctly predicted samples

This is indeed one of the interesting findings from our work: it is sufficient to learn a strong correctness predictor using only correctly predicted examples. This finding might make more sense if we consider confidence score methods, where correctness predictions are made using the probabilities output alone without using any additional samples.

The question of why Counterfactual Simulation works is in fact tied to the core question we are asking in this work: does understanding the internal causal mechanisms allow us to better predict model behaviors, especially under distribution shift?

Intuitively, if we know that (1) for all correctly predicted examples, the model behaves according to a high-level causal model $H$ (i.e. implements a systematic solution), and (2) for an unseen example, the model does not implement the same solution, then it is very likely that the model is predicting something abnormal and so the prediction is likely wrong. Mathematically, this intuition is formalized in Eq 7-9, where the counterfactual simulation measures whether the model implements the same solution as we have observed on the correct samples.

---

Q6. Presentation

We appreciate the reviewer’s suggestion to clarify the definitions of "Distributed Alignment Search" and "Interchange Intervention Accuracy."

We will include preliminary explanations of these terms in the revised manuscript!

We thank the reviewer for their insightful comments! We are encouraged that they recognized the novelty of our methods, the strength of our empirical results, and the significance of bridging the explanatory and predictive aspects of interpretability analysis.

---

---

Q1. Clarification on the linear decomposition assumption

Linearity is indeed an assumption made in DAS. We have discussed its implications in the presence of non-linearity and potential remedies in our response to reviewer w7xm 2.4 and offer additional evidence in Q4 below.

---

Q2. Clarification on the causal assumptions and causal model

2.1 Causal structures in causal abstraction

We also appreciate the reviewer’s request to clarify the causal aspects of our approach, which we recognize does rely on previous work on causal abstraction. There are in fact several causal structures at play, and causal abstraction is about the relationship between different causal structures. Most basically, there is the causal structure of the neural network itself: a sequence of layers of variables with each layer depending functionally on the previous. But there is also a second causal structure given by the “high level” model $H$ that, in our running example, is given by just three variables, $X, X_T$, and $Y$. The variable B does not occur in $H$. Moreover, $X_T$ is not unique in the sense that we can have different high-level models representing different levels of abstraction.

DAS then involves searching for $X_T$ somewhere in the network, such that the interchange interventions are successful. When that happens, we say that H is an approximate abstraction of the network, in the sense that $X_T$ mediates the transformation from task $T$ input to output. In the limit, DAS is guaranteed to find this encoding of $V_T$, provided it exists. But note that basis $Q$ is not guaranteed to be unique.

2.2 The background variable $B$

With this much, we can say exactly what $B$ is. $B$ is just the orthogonal complement of the identified representation of $X_T$ (e.g., somewhere in the residual stream). It lives outside the simple high-level structure $H$ and is determined by the structure of the network, namely how it encodes $X_T$. We call it a background variable because it encompasses everything about the input that does not feed into $X_T$ (and ultimately determine the output $Y$).

2.3 Identifiability Please see our response to reviewer w7xm 2.3.

2.4 Modeling more complex task structures

We totally agree that this three-variable high-level structure can be overly restrictive. Indeed, most tasks studied in our paper involve significantly more complex structures. For instance, in the IOI task, at least 6 causal variables are involved (Figure 2, Wang et al. 2023), including 3 input variables (IO, S1, S2), 2 position variables (outputs of the Induction and the S-Inhibition Heads), and the output. We use the second position variable as $X_T$ to perform counterfactual simulation. We will clarify in the revised manuscript that the three-variable model is for illustrative purposes, and our methodology generalize to more complex settings as demonstrated in our experiments.

---

Q3. Method assumes knowledge of task structure to define $X_T$

This is a valid concern, and indeed a known limitation. However, the interpretability community has proposed auto-interp methods to reduce the reliance on human priors (Mu et al. 2021; Hernandez et al., 2022; Bills et al., 2023; Conmy et al., 2023; Rajaram et al., 2024). Our proposed methods work well with these methods.

For example, an auto-interp pipeline produces a natural language description of a feature subspace, which can be translated into a high-level causal model $H$: input → concept → output. We can then (i) verify whether $H$ is faithful via interchange interventions, and (ii) use $H$ to perform the counterfactual simulation defined in Eq (7).

---

Q4. Reasons for why value probing underperform in many settings

We hypothesize two reasons:

• Incomplete coverage of the causal pathway. Unlike Counterfactual Simulation, value probing does not cover the full causal mechanism from X to Y and thus might fail to detect errors in the $X_T→Y$ computation.
• Complexity of decision boundaries. Linear probe requires not only that the variable can be encoded in a linear subspace, but also that the individual values of variables be linearly separable (L222-225)--a stronger assumption than DAS’s (Eq. 5). Non-linearity and high-dimensional variables generally make the geometry of the decision boundaries more complex and harder to learn. A great example is the country variable in RAVEL, which has over 200 unique values. Although the most frequent countries are linearly separable (as visualized via PCA), the long tail distribution makes learning the complete decision boundaries challenging. This likely explains Value Probing’s underperformance on RAVEL.

We thank the reviewer for their detailed feedback! We are glad they found our proposed methods “well formalized” and “soundly incorporate the notion of causality” and appreciated our task design, metrics, and supporting results.

---

Q1. Experimental Details

We provide detailed experimental setup below to further assure reviewers that our baseline comparisons are fair.

• Datasets
  • For each task, we randomly sample 3 folds from all prompts (i.e. the Cartesian product between templates and variable values). Each fold has 2048/1024/1024 examples for train/val/test sets (1024/512/512 for UnlearnHP due to limited source texts). Our dataset sizes are comparable to prior work on truthfulness prediction (e.g., Orgad et al 2025, Gottesman et al 2024).
• Training setup and hyperparameters (selected based on val set accuracy)
  • Confidence Score
    • Temperature scaling: Search over T = 0.5, 1, 2, 3 and report results for the best two (T = 1, 2).
    • Output token selection: Search over first-N, N ∈ [1, 100] and answer tokens. Answer tokens are identified via regex.
    • Aggregation: Experiment with mean (Kadavath et al., 2022; Guerreiro et al., 2023) and min (Varshney et al. 2023), reporting mean as it outperforms min on most tasks.
  • Probing:
    • Model: Experiment with sklearn.linear_layer.LogisticRegression with default settings (follow Orgad et al. 2025) and sklearn.svm.LinearSVC (which has lower accuracy)
    • Feature location: Search across all layers and tokens.
  • Counterfactual simulation:
    • Training data: 10K pairs randomly sampled from 1024x1024 correct examples
    • Intervention dimension: Search over powers of 2; Use 1/256/1024/4/4 for the five tasks.
    • Intervention location: Use locations identified in prior work or search over variable tokens, their immediate successors, chat template tokens, and the last token across layers.
    • Optimizer: AdamW with constant LR = 1e-4, no weight decay; trained for one epoch.

Q2. Linearity Assumptions on Representations

2.1 Presentation of assumptions

We fully agree and have made the linearity assumption for value probing explicit in L222-228. We also take this opportunity to clarify the assumptions behind counterfactual simulation below.

2.2 Counterfactual simulation does not require linearity

We would like to clarify that Eq. (7) does not rely on any assumption of linearity in model representations. The linearity assumption appears only in Eq. (5), which defines one possible localization method. Crucially, our formulation in Eq. (7) is designed to be general and agnostic to the choice of localization methods. It remains valid even when using localization methods that operate in non-linear spaces.

2.3 Causal abstraction assumes access to the full causal structure, and is not subject to the identifiability concerns raised in [3–5]

The identifiability problems in causal settings arise when some fundamental aspects of the causal structure is unknown, as in the important literature references by the reviewer [3, 4, 5]. By contrast, in our setting, the relevant causal structure--the neural network itself--is assumed to be fully accessible to use. The task of DAS, and related techniques for finding causal abstractions, is simply to check whether a "high-level" causal structure is implicitly implemented in the network. This is not an inference or identification problem, but rather a search problem whose goal is to help us understand how the model represents examples and makes predictions. We will clarify this important distinction at the beginning of the paper.

2.4 Generalizability to non-linear representations: Transformer representations are not fully linear, yet localization methods like DAS can still have partial success

We thank the reviewer for raising this point. We agree that the linear subspace approach in Eq. (5) has limitations when high-level concepts are non-linearly encoded. In this case, we can switch to a more suitable localization method (e.g., involving non-linear mappings) and our method in Eq. (7) remains compatible.

If we stick with current linear methods like DAS, generalizability becomes an empirical question. Since Transformer-based LLM representations are neither fully linear nor fully non-linear (Park et al. 2024; Smith 2024; Engels et al. 2024), our results already offer empirical insights—we expect a partial success rather than a complete failure with non-linear representations.

Q3. Generalizability to more complex causal structures

3.1 Expressiveness Please see our response to reviewer uRDA Q2.4.

3.2 Variable identifiability Please see our response to Q2.3

Q4. Presentation

We sincerely thank the reviewer for the suggestions on engaging the broader audience who are less familiar with the interpretability literature. We will add a preliminary on localization methods and evaluation metrics.