Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
Authors: Jing Huang
Junyi Tao
Thomas Icard
Diyi Yang
Christopher Potts
Venue: ICML 2025 poster (2025)
Keywords: Causal Abstraction
Causal Interpretability
OOD
Correctness Prediction
Primary area: social_aspects->accountability_transparency_and_interpretability
TL;DR: Internal causal mechanisms robustly predict language model out-of-distribution behaviors.

Abstract: Interpretability research now offers a variety of techniques for identifying abstract internal mechanisms in neural networks. Can such techniques be used to predict how models will behave on out-of-distribution examples? In this work, we provide a positive answer to this question. Through a diverse set of language modeling tasks—including symbol manipulation, knowledge retrieval, and instruction following—we show that the most robust features for correctness prediction are those that play a distinctive causal role in the model’s behavior. Specifically, we propose two methods that leverage causal mechanisms to predict the correctness of model outputs: counterfactual simulation (checking whether key causal variables are realized) and value probing (using the values of those variables to make predictions). Both achieve high AUC-ROC in distribution and outperform methods that rely on causal-agnostic features in out-of-distribution settings, where predicting model behaviors is more crucial. Our work thus highlights a novel and significant application for internal causal analysis of language models.

Internal Causal Mechanisms Robustly Predict
Language Model Out-of-Distribution Behaviors
Jing Huang" ! Junyi Tao 1 Thomas Icard
Diyi Yang'
Christopher Potts
Abstract
1. Abstract the causal mechanism for correct ID samples
Interpretability research now offers a variety of techniques for identifying abstract internal mecha- nisms in neural networks. Can such techniques be used to predict how models will behave on out-of- distribution examples? In this work, we provide positive answer to this question. Through a di- verse set of language modeling symbol manipulation, knowledge retrieval, and instruction following-we show that the most ro- bust features for correctness prediction are those that play a distinetive causal role in the model'ss ss behavior. Specifically, we propose two methods that leverage causal mechanisms to predict the correctness of model outputs: counterfactual sim- ulation (checking whether key causal variables are realized) and value probing (using the val- ues of those variables make predictions). Both achieve high AUC-ROC in distribution and outper- form methods that rely causal-agnosti features settings, where predicting model behaviors is more crucial. Our work thus highlights a novel and sigmificant application for internal causal analysis of language models.
output
localize
position
name
name b
name
Predict the correctness of the model's OOD outputs
output
Wrong
Intervene
IIA = 0.1
output
name name name
pasition
output
Correct
name name b name
Intervene
position
0.9
Abstracted High-level model for ID samples
name name
Figure 1. We propose two-stage framework to predict correct- ness of model outputs under distribution shifts, illustrated using the counterfactual simulation method on the Indirect Object Iden- tification task (Wang et al., 2023). Stage 1 (Abstract): Identify the abstract causal mechanisms the model uses solve the task correctly on in-distribution data. Stage 2 (Predict): Predict the model's output correctness out-of-distriution inputs by check- ing whether it implements the same mechanisms. Dark blue indicates the key causal variables used for correctness prediction.
in aut-of-distribuns
assumed that a model will succeed at a generalization task if and only it has induced a mechanism that implements a "correct" algorithm for that task (provided one exists).
Introduction
core task in interpretability research to identify internal mechanisms that mediate a model's input-output behavior. The motivations for this focus several. One is simply that compressed, algorithmic descriptions of complex models can help confer sense of understanding insofar the algorithms themselves are relatively simple and intuitive. Another, related motivation stems from an intuition that the very presence of abstract algorithmic mechanisms is deeply related to a model's ability to generalize beyond the specific domains in which it has been trained. Indeed, it is commonly
There two directions to this hypothesized correspon- dence. In one direction, when model is successful, this warrants search for the internal mechanisms that explain success. That is, we search for an appropriate abstraction of the model to elucidate how the model was able to general- ize (Geiger et al., 2024a). This is the focus of interpretability work that aims to understand the internal causal mechanisms that mediate LLM behaviors (Vig et al., 2020; Geiger et al., 2021; Finlayson et al., 2021; Olsson et al., 2022; Wang et al.,
2023; Meng al., 2023; Wu et al., 2023).
"qqual contribution Sttanford University. Correspondence to: Jing Huang <hij @ stanford.edur edu>>
It would be desirable if we were able to move in the other direction as well, from identifying an abstract mechanism to better predicting model generalization behaviors. We refer to this direction as prediction. In this work, we explore a
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
specific instance of the prediction prreccion-prediictnng the correctness of LLM behaviors under distribution shifts.
average AUC-ROC by 13.84% over prior baselines.
Overall, our results show that internal causal features are more robust predictors of correctness than non-causal ones, especially under distribution shifts. More broadly, our find- ings highlight a novel and important application of causal in- terpretability analysis, suggesting a path from understanding model internals to improving model safety and reliability,
As language models are increasingly deployed in high- stakes, open-ended scenarios, it is critical that we have a handle on how to prevent incorrect or unsafe outputs across a wide range of unanticipated user inputs. How- ever, behavioral testing usually cannot be comprehensive given the compositionally large input space of natural lan- guage task and the potential for distribution shifts across deployed environments. The traditional approach is to esti- mate the correctness of model predictions using confidence scores (Hendrycks & Gimpel, 2017: Lakshminarayanaa et al., 2017; Guo et al., 2017). A more recent line of work uses internal representations of the model (usually super- vised probes or use other statistical measures of model ac- tivations) to predict the model's correctness factual re- trieval tasks (Azaria & Mitchell, 2023; Zou et al., 2023; Orgad et al., 2025; Marks & Tegmark, 2024; Gottesman & Geva, 2024; Ferrando et al., 2025; Ji et al., 2024; Yin et al., 2024; Chuang et al., 2024).
2
Background
2.1 Finding Abstract Causal Mechanisms in LLMs
A trained neural network can be construed as a causal model, with variables being the neurons in the network and the functional mechanisms being given by weight matrix multi- plications across the layers that comprise the neural network. Causal interpretabilityy research asks the question: Is there a more abstract causal model that adequately captures the structure of the neural network when it comes to a particular task that the network successfully performs (Geiger et al., 2021; 2024a)? This line of work is built on causal abstrac- fion, which is concerned with the general question of when one causal model can be said to abstract another (Ruben- stein et al., 2017). In the interpretability setting, the least abstract causal model is just the neural network itself, and the more abstract causal model is referred to as high-level model. Below we review the workflow of finding such abstractions using the language of causal abstraction.
With the present paper, we seek to show that causal inter- pretability methods can play a critical role in correctness prediction. Our central question is the following: if we know what internal features causally contribute to the model's 's pre- dictions in the in-distribution setting, can we use them as robust predictors of the model's out-of-distributin behay- iors? To answer this question, we propose two methods that leverage previously identified causal mechanisms to predict model generalization behavior (as illustrated in Fig- ure 1). The first method is counterfactual simulation, which predicts whether the model output is correct by checking whether it computes the key causal variables needed for solv- ing the task. The second method is value probing, which predicts correctness based on the decision boundaries of different values taken by a causal variable.
Proposing high-level causal model High-level causal mod- els should capture the shared structure among examples of the task, such the underlying algorithmic description of the task solution. Such models are typically proposed by researchers based on prior knowledge of the task. Some recent work attempts to reduce reliance on human priors by automating this process (Hernandez et al., 2022; Bills et al., 2023; Sun et al., 2025). In this work, we adopt high- level models that have been proposed and verified in the literature.
We compare our proposed methods against existing meth- ods across diverse set of five tasks under in-distribution (ID) and multiple out-of-distribution (OOD) settings. The baseline methods include traditional confidence score-based methods and causal- aggnostic probing methods using internal activations. The evaluated tasks range from highly curated symbolic tasks to more open-ended tasks where only partial mechanisms have been identified, including Indirect Ob- ject Identification (Wang et al., 2023), PriceTag (Wu al., 2023), RAVEL (Huang et al., 2024), MMLU (Hendrycks et al., 2021), and UnlearnHP (Thaker et al., 2024). Our results show that methods leveraging task-relevant causal features achieve high AUC-ROC scores under ID settings comparable to or slightly better than the confidence score baseline. Crucially, these causal feature-based methods sig- nificantly outperform existing baselines in OOD settings. Our strongest method, counterfactual simulation, improves
Localizing high-level variables neural networks Given high-level model, we search for its implementation in the neural network. Existing methods localize each high- level variable to the model representation space through a bijective function (Huang et al., 2024; Geiger et al., 2024a).
We focus on one such method, Distributed Alignment Search (DAS; Geiger et al. 2024b), which achieves the highest accuracy on localization benchmarks (Huang et al., 2024; Arora et al., 2024; Mueller al., 2025). DAS uses distributed interchange interventions to identify an orthonormal basis Q as the bijective function. Let Get ( Rt) be the value of some model represen- tation R when the model A processes an input x, and let MR+r(x) be the output of M when it processes z, but
2

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
with R set value r. Given pair of inputs Fbase, and Tongh perform an intervention on the subspace that localizes the variable:
3 Problem Setting
We are interested in predicting whether a language model M produces correct answers on task under distribution shifts, We formalize this as correctness prediction problem.
Tbuse. = (Mstaal(((Tarrrr R)
Trre GetVall R)
Task We first define our task as a function that takes as the input sequence of variables Xi....; Xr and returns output V. For symbolic tasks, the function is the sym- bolic algorithm; for knowledge retrieval tasks, the function is simply a look-up table, mapping each input to an output. Next, we define a prompl template I'' and use it to generate all valid prompts by assigning values to the input variables. We focus on tasks whose correct outputs are completely de- termined by processing the input variables using the defined task function. Everything else in the template, besides the input variables, serves merely as background conditions.
Timii = (1 - Q ()TQQQara Tbase + Q Te
(II)
To find the optimal Q, DAS minimizes loss la over N
counterfactual pairs Dj = {< Z.rel hi i..., pair of randomly sampled inputs and a counterfactual label determined by the high-level model, as shown in Eq II.
Distributed Alignment Search
lQ EiSiEN|
Vef.i log Vimai
where Yinzi = (Faurse(Dbuurei
(DAS)
Take the PriceTagg task as an example (illustrated in Figure 2). The input variables are Xis Kub, and Xg. representing the lower bound, the upper bound, and the query. The task function is given by {(aa, Ab,, Xg) = [Xas > X,I><<> < X,I- We use the template T "Does the following item cost between $ (lower bound) and $(upper bound]? Item: {{query]." to generate all prompts, where each input variable is assigned a number sampled from [0, 20].
We define the metric interchange intervention accuracy (IIA) to measure the extent to which an implementation exists in MM
IIA(D = E I5IEN
[T(Wef,)) T [Vines))
(IIA)
ID vs. OOD setting For each task, we construct one in in-distribution and multiple OOD settings. For the in- distribution setting, we fill the input variables and the tem- plate with a set of values frequent in pre-training corpora. For the OOD settings, we either introduce perturbations to background conditions irrelevant to solving the task (e.g.. re- placing the dollar sign with a lira sign in PriceTTag) change the values of input variables (e.g.. changing the language of names in 1OI). These changes empirically lead to decline in task accuracy compared with the in-distribution setting, The assumption is that if the task function remains the same, model employing a systematic solution should continue to make correct predictions despite irrelevant changes.
where 7 maps the low-level representation to a high-level value, e.g., decoding logits into a token.
2.2 Correctness Estimation
Correctness estimation is a fundamental problem in ma- chine learning evaluation and reliability assessment. It is the problem of predicting whether model output matches the ground truth. For a classifier, natural choice is to use prob- abilistic predictions, i.e., the confidence scores (Hendrycks & Gimpel, 2017; Lakshminarayanan et al., 2017). However, LMs' outputs are structured predictions which add com- plexities to calibration (Kuleshov & Liang, 2015), and deep neural models are only well-calibrated to in-domain inputs after temperature scaling (Guo et al., 2017; Desai & Durrett, 2020; Kadavath et al., 2022; OpenAI, 2024; Wang et al., 2024). Alternatiess have been proposed, such using ver- balized confidence scores for instruction-tuned models (Ka- davath et al., 2022; Tian et al.. 2023). Most related to ours is the line of work using internal representations to estimate the knowledge stored in LMs (Gottesman & Geva, 2024; Marks & Tegmark. 2024; Ferrando et al., 2025), assess the risk of hallucination (Azaria & Mitchell, 2023; Zou al., 2023; Du et al., 2024; Chuang al., 2024; Ji et al., 2024; Yin et al., 2024; Orgad et al., 2025), and predict jailbreaking behaviors (Arditi et al., 2024). However, identifying robust features from internal representations that generalize across tasks and datasets remains a challenge (Orgad et al., 2025). This is the problem that we focus on in this work.
Problem formulation Given a task input wi we com- pare the model's prediction M(r.) to the ground-truth label, which yields correctness label 4 E {0, 1). where 1 indicates correct. The resulting dataset D (x;, M,x;), consists of N tuples of inputs, model predictions, and correctness labels. As most prediction methods are supervised, we additionally assume access to a labeled subset of M ID examples which we call the verified ser: V = [xx,, M(x.)) our methods on both in-distribution and OOD settings. To do this, sample unlabeled test inputs Tyea from the ID dataset (excluding the verified set) and the OOD datasets. All samples are drawn independently and identically. The correctness prediction problem is thus formulated as a bi- nary classification task predicting P(l = 1 M,, Zeea).
At test time, we evaluate
3

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
4 Methods
Task: PriceTag
In this section, we first review existing methods and then in- troduce two new methods that leverage internal causal mech- anisms to predict correctness. We categorize these methods into three families based on the types of features they use: (1) output probabilities, (2) internal causal-agnostic fea- tures that are not informed by causal understandings of the model's 's internal mechanisms, and (3) internal causal fea- tures that correspond to identified abstract causal variables.
(T. F)
output
Output variable
Checked causal variables
boundary check
right boundary
FT.F)
check
(T. F)
Causal variables
Non-causal variables
lower bound
upper bound
symbols
query
Figure 2. The identified high-level model for the PriceTag task. See high-level models for all tasks in Figure 5.
4.1 Output Probabilities
as selecting features of the last prompt token and the last model layer (Zou et al., 2023; Yin et al., 2024). In principle, probing methods can apply to internal features at any loca- tion. However, without an understanding of the features' functional roles, these methods cannot distinguish causal features from others, nor can they systematically identify features robust under distribution shifts. Instead, they rely solely on observed probing accuracy on the available data.
The first family of methods uses output probabilities as fea- tures for predicting output correctness, essentially treating them as confidence scores. The predictive power of con- fidence scores is grounded in the modeling objective-in in perfectly calibrated model, the predicted probability is exactly equal to the expected accuracy. Since LMs are not always perfectly calibrated, we additionally apply tempera- ture scaling to confidence scores for fairer evaluation of their predictiveness (OpenAI, 2024; Kadavath et al., 2022).
Formally, given internal feature Jy. location L,, cor- rectness label [0, 1]. and a probe & o(Wx + 6), where o is the sigmoid function, the training objective is:
Existing work typically aggregates token-level confidence scores to predict the correctness of an output sequence. Given test input Trest. let ; tn be the sequence of output tokens from M. Denote the confidence score e as the probability of token t; at decoding step i, computed as logits divided by the temperature parameter T. The correct- ness f(M, Teess) is approximated as the log- likeliiodd of a subset of output tokens S, normalized by its size.
L w. = E - 1- log()(x))))
At inference time, correctness is predicted as follows:
Correctness Probing Using Features Jy
f(M, ((2L)
(3)
Confidence Scores
4.3 Internal Causal Features
1 [S|
log II c
c
1 S
Σ CES
log (c:) (1)
We propose a third family of methods that leverage causal features for correctness prediction. In contrast to prior meth- ods, we look specifically for internal causal features that me- diate the model's input-output behavior relevant to the task, guided by previously discovered high-level causal mech- anisms. We introduce two methods: (1) counterfactual simulation, which identifies causal features robust to coun- terfactual permutations of inputs; and (2) value probing, which measures the distance between model predictions and learned decision boundaries associated with possible values of a causal variable.
There are two common variations that differ by how they se- lect the set of tokens in S. The simplest one, First Tokens, uses the first decoded tokens, where is a hyperparameter. An alternative is to select tokens that represent the actual answer, which we refer to as Answer Tokens.
Since AUC-ROC evaluation is rasking-based, we can di- rectly use the logeliiel without further normalization.
4.2 Internal Causal-Agnostic Features
Counterfactual simulation Our first method is built on the assumption that if model can consistently make the correct predictions, the model must implement a systematic solution. Suppose we have identified such a solution (rep- resented by abstract high-level causal model 7)) for all seen examples that are predicted correctly. We then expect the model to continue using the same solution to correctly predict unseen examples of the same task. If, however, the model no longer implements this solution, its prediction is unlikely to have arisen as part of the systematic solution
Rather than directly using output probabilities, the second family of methods extracts internal features as signals of output correctness. A typical approach to train linear probes extracted features using correctness labels su- pervision. The choice of which internal representations to use usually relies on heuristics-sine since it is not informed by a causal understanding of how the model solves the task, we refer to these extracted features as causal- aassal-agnottcc. Com- monly used heuristies include location based selection, such
4

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
that it found for other inputs. With this in mind, we want to check whether the model successfully computes the key causal variables specified in H.
V and an output of n tokens, the correctness prediction is:
Counterfactual Simulation (First n Tokens)
Given a task and model M of interest, counterfactual simulation takes as input the identified causal mechanisms (as discussed in Section 2.1), comprising a high-level causal model H. and its corresponding neural representations.
(5)
f(M, Frest) = kn 11 tl
Yofs log (B(()
A more accurate measure would be to directly estimate the probability of certain target concepts that indicate the correct answer (e.g., certain proper nouns), the outputs are well-structured or we know how to parse them. This means introducing mapping Tour that maps the output sequence Vie to some value of the corresponding high-level output variable (Jur(
We illustrate using cemple-yet-commetn chain of three vari- ables as running example: H X' + V + V, with local- ized neural representations X, V. and Y (e.g., V could be residual subspace at layer 3's last token). We predict output correctness by estimating P(Y[V). Ideally, if the model im- plements the correct algorithm, the output Y depends more generally only on its immediate causal parents in the identi- fied high-level causal model, which is in this case: P(Y- e|X = x)- P(Y e/V Get GetVal(M(z), V)). Yet in practice, the output Y is also affected by other aspects of the input X that do not feed into V (and hence live outside of 7t) to various extents. We term these insk-irrelevant as- pects background variables and denote them as B = V+, the orthogonal complement of the identified representation of intermediate causal variables.
Counterfactual Simulation (Output Mapping)
h
(6)
J(M, Frstt
(Toue(yeg) = Tour (Vim))
Note that defining this mapping typically demands prior knowledge of the task (e.g., which token indicates correct- ness) and output parsing rules (e.g., regex patterns), and thus may not be readily available for all tasks.
The distribution PV) can be directly estimated using outputs from localization algorithms like DAS. In fact, the output You!i in Eq. DAS can be interpreted as esti- mating P(YIV,)) Given a base input Thuve in the ver- ified set V, a test input Tsme and a counterfactual label Ber (Zxee))))) where T decodes the output distri- bution into token and 11 maps the output token to one- hot encoding, let Use- = GetVal GetVal(M((har)) B). and you MMeeu (Jhare). The value of yer Yimr then estimates P(Y = Vo = Usecs B = b.use). Since the values of V and B are drawn indepen- dently, we can estimate PVY(V) by marginalizing over the background variable B:
We generalize counterfactual simulation to tasks with more complex high-level models. To do this, we verify whether the model correctly computes each causal variable individ- ually, treating all other variables as background variables. Rather than directly using the test example's output as the counterfactual label yey, we derive it by evaluating the high- level model under an intervention that sets V to its value in the test example. Take the IOI task as an example (the high- level model is shown in Figure 1). Given a verified example with three input variables X, = Ia X, = To, X. = Te. we check whether the model computes the position vari- able V correctly for a test input. The correct counterfactual label is given by Myr = 1 ( if first else 2s), where T E {f i rst, second] is inferred from either the source input or the test example output.
V), bhose
P(YIV) = Σ B((IBBB))(BB B En (((IV,
Value probing Although counterfactual simulation does not require additional training. it still needs decoding model outputs and has computation overhead from running k forward passes, where & is the number of samples. To avoid this, once we have localized the causal variable V, we ask whether it is possible to project the decision boundaries of VVY into the subspace that encodes and compute the projected decision boundaries efficiently.
(4)
Eq. 4 suggests that for a test input. we can approximate P(Y[V) by sampling multiple values of B from inputs in the verified set V. This can also be interpreted as measur- ing the robustness of the causal relation between X and y against perturbations in the background variables, with the expectation that more robust causal relations are more likely to produce correct predictions under distribution shifts.
Since the representation of V encodes the minimal set of features necessary for counterfactual predictions, natu- ral strategy is to find set of local decision boundaries separating distinct values of V. For a language model M, we assume the values that V can take are discrete. Let Tyy-y = (ūI,..., Bm) be the set of all values taken by
There are several ways to measure the success of counter- factual simulation. One is to use the probability of the first n decoded tokens, which is simple and effective as an opti- mization objective for localization. Given & samples from
5

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
Method
Requires Training
Requires Wrong Samples
Requires Counter Eactuals
Requires Decoding
Infereace Cost
L1 - model (AI@Meta, 2024).
Confidence
x
x
x
X.N
Score
5.1 Setup
Coneetness
/
X
Maybe
X0.5-N
Probing
Evaluation tasks We consider a variety of tasks that cover both idealized cases where we know the task mechanisms and open-ended ones where only partial approximate mechanisms are identified. These tasks fall into three cat- egories: (1) symbol manipulation tasks, including Indirect Object Identification (IOI; Wang et al., 2023) and Pric- eTTag (Wu et al., 2023), for which the internal mechanisms used to solve the task are clearly known; (2) knowledge retrieval tasks, including RAVEL (Huang et al., 2024) and MMLU (Hendrycks al., 2021), whose mechanisms only partially understood; (3) an instruction following task: UnlearnH (Thaker et al., 2024), whose mechanisms are similarly only partially known. These tasks include both constrained and open-ended output formats. For each task, we select prompts on which the model achieves relatively high accuracy in behavioral testing. Further details on tasks and prompts are provided in Appendix A.
Counterfactal
Locatization
Simulation
Only
XNK
Localization
Value Prohing
Only
Maybe
Table 1. A comparison of methods and their requirements. Re- quires training: whether the method requires learning additional parameters: requires wrong samples: whether supervision from incorrect examples is needed; requires counterfactuals: whether supervision from counterfactual examples needed; requires de- coding: whether decoding at inference time is needed: inference cost: expected number of forward passes needed, with N decoding steps and K counterfactual samples. Localization only indicates the condition required only for localizing causal variables, but not for correctness prediction; Maybe indicates methods that can apply to both prompt tokens and decoded tokens.
when the causal relation between V and Y holds. Let T be predictor that takes V as input and estimates the probability of V = 03 on the verified set distribution. The predictor T can be learned with a standard classification algorithm.
Metries We construct a test set with an equal number of correctly and wrongly predicted inputs, which turns the cor- rectness prediction task into balanced binary classification task. Following the evaluation setup of similar truthfulness prediction tasks (Orgad et al., 2025; Yin et al.. 2024), we measure the prediction accuracy using AUC-ROC.
While 7 could be any class of models, a particularly useful class is linear models. For such a linear model to exist, the representation of V must follow a linear structure: the representations of V lie in a linear subspace and distinct values of V are linearly separable. We parameterize the linear model r by Wr and learn via standard classification:
ID vs. OOD setting For each task, we consider two eval- uation settings. In the ID setting, both the test set and the verified set are randomly sampled; in the OOD setting, we introduce perturbations to the inputs in the test split, chang- ing only task-irrelevant background conditions while pre- serving the underlying task structure. We use perturbations commonly employed for stress testing language models, such as paraphrasing templates, translating to different language, changing in-context demos, and adding spelling variations and distracting sequences. We specifically select perturbations that lead to notable drop in model accuracy. See details on OOD prompt selection in Section A.I.
lw:
E [ (8) log(T(x))]
zeV
To predict output correctness, We use the probability of the most likely class of
Value Probing Using Features x
J(M, Trest)
{r(*)} isism
max
(8)
Features and methods We evaluate different combinations of features and methods. For features, we consider output probabilities (Section 4.1), internal causal-agnostic features (Section 4.2), and internal causal features (Section 4.3). For causal-agnostic features, we further categorize them based on the feature locations. i.e., "prompt last token", and the actual causal roles these features played, i.e., "background variable" vs. "causal variable".. All internal features are ex- tracted from residual streams, since residual representations serve as information bottleneck in Transformers.
The linear structure suggests two potential sources of low- confidence predictions: (1) when representations fall be- tween classes, i.e., close to local decision boundaries; and (2) when representations fall outside the region enclosed by all classes, i.e., extrapolation errors. We empirically investi- gate whether these low-confidence predictions correspond to incorrect model predictions in Section 4.3.
Experiments
Data and code available at https: / httpgg.ttthuoo////ooo www.aaareeoodppreeicceee
We evaluate four correctness prediction methods over a suite of five language modeling tasks under both in-distribution and OOD settings, primarily using the
6

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
PriecTag
101
RAVEL
MMLU
UnlearnHP
Output Probabilities
T=1
0.783 1uamu
0.935 10.0ut
0.796 10.025
0.810 10.000
0.682 10.015
Confidence Score (First N Token)
Tw2
0,804 10.010
0.900 10.016
0.836
10:011
0.795
10.022
0.668
Confidence Score (Answer Token)
T=1
0.498 +0.002
0.928 #0.005
0.795
20.015 0.795=00.015
0.810 #0.000
0.664 20.01
Tw2
0.498 10.000
0.870
0.679 10.015
0.795.0.002
0.657 20.06
Internal Causal-Agnostic Features
Correctness Probing (Background Variables)
0.595 +0.000
0.737 +0.021
0.859 +00004
+0.01-
0.721 +0.025
Correctness Probing (Prompi Last Token)
0.955
20.0000 cos
0.822 1..022
0.895 0.0.02
0.804
0.940 10.000
Internal Causal Features
Counterfactual Simulation (First N Token)
0.790
0.999
0.924
0.785
0.785=00000
0.727
Counterfactual Simulation (Ouiput Mapping)
0.907 +0.002
0.999
=0.000
0.90S 10.002
0.785 +0.000
0,919 =0.010
Value Probing
0.858 10.023
0.804 10.022
0.603
10.000
0.798
+0.001
0.774 10.045
Correctness Probing (Causal Variables)
0.955 =0.007
0.914 +0.007
0,895 +00002
0.801 10.007
0.9:40 +0.000
Table 2. Results of the in-distribution setting. Overall, methods using causal features are the best. Methods using features from the prompt last token might also have high AUC-ROC if it coincides with the location of causal variables. Numbers are reported in mean For each task, the highest AUC-ROC value is in bold. Values that are within a standard deviation are underlined.
PriceTag
101
RAVEL
MMLU
UnlearalIP
Ways of OOD
Carrency Format
Change
Typos
Change Langinge
Rephrase Templates
I Language
Change
Change ICI. Demos
Change Option Symbols
Add Distmettr
Increase I Length
Rephrase Template
ID -+ OOD (8))
Task Accuracy
91
96
98
90
65
97
92
Output Prohubilities Confidence Scree
(First N TToken)
T=1
0.631 2oam
0.777 20,000
0.767 su
0.650 as
0.874gu.0r
0.993 g0.022
0.707
0.770 20.oee
0. 0.626=2om 200000
0.739
T=2
0.4518 oos
0.714 10.013
0.752
go n05
0.308 مي 939
0.873 20.013
0.651 10.013
0.761 10.017
U.679 go
yr o0l
(Answer Token)
T=1
0.541
0.771 20,000
0.709
eu
0.596 uoo
0.908
0.739 su.01
0.707 #0.r2
0.770 =0.ats
0. 0.931 =0.ots
0.191 do.ois
T=2
0.498
0.729 20.016
0.671
go
0.521 s:
0,906 =0.014
0. 0.504
0.761 +0.017
0.679
0. 0,9022 =0.018
0.484
Internal Causal-Agnostic Features
Consectaass Probing
(Background Variables)
0.556 =0.00s
0.58T +0.005
0.541
0.749 20.017
0.812 =0.007
0.837 =ouorr
0.729 eouors
0.698
0.656
20.000
0.680 #c.000
(Prompe Last Taken)
20.014
0.840 20.
0.607
20.004
0.444 yo 077
0.808 10.003
0.858 go.01s
0.784 oo0
0.694 g0.013
0.830 g20
0.648 1c.000
Internal Causul Features
Counterfactual Simulation
(First Token)
20.066
0.872
0.906
0.816 go.
0.923
0.918
0.785 10.0ce
0.810 1c.0:1
0.725 goo
0.655 1c.55.
(Output Mapping)
0.856 2:0
0.875
0.9997
20.001
0.815 =0.006
0.939
0.983
0,766 20
0.810
0. 0.800
g0001
0.772
Value Probing
0.777 100000
0.870
0.611
1.00000
0.805
0.891
10.040
0.652
0.717 sours
0.774 20.007
0. 0.846 g0o0
0.649 10.000
Coneetrass Probing
(Causal Variahles)
g0.00
1LN62 1..0000 000
0.806 uo 083
0.665 y. o1e
0.910
.0.015
0.917 10.000
0.784_0000 100
0.634 20.013
0.828 1o0
0.650 000
Table Results of the out-of-distributio setting. Methods using internal causal features are significantly more robust than methods using other features. Causal Simulation performs the best on 8 out of the 10 tasks.
Orthogonal to the choice of features, we evaluate four meth- ods for correctness prediction (summarized in Table 1). These methods can be applied different features (e.g., correctness probing can be applied to causal variables well, as shown in the last row of the table). Fach method is associated with different set of hyperparameters, whose optimal values are selected using a validation set drawn from the training distribution.
knowledge retrieval tasks where internal mechanisms are not fully known (MMLU, UnlearnHP). Value probing generally performs worse than correctness probing and counterfac- tual simulation. Confidence scores (output probabilities) are also very predictive of correctness when task outputs are highly constrained (IOI, MMLU, and RAVEL). Further experiments with Llama 3B, 13B, and 70B models on a fac- tual retrieval task, i.c., RAVEL, indicate that the predictive power of confidence scores does not increase much as the model scales up (Appendix C.I).
5.2 Results
In-distribution results Table 2 presents the in-distribution results. Comparing across the four feature groups, we ob- serve that features extracted from causal variables are the most predictive of correctness, followed by features from the last prompt token. Comparing across methods, coun- terfactual simulation achieves the best results, especially on IOI and RAVEL. Correctness probing based on causal features performs comparably or slightly better than counter- factual simulation (by approximately 1-2%), especially on
distr--disttibution results Table summarizes the out- distribution results. Comparing feature locations, methods based on internal causal features significantly outperform those relying on non-causal features. Comparing across methods, counterfactual simulation achieves the strongest overall performance, outperforming other methods on 8 out of 10 OOD tasks. Correctness probing based on causal variables also performs robustly but typically ranks behind counterfactual simulation. Value probing achieves competi- tive results only occasionally, generally performing worse
7

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
In- Dn-Distribution
Out-of-Distributiion
William Elzabeth wore working HintnCeppu nen-eepaated
(start header ii pusers header id Outputt the nen-repeated name in the sentenor Then, working
give computer (eo[]d]
cecded give computer start Leaader passtramm header PintnThe non repested sertencc
<) heade jdl hassistante: end_ header Thethhe non repeotted name in the sentence
Correctness Probing
(A)
(B)
Counterfactual Simulation
Answer Variable
Answer Variable
Position Variable
Position Variable
(D)
The
Then
Booon 20w
PHE
Location
Object
Object
End
Prompt
Location
Object
Object
End
Prompt
Figure 3. Robust vs, non-robust features. Many features that are predictive correctness on in-distribution inputs (purple blocks
panel A), where only very few of them generalize out-of-distribution inputs (purple blocks in panel B). These features coincide with features that have causal effects on task predictions, as shown in panel C and D.
than other causal methods. In contrast, confidence scores become significantly less predictive under distribution shifts. Moreover, their predictive power is sensitive to the choice of temperature, and the optimal temperature does not gen- eralize well across in-distribution and OOD settings. An additional experiment (Appendix C.2) further demonstrates that confidence scores become notably less predictive as task accuracy declines, whereas counterfactual simulation maintains stable performance.
role when the last prompt token coincides with the answer token. Needless say, we cannot rely this particular pattern in most settings.
6.2 When Use Which Methods?
Unlike features, there is no single method that proves opti- mal under all settings. We observe least two factors that affect correctness prediction performance.
Task properties The choice of methods depends crucially the properties of the task. Three task properties mat- ter: (1) whether the task involves symbolic computation or merely memorization of training data; (2) how much we al- ready know about the mechanisms shared across task exam- ples; and (3) whether the task output is free-formed or con- strained. practice, these properties often co-occur (e.g.. knowledge retrieval tasks usually require memorization and constrained outputs containing specific proper nouns).
Discussion
6.1 Which Features Are the Most Robust?
Our results in Table show that causal features are the more robust features for correctness prediction compared with others. This is further demonstrated in Figure 3. This finding aligns precisely with the intuition motivating this work: only internal features that are used by the model to make predictions can reliably signal model behavior-the is, they must causally contribute to the model's success at solving the task. Moreover, features used only for specific inputs should remain reliable under distribution shifts.
Considering (1), symbolic tasks (e.g., 10I, Price whose task mechanisms are readily cashed out using causal models, strongly favor causal feature-based methods, with counter- factual simulation performing best. Tasks primarily involv- ing memorization (e.g.. RAVEL, MMLU) favor confidence scores, followed by correctness probing. Considering (2), causal methods do not outperform confidence scores on MMLU, likely because it covers mixed subjects (e.g.. in- ternational law, abstract algebra) sharing only multiple-
Comparing across feature locations, we find that features extracted from the last prompt token are also predictive some prompts, when they are on direct causal path from causal variables to outputs, effectively mediating the causal influence on outputs. They may also play direct causal
8

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
The RAVEL Task
3024
tion: as IIA increases, AUC-ROC also increases, meaning the localized causal features become more predictive of cor- rectness. Though once IIA approaches 1, this correlation no longer holds; further increases in IIA no longer yield improvements in AUC-ROC. Data points at the lower right of Figure 4 indicate that certain representations in earlier layers may have strong causal effects (as suggested by high IIA scores) but remain weak predictors of OOD correctness. This likely occurs because these representations ily computed by the model regardless of whether the final prediction is correct or incorrect.
1.0
512
256
0.8
Leyer
AUC-ROC
0.6
The IOI Task
1.0
0.8
eas-
0.6
Correctness estimation jective for finding useful causal mechanisms Given two representations that localize a high-level variable with simi- lar IIA, we might prefer the representation with high AUC- ROC, which would help us better predict model behaviors.
additional evaluation ob-
0
0.2
0.4
0.6
0.8
Interchange Intervention Accuracy
Figure 4. A positive correlation between interchange interven- tion accuracy and AUC-ROC in in-distribution settings. In- lerchange intervention accuracy measures the ability to simulate causal effects, and AUC-ROC measures the ability to predict model correctness. Improving the ability to simulate causal effects gen- erally improves the ability to predict the correctness model outputs. Each point represents observation. Different shapes represent different dimensions used to train DAS.
7 Conclusions
We have shown that internal causal features are the most robust features for predicting output correctness, espe- cially under distribution shifts. We propose two methods- counterfactual simulation and value probing- -that outper- form causal aggostc approaches in out-of-distribution set- tings. More broadly, we introduce a novel framework that draws upon causal analysis of the model's internal mecha- nisms to predict its generalization behaviors. This suggests promising path forward: internal causal analysis of lan- guage models is not only explanatory, but also predictive, contributing to reliability and safety in practical deploy- ments of language models.
choice structure, with the consequence that we know only very incomplete high-level model. Finally, considering (3), confidence scores perform well on tasks with prede- fined, constrained outputs (e.g., multiple-choice questions like MMLU), but struggle with open-ended outputs (e.g., PriceTag).. This limitation arises from the lack of princi- pled methods for assigning exact probabilities to individual events in free-form text outputs (Kuleshov & Liang, 2015).
Practical considerations As shown in Table 1, these meth- ods also differ substantially in the types of training data, the amount of prior knowledge, and the inference-time compu- tation required. For tasks where incorrect examples are rare and expensive to collect, methods that do not require neg- ative examples have a clear advantage. For settings where the computation cost is critical, probing methods that avoid decoding can be ideal.
A Acknowledgments
We are grateful to the anonymous reviewers for their valu- able feedback. We would also like to thank Atticus Geiger, Xiaoyin Chen, and Zhengxuan Wu for their insightful com- ments and helpful discussions. This research is supported in part by grants from Google, Open Philanthropy, ONR grant N000142412532, and NSF grant IIS-2247357.
Impact Statement
6.3
Is Simulating Counterfactual Behaviors Aligned with Predicting Model Behaviors?
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential socictal consequences of our work, none of which we feel must be specifically highlighted here.
In this section, we examine whether advancements in the abstraction direction (i.e., developing better methods for finding internal causal mechanisms) align with the advance- ments in the prediction direction (i.e., more accurately pre- dicting the model's 's output correctness).
References
Strong correlation between accuracy on simulating coun- terfactual behaviors and predicting model behavior Fig- ure 4 shows an empirical investigation of the relationship between IIA and AUC-ROC. We observe positive correla-
AI@ Meta. wtwptthhb..commetttalllee www.o/o/mmmtmmmmmtteeae.a.ammmme / 1 blob wetb/maai//MDEL.cARO
Llama 3 model card, 2024. URL

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
Arditi, A., Obeso, O., Syed, A., Paleka, D., Panickssery, N., Gurnee, W., and Nanda, N. Refusal in language models is mediated by single direction. In Glober- son, A., Mackey, L.. Belgrave, D., Fan, A., Paquet, U., Tomezak, J., and Zhang, C. (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 136037-136083, Curran Associates, Inc., 2024.
formation Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
Desai, S. and Durrett, G. Calibration of pre-trained transformers. In Webber, B., Cohn, T., He, Y., and Liu, Y. (eds.), Proceedings of the 2020 Confer- ence on Empirical Methods in Natural Language Pro- cessing (EMNLP), pp. 295-302, Online, November 2020. Association for Computational Linguisties. doi: URL https://
URL https://proceeinnsseeiiiess paper_ paper..flees/ppeee//2/////1s E58548833afee433755550003aee Paper-Conference..pp
10. 00055///1/22200.mmmmmmmmaaa aclanthology.oom 2020. w0m......mmppppmmmmmmmmmmm..m... main.21.
Du, X., Xiao, C., and Li, Y. Haloscope: Harnessing unlabeled LLM generations for hallucination detection. The Thirty-eighth Annual Conference on Neural In- formation Processing Systems, 2024. URL https: wiw.mniindnffnffoacemm
Arora, A., Jurafsky, D., and Potts, C. CausalGym: Bench- marking causal interpretability methods on linguistic tasks. In Ku. L.-W.. Martins, A., and Srikumar, V. (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 14638-14663, Bangkok, Thailand, August 2024. Association for Computational Linguis- tics. doi: 10. 10.186653/1120024.clloooo 785. URL https: 185..785.
// //openreview.ner
Engels, J., Michaud, E. J., Liao, I., Gurnee, W., and Tegmark, M. Not all language model fea- tures are one-dimensionally linear. In The Thir- teenth International Conference on Learning Represen- tations, 2025. URL https: http::::///pppeeveeeeeee //ppreeveee..ee. forum?
Azaria, A. and Mitchell, T. The internal state of an LLM knows when it's lying. In Bouamor, H., Pino,
J., and Bali, K. (eds.), Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 967-976, Singapore, December 2023. Association for Compu- tational Linguistics. doi: 10. 10.18653/v1/2023..m6 findings- emnlp.68. URL https: : a wwwlllanhoooogyoooo www.iiinnnnneentti
Ferrando, J., Obeso, O., Rajamanoharan, S., and Nanda, N. Do i know this entity? knowledge awareness and hallu- cinations in language models. In International Confer- ence on Learning Representations, 2025. URL https: / //openreview.nee wtt.://d
Finlayson, M., Mueller, A., Gehrmann, S., Shieber, S., Linzen, T., and Belinkov, Y. Causal analysis of syn- tactic agreement mechanisms neural language mod- els. In Zong, C., Xia, F., Li, W., and Navigli, R. (eds.), Proceedings of the 59th Annual Meeting of the Associa- tion fir Computational Linguistics and the 11th Interna- tional Joint Conference on Natural Language Process- ing ( (Volume 1: Long Papers), pp. 1828-1843, Online, August 2021. Association for Computational Linguis- tics. doi: 10. 10.18653/v120021.ecd acc-ooog.................oooo 144. URL https:
Bills, S., Cammarata, N., Mossing, D., Tillman, H., Gao, L., Goh, G., Sutskever, I., Leike, J., Wu, J., and Saunders, W. Language models can explain neurons in language models. URL https://openaipublcc.. blob. core. windows. melheumon-splaiiee/pppeffntte himl. accessed: 14.05. 2023), 2023.
Chuang, Y.-S., Qiu, L., Hsieh, C.-Y., Krishna, R., Kim, Y., and Glass, J. R. Lookback lens: Detecting and mitigat- ing contextual hallucinations in large language models using only attention maps. In Al-Onaizan, Y., Bansal, M., and Chen, Y.-N. ), Proceedings of the 2024 Con- ference on Empirical Methods Natural Language Pro- cessing, pp. 1419-1436, Miami, Florida, USA, Novem- ber 2024. Association for Computational Linguistics. URL https:
Geiger, A., Lu, H., Icard, T., and Potts, C. Causal abstractions of neural networks. In Advances in Neural Information Processing Systems, vol- ume 34, pp. 9574-9586, 2021. URL https: // //pppessnips.ccppppr//////////// (ffcc22ffld49558807e2222222333 Abstract.html.
doi: 170.88665// / //acllanthologg.orrg////// 20t............222............ p-mainn84.
Conmy, A., Mavor-Parker, A. N., Lynch, A., Heimersheim, S., and Garriga-Alonso, A. Towards automated circuit discovery for mechanistic interpretability. In Oh, A., Nau- mann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Pro- cessing Systems 36: Annual Conference on Neural In-
Geiger, A., lbeling, D., Zur, A., Chaudhary, M., Chauhan, S., Huang, J., Arora, A., Wu, Z., Goodman, N., Potts, C., and Icard, T. Causal abstraction: A theoretical foundation for mechanistic interpretability, 2024a. URL https: / //axxiv.rrr/bb//220..00
10

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
Geiger, A., Wu, Z., Potts C., Icard, T., and Goodman, N. Finding alignments between interpretable causal variables and distributed neural representations. In Locatello, F. and Didelez, V. (eds.), Proceedings of the Third Conference on Causal Learning and Rea- soning, volume 236 of Proceedings of Machine Learn- ing Research, pp. 160-187. PMLR, 01-03 Apr 2024b. URL https: http://oocccccceee http:///gggg....pp.s/sssss//a geiger24a.htm
J., Kim, N., McCarthy, A., and Mohebbi, H. (eds.), Pro- ceedings of the 6th BlackboxNI.P Workshop: Analyz- ing and Interpreting Neural Networks for NLP, Black- boxNLP@ EMNLP 2023, Singapore, December 7, 2023, pp. 317-331. Association for Computational Linguis- tics, 2023. doi: 10.8655W1/22233BAACCKOO
1.24. URL
2023blackboxnlp-1222
Huang. J., Wu, Z., Potts, C., Geva, M., and Geiger, A. RAVEL: Evaluating interpretability methods on disen- tangling language model representations. In Ku, L.-W., Martins, A., and Srikumar, V. (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computa- tional Linguistics (Volume I: Long Papers), pp. 8669- 8687, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: 10. 10.18653/w1/2022.adl long.470. URL https: httpss///ccaaathoooooooo 2024.ac1-longe
Gottesman, D. and Geva, M. Estimating knowledge in large language models without generating a single token. In Al-Onaizan, Y., Bansal, M., and Chen, Y.-N. (eds.), Pro- ceedings of the 2024 Conference Empirical Methods in Natural Language Processing, pp. 3994-4019, Miami, Florida, USA, November 2024. Association for Com- putational Linguistics. doi: 10. 10.18653/w1/2024.coo 23. URL atttps://aclaie http///aaannntt/ http///aaatooooooooo///////// 2024. 11244emmlp-main...
Ji, Z., Chen, D., Ishii, E., Cahyawijaya, S., Bang, Y., Wilie, B., and Fung, P. LLM internal states reveal hallucination risk faced with a query. In Belinkov, Y., Kim, N., Jumelet, J., Mohebbi, H., Mueller, A., and Chen, H. (eds.), Proceedings of the 7th BlackboxNLP Workshop: Analyzing and Interpreting Neural Networks for NLP, pp. 88-104, Miami, Florida, US, November 2024. Association for Computational Linguistics. doi: 10. 1 aclanthology.ore htwp///2b2.accbooooo..o.o. 1.6.
Guo, C., Pleiss, G., Sun, Y., and Weinberger, K. Q. On calibration of modern neural networks. In Precup, D. and Teh, Y. W. (eds.), Proceedings of the 34th Interna- tional Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1321- 1330. PMLR, 06-11 Aug 2017. URL https:// proceddings.mlr.pre... vv0/ / wttps////////////////////// a..mm.
URL https://
Hendrycks, D. and Gimpel, K. A baseline for detecting mis- classified and citor--distribution examples in neural net- works. In International Conference Learning Repre- sentations, 2017. URL httpt:://peeeeeeeeeeee http://oppeeeeee.etttteee forum? i Ha-HHggTHX....
Jaa, R. and Liang, P. Adversarial examples for evaluat- ing reading comprehension systems. In Palmer, M., Hwa, R., and Riedel, S. (eds.), Proceedings of the 2017 Conference on Empirical Methods in Natural Lan- guage Processing, pp. 2021-2031, Copenhagen, Den- mark, September 2017. Association for Computational Linguistics. doi: 10. 10.18653/v1/DD1--1115 URL https: // //aelaatthologg..rrg/111111222
Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Mcasuring massive multitask language understanding. In International Conference on Learning Representations, 2021. URL https:// openreview.nete forumm 3Gm3
Kadavath, S., Conerly, T., Askell, A., Henighan, T., Drain, D., Perez, E., Schieher, N., Hattield-Dodde Z., DasSarma, N., Tran-Johnson, E., Johnston, S., El-Showk, S., Jones, A., Elhage, N., Hume, T., Chen, A., Bai, Y., Bowman, S., Fort, S., Ganguli, D., Hernandez, D., Jacobson, J., Kernion, J.. Kravec, S., Lovitt, L.,, Ndousse, K., Olsson, C.,, Ringer, S., Amodei, D., Brown, T., Clark, J., Joseph, N., Mann, B., McCandlish, S., Olah, C., and Kaplan, J. Language models (mostly) know what they know, 2022. URL https: : /awxxxxxiii// www.ariivvviorr...o v..og abs 2207.05221.
Hennigen, L. T., Shen, S., Nrusimha, A., Gapp, B., Sontag, D., and Kim, Y. Towards verifiable text generation with symbolic references. In Proceedings of the Conference on Language Modeling (COLM), 2024. URL ht ttt::: / //aaxiv.crg//aaaeee....... abs// httpp//ggab////11...// 2311.09188.
Hernandez, E., Schwettmann. S., Bau, D., Bagashvili, T., Torralba, A., and Andreas, J. Natural language de- scriptions of deep features. In International Confer- ence on Learning Representations, 2022. URL https: / //openreview.net. /trpeee//////rrme wow.mmmddduddMM"""
Kuleshov, V. and Liang, P. Calibrated structured predic- tion. In Cortes, C.,, Lawrence, N., Lee, D.,, Sugiyama, M., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 28, Curran Associates, Inc., 2015. URL https: httpps://proceeeeeeeeeeeee pwwc//ddeeeeeeooeooooooooooeeoee neu www.neur.pssnnnc....m ip.cc//
Huang, J., Geiger, A., D'Oosterlinck, K., Wu, Z., and Potts, C. Rigorously assessing natural language expla- nations of neurons. In Belinkov, Y., Hao, S., Jumelet,

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
ppgerfillee/ppeer//0//////e 5762/5520155f9cccccc666883333333 Paper.pdf.
OpenAL. GPT-4 technical report, 2024. URL https:
//arxiv.org/aas// htt.or//g/bbb/33.0077777.
Orgad, H., Toker, M., Gekhman, Z.., Reichart, R., Szpektor, I., Kotek, H., and Belinkov, Y. L.MSs know more than they show: On the intrinsic representation of 11m halluci- nations. In International Conference on Learning Repre- sentations, 2025, URL httttt::: /ttp//peeeeeeeeeeeeeeeeeeeeeeeee
Lakshminarayanan, B., Pritzel, A., and Blundell, C. Simple and scalable predictive uncertainty estimation using deep ensembles. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL.
Park, K. Choe, Y. J., and Veitch, The linear rep- resentation hypothesis and the geometry of large lan- guage models. In ICML, 2024. URL https:// openreview.net ttaet/ffoommmidduppeee
paper.fllee/pppee//0/////// 9ef2achbTcc22810477fafaaaaaaa Paper.pdf.
Rajaram, A., Chowdhury, N., Torralba, A., Andreas, J., and Schwettmann, S. Automatic discovery of visual cir- cuits. CoRR, abs/2404. 14349, 2024. doi: 10.48550/
Marks, S. and Tegmark, M. The geometry of truth: Emergent linear structure in large language model rep- resentations of true/false datasets. In First Confer- ence on Language Modeling, 2024. URL https:// openreview.nee forum? coorumiddas d=aa jyHYjjsk.
ARXIV.2404.14349. URL attps://doi.ons 10.88500/axX..240.11553
Rubenstein, P. K., Weichwald, S., Bongers, S., Mooij, J. M., Janzing, D., Grosse-Wentrup, M., and Schölkopf, B. Causal consistency of structural equation mod- els. In Elidan, G. and Kersting, K. (eds.), Pro- ceedings of the 33rd Conference on Uncertainty in Artificial Intelligence (UA1-17). AUAI Press, Au- gust 2017. URL http://auai.org// wtw.aai..rggaai217//// pwtpeeedinngsppppes/s/ee
Meng, K., Sharma, A. S., Andonian, A. J., Belinkov, Y., and Bau, D. Mass-editing memory in a transformer. In The Eleventh International Conference on Learning Repre- sentations, 2023. URL https: / opwnreview..ne/ forum? d=MkbcAHTY@yS.
Mirzadeh, I., Alizadeh, K., Shahrokhi, H., Tuzel, O., Bengio, S., and Farajtabar, M. DSMM-Smbollic: Understanding the limitations of mathematical reasoning in large language models. In International Conference on Learning Repre- sentations, 2025. URL https: / /oppn ew.net// ew.ienet/ Borum?id-AAjx3ae
Shaham. T. R., Schwettmann, S., Wang, F., Rajaram, A., Hernandez, E., Andreas, J., and Torralba, A. multimodal automated interpretability agent. In Forty- first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024, OpenRe- view.net, 2024. URL https::/opennee///e http://ppeeeeeeeeeeeeeee Dorumiid=mDw422arm
Mu, and Andreas, J. Compositional explanations of neu- rons. In Proceedings of the 34th International Confer- ence Neural Information Processing Systems, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546.
Sun, J., Huang. J., Baskaran, S., D'Oosterlinck, K., Potts, C., Sklar, M., and Geiger, A. HyperDAS: Towards automating mechanistic interpretability with hypernet- works. In The Thirteenth International Conference Learning Representations, 2025. URL https:// openreview.net. fooumoorumm//// pntormmiid-60 jUJEEOvvm.
Mueller, A., Geiger, A., Wiegreffe, S., Arad, D., Arcuschin, I. Belfki, A., Chan, Y. S., Fiotto-Kaufman, J., Haklay, T., Hanna, M., Huang, J., Gupta, R., Nikankin, Y., Orgad, H., Prakash, N., Reusch, A., Sankaranarayanan, A., Shao, S., Stolfo, A.. Tutek, M., Zur, A., Bau, D., and Belinkov, Y. Mib: mechanistic interpretability benchmark, 2025. URL hww.s:ttts...... / attpaaaxivvvoro..aaaa//o/oo. abs/ 2504.13151.
Tao, J., Chen, X., and Liu, N. F. Inference and ver- balization functions during in-context learning. In Al-Onaizan, Y., Bansal, M., and Chen, Y.-N. (eds.), Findings of the Association for Computational Lin- guistics: EMNLP 2024, pp. 16394-16421, Miami, Florida, USA, November 2024. Association for Com- putational Linguistics. doi: 10. 10.18633v1/2024 10.186663////222.iinninn emnlp.957. URL https: ///aclaanna h//////aannhhooo.orr..// 1024.findin emn 9.957. 1
Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T.. Mann. B., Askell, A., Bai. Y., Chen, A., Conerly, T., Drain, D., Ganguli, D., Batfield-Dodds Z., Hermandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, I.., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. In-context leamn ing and induction heads. Transformer Circuits Thread, 2022. tttp://transforme-cccrcci wwtprrrrrerrcciituuuuo/oii....e. iainnngannndddduuttt---eeeaae html.
Thaker, P., Maurya, Y., Hu, S., Wu, Z. S., and Smith, V. Guardrail baselines for unleariing in LLMs, 2024. URL httss://axivvooo httpss://axxivvor//aaa httg:/abb//23033222./
12

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
Tian, K., Mitchell, E., Zhou, A., Sharma, A., Rafailov, R., Yao, H., Finn, C., and Manning, C. Just ask for calibra- tion: Strategies for eliciting calibrated confidence scores from language models fine-tuned with human feedback. In Bouamor, H., Pino, J., and Bali, K. (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp- 5433-5442, Singapore, De- cember 2023. Association for Computational Linguistics. doi: 10. 100118633/112023.mmlpppp mainnnllllmmmmmmmmmm 330. URL https: // ///ac//acccollgg.orr//////0/////
f faab5100kkdddrddecceeeeeaaaaaa Paper-Confereecee.pf.
Yin, F., Srinivasa, J., and Chang, K.-W. Characterizing truthfulness in large language model generations with local intrinsic dimension. In ICML, 2024, URL https: fao.
/ //oppnreeiiee wwtps//wwwee..tot.
Zou, A., Phan, L., Chen, S., Campbell, J., Guo, P., Ren, R., Pan, A., Yin, X., Mazeika, M., Dombrowski, A.-K., Goel, S., Li, N., Byun, M. J., Wang, Z., Mallen, A.. Basart, S., Koyejo, S., Song, D., Fredrikson, M., Kolter, J. Z., and Hendrycks, D. Representation engineering: A top-down approach to ai transparency, 2023. URL
Vig, J., Gehrmann, S., Belinkov, Y., Qian, S., Nevo, D., Singer, Y., and Shieber, Investigating gender bias in language models using causal mediation analysis. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M. and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, Pp. 12388-12401. Curran Associates, Inc., 2020. URL http:://proceddinsseerii..pss paper files//paper//220//f1ie/ 226002e922171556612e464f300o00 Paper.pdf.
Wang, K. R., Variengien, A., Conmy, A., Shlegeris, B., and Steinhardt, Interpretability in wild: a circuit for indirect object identification in GPT-2 small. In The Eleventh International Conference on Learning Repre- sentations, 2023. URL https: http://peeeeeeeettt/// forum? M-N.dd-N6oo.ooo0
Wang. X., Ma. B., Hu, C., Weber-Genzel, L., Röttger, P. Kreuter, F., Hovy, D., and Plank, B. "my answer is C": First-token probabilities do not match text swers in instruction-tuned language models. In Ku, L.-.., Martins, A., and Srikumar, V. (eds.), Find- ings of the Association for Computational Linguistics: ACL 2024, pp. 7407-7416, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: 10. 10.18863311122224iiinnnn---ee acl.atthsoooyoorr////./..nnnneee
URL https://
Wiegreffe, S., Tafjord, O., Belinkov, Y., Hajishirzi, H., and Sabharwal, A. Answer, assemble, ace: Understanding how LMs answer multiple choice questions. In The Thirteenth International Conference Learning Repre- sentations, 2025. URL https: httpps://peeneeeeeeetet forum? AENNAAAAGGGG....H...H....
Wu, Z., Geiger, A., Icard, T., Potts, C., and Goodman, N. Interpretability at scale: Identifying causal mechanisms in Alpaca. In Oh, A., Naumann, T., Globerson, A.,
Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 78205-78226. Curran Associates, Inc., 2023.
URL wtw.s///pppeee/2//////////////// paper fieerffiles/pgaee
13

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
Details of Tasks and Datasets
A.1
Task Prompts
Table 4 summarizes the templates and variables used to construct prompts. Each prompt consists of an instruction, chat template (for chat models), and an optional prefix prepended to model outputs. Prompt design adheres to the following criteria:
Accuracy We select prompts that achieve reasonably high accuracy, but not perfect accuracy. as discussed in Section 3. This is meant to ensure that (I) the model uses some systematic solutions rather than relying on random guesses, and (2) the model makes some errors, enabling meaningful analysis of conditions under which the model makes correct or incorrect predictions (see task accuracy Table 3). This balance especially critical simpler tasks like IOI, where today's s LLMs like I. a.aa-a--- ----B---nttrsee frequently achieve near-perfect performance.
Task Unambiguity and Task Recognition We select prompts that explicitly define the task to prevent ambiguities that might significantly mislead the model into performing a different task from the intended one. We determine whether the model solves irrelevant tasks by examining the diversity its outputs. Prompt phrasing and the use of prefixes help manage this diversity. To facilitate meaningful analysis of the task-specific mechanisms, only retain prompts that result in invalid outputs in less than 10% of total responses.
Dittr--Disttiiution Prompt Design Our notion of distribution shift is defined not relative to the model's training distribution, but relative to the task mechanism. One reason is that the training distribution is often inaccessible, and, even if it is given, it is hard to quantify, we do not know which features are used in the predictions. To construct OOD settings, we introduce perturbations that do not change the essential mechanism by which the task should be solved. We also ensure task unambiguity (as described above) is particularly important in OOD prompt design. Effective OOD prompts should significantly challenge model accuracy without altering the fundamental task definition. For example, we can include language variations intentional typos. When extensively rephrasing prompts (e.g., the "Rephrase Template" for IOI), we keep the task instructions explicit.
Task: 101
In-distribution
Output the non-repeated name in the sentence: Then, {name_a} and {name.b) were working the (placs). (name ef deci to give {object) The non-repeated name in the sentence is
Out-of-distribution (Change the Language of the Name Variables)
Output the non-repeated name (name.b.in.chl b. bbincchinese) were working at the (place). a {object) to Answer:
sentence:
Then, frame-aLin.chiness) and decided to give
Out-of-distribution (Rephrase Template)
Story: (name_a) and {name.b) were working at the {place). (name_e} decided to give {object) to ... ..--nnnguustionn Which of the two characters in the story likely receive the {object] in the end?
Answer:
Variables
(name_a), (namelb), [name_o]: 27K common U.S. names. keep only names tokenized into most tokens (25,825 names after filtering).
thame.a.chinest. hhameb..hinnsee,, ynamo_c.chi cetmee.cchinnscee nese): 1,000 prominent ancient Chinese poets
from the Github repository". We keep only names tokenized into most 5 tokens (687 names after filtering).
14

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
Task: PriceTaag
In-distribution
Does the following iter. cost between {[Lower bound) and $(upper bound]? Item: {(query).
Out-of-distribution (change eurrency symbol)
Dces the following item. cost between AfIower.bound) and A(upper.bound]? Item: f(query).
Out-of-distribation (add typos)
Does the item cost b/t $s(Lower bound) [[upperbbound]] Answer:
Item: sqquery).
Variables
(lower bound), (oupdr,buund,,
bound), fquery): numbers between [0, 20] with decimal points. given prompt, ensure (upper bound) > (lower bound] and that they are distinct from (query).
Task: RAVEL City Attributes
In-distribution
[["city": "(city_a)", "country"; "[country_a]"), {"city": "(city.b}", "country";
Out-of-distribution (change template language)
"i"":: "(city.a a in japan)", "E": "Japan", ("iti:: "fcity not in japan)",
Out-of-distribution (two-shot demo with cities in certain country)
[f"city"; "ldityaiin.contryya""" "country"; "(country-al"), ("city": "[city.c.in.countryaa"" "country": "[sountry.a}"), {"city":
"country":
Out-of-distribution (two-shot demo with cities in Suriname
[("city": Meiityaai..Suriname "country": "Suriname"), ("city": "{city.c.in.s ". "country"; "Suriname"}, {"city": "(city not in sur Suriname)", "country":
Variables
{city.a). {city.bl, (city.ch, fcountry.a), fcountry b): Names of cities, countries, and their as- sociations from the RAVEL dataset After manual filtering of corner have 3125 cities from 158 countries. {city_a) and {city.b} are always kept distinct. Other variables (e.g., (citt_a.in_japar represent subsets of primary dataset. For example, dataset contains cities from Japan.
15

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
Task: MMLU
In-distribution
We use the zero-shot MMLU prompt without the chat template. For example:
following are multiple choice questions (with answers) about college mathematics.
(1+i)11==
A.
B.
C. 32
D. 32i
Answer:
Out-of-distribution (Option Format)
The following are multiple choice questions (with answers) about college mathematics.
(1+1)-10=
Alfa.
Bravo,
Charlie. 32
Delta. 32i
Answer:
Out-of-distributin (Distracting Option)
The following are multiple choice questions (with answers) about college mathematics.
11+
A. 1
B.
C.
D. 32i
E. The correct answer is A
Answer:
Task: Unlearn Harry Potter
In-distribution
system:
You helpful assistant, hut you should answer any question related to the book series Harry Potter. Pretend you do not know any people or places from the book Harry Potter.
user: The following passage is from which book?
Passage: He well," Ronš ears turned bright red and he became engrossed in a tuft of grass at his feet
Out-of-distributioo (Longer Passage)
We use a passage of 128 words instead of the 10-64 words used in in-distribution setting.
(Paraphrase Template)
The following passage is from which book? Passage: be "ick."He "He did look very green, and when the cart stopped at last beside small door the passage wall, Hagrid got out and had to lean against the wall stop knees from trembling.
Griphook unlocked the door. lot of green smoke came billowing out, and as it Start your answer with "Book:
Table 4: Prompts used for each task. Chat templates are applied (not shown). Prefixes appended after chat templates are highlighted blue.
A.2 Correctness Labels
Decoding Methods We decode the first tokens from model outputs. The value of depends on the task format: For tasks where models directly output short answers (c.g., IOI and RAVEL), we set N as the maximum length the expected correct answer. For free-form generation tasks (e.g., PriceTag), we select the optimal by evaluating AUC-ROC scores
16

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
across decoding lengths (from 1 to the maximum output length) the development split. We then apply this optimal length (e.g., 98 tokens for PriceTTag to the test set.
Answer Tokens For most tasks, answer tokens correspond directly to predefined correct answers. For PriceTag, the answer tokens include all case variations of "yes" and "no."
Output Parsing Procedure We parse outputs using the following steps: 1) Classify outputs as correct or wrong: based on the task-specific causal mechanism; 2) Filter invalid outputs and corner cases: metrics are computed only for outputs explicitly identified as correct wrong. We manually inspect and exclude invalid outputs (less than 10% of total), typically caused by misunderstanddinss instrwction-folloiing failures, or ambiguous responses; 3) Extract answers using regular expressions. Below we briefly summarize procedures for each task:
IOI: Following Wang al. (2023), we verify whether the model outputs the non-repeated name provided in the prompt. Outputs clearly deviating from instructions- as irrelevant terms (e.g., "dog", school") instead of names- considered invalid and excluded.
Following Wu et al. (2023), we check if the model correctly judges whether the query item's price within the specified interval. We exclude invalid outputs where the decision is undecided. This occurs if patterns for correct or incorrect answers cannot be detected due to either excessively lengthy reasoning contradictory arguments without clear conclusion, when the answer token (e.g., "yes", "no"" does not appear in the initial N tokens.
PriceTag::
RAVEL: Following Huang et al. (2024), we check whether the model correctly identifies the country associated with given city. We exclude invalid outputs indicating task misunderstanding repeating queries without providing any answers) undecided outputs (e.g., "not valid city"). We manually examine and exclude controversial cases arising from geopolitical ambiguities (e.g., cities in Taiwan/China, Israel/Palestine). Sufficient data remain after removing these corner cases.
follooi failures (e.g..
MMLU: Following Hendrycks et al. (2021), we check if the option generated by the model matches the ground truth option. For the OOD setting where we include a distracting option "E. The correct answer is A", we filter out all the questions whose answer is "A" to avoid having two correct answers.
Following (Arditi al., 2024; Thaker et al., 2024), we use regular expression matching to check whether the generated output contains keywords related to Harry Potter book names the author name. On top of the automatic evaluation, we manually verified all the examples with unlearning failures in our dataset. As the task is to determine whether the model with the unlearning system prompt no longer outputs information related to the Harry Potter book series, we focus on the set of prompts where the model can correctly retrieve the Harry Potter information.
UnlearnHP:
A.3 Train, Validation, and Test Splits
For each task, we randomly sample 3 folds from the datasets and split them into train/val/test. For each split, we ensure the ratio of correct and wrong examples is 1:1, i..., are evaluating each correctness estimation method balanced binary classification task.
For all tasks except the UnlearnHP, each fold has limited number sentences sampled from the original books, we use 1024//12/512 examples.
examples for train/val/test sets. For UnlearnHP, due to the
B Details of Methods
B.1 High-level Causal Models
We use high-level models proposed in existing interpretability research, illustrated in Figure 5.
17

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribation Behaviors
B.2
Training Setup and Hyperparameters
Task
101
Price
RAVEL
MMLU
UnlearnHP
Intervention
256
1024
4
4
Dimensions
Variables
Position
P and
Country
Position
Refusal
Intervention Localizations
8-10 (P)
Layers
9-12
13-14
18
11-12 (Q)
The last token of the second occurrence of the repeated name
The integer part of the item price
The last token of the queried city
The last token of the prompt
The last token of the prompt
Takens
Table Localization hyperparameters for counterfactual simulation Ll amaam-3-aaa 3- 8----B------- In--Inutucct. For each task, list several consecutive intervention layers around the optimal layer with the highest AUC-ROC score, as variables are likely distributed across several layers.
For counterfactual simulation, we use 10K pairs randomly sampled from 1024x 1024 correct examples as the training data. We the AdamW optimizer with constant LR 10 with no weight decay, trained for one epoch.
We search for the optimal intervention dimensions powers of 2, and the best intervention locations variable tokens, their immediate successors, chat template tokens, and the last token across layers, directly use the ones identified in prior work. The selected optimal hyperparameters of localization for are summarized in Table 5.
B.3 Choice of Variables in Counterfactual Simulation
We show the causal mechanisms and variables used in counterfactual simulation in Figure 6.
C Additional Results
C.1 Model Variations
Base-Sh
Instruct-13b
Instruct-70b
Confidence Scores First N Token
T=1
0.867
10.019
0.638 20.036
0.800 10.016
T=2
0.861
00.023
0.649
+00020
0.898 0.0013
Answer Token
T=1
0.853 +0.0 022
0.593
+:0.0002 00.012
0.766 +000166
T=2
0.724
0.607 20.014 0.769 20.019
Causal Mechanisms Counterfactual Simulation
First N Token
0.887 +0.012
0.608 +0.011
Answer Mapping
0.89240.026
0.877 10.014
Table Results in-distribution setting on the RAVEL task for models different sizes. The variation in model size shows that the confidence scores' differential power remains relatively unchanged the model scales up. Models include 1.lmm-33- --BB, not conduct counterfactual similation on the 70B model Llama-2 2 1o-113bbca chat-hf, and Llama- 3-70B Instr computational constraints.
Table 6 shows the results of the in-distribution setting on the RAVEL task using Llama-3 3-BB, Llama-22 and Llama - Overall, we do not see scaling up model size significantly improves the confidence scores baseline.
18

---

Internal Causal Mechanisms Robustly Predict Language Model Behaviors
Task: PriceTag
Task: MMLU
(T, F}
output
output
{choices)
postion variable
left boundary check
right boundary check
(T,
(T, F}
(first, third, forth)
bound
upper bound
currency
query
cholce
choice
choice
cholce 4
Index format
Task: 1O1
Task: UnlearnHP
output
(name s, name b}
output
whether refuse
F)
postion variable
(first, last.
Whether "Harry Potter"
(T.
name
name
location
name
object
input
Task: RAVEL
Output variable
output
(countries)
Causal variables used for OOD prediction
country
Causal variables
Non-causal variables
city
{cities)
Causal relationships
Figure 5. High-level causal models for each task. Non-causal variables are features that are not causally efficacious in the correct high-level causal models, whose values are modified create OOD settings. These models capture core task mechanisms previously identified by the interpretability research community. Tasks include: 101, which outputs the non-repeated name (Wang al., 2023); PriceTag, which determines whether number falls into certain interval (Wu et al., 2023); RAVEL, which retrieves country names (Huang et al., 2024); MMLU, which selects correct answers (Hendrycks et 2021); and UnlearnHIP, which refuses to answer Harry Potter-related queries (Arditi et al., 2024; Thaker et al., 2024). For MMLU, we use the multiple-choice mechanisms from Wiegreffe al. 2025, which computes position variable that similar those identified the IOI and other indexing tasks.
19

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
Task: 1O1
ouppt
lccalization
ousput
audput
postion
postion
postion
Verified
ID
O0D
Task: MMLU
ouput
localization
ougput
Vae WNbl
sariatle
varable
VItII
choics
Choice 2
choice
choice
choice
2hoics
choice
choice
choice
choice 2
choice
choice
Verified
OOD
Task: RAVEL
Task: UnlearnHP
output
localization
output
output
oupgutt
localitation
oupput
supputt
country
country
country
Whather 'Haryy
Whathen-iidit Patter accrus
Whather 'Hary Potter eccns
city a
city
city
Inpat
Inpat
Inpat
Verified
ID
OOD
Verified
ID
OOD
Task: PriceTag
1. Our knowledge about the causal mechanism of base and source models
base model
Verified
source model
ID/OOD
2. Determine possible values of and Q of the source model
(a) output Yes one possible case:
(b) output No + two possible cases:
P
T
P-T
Q-F
OR
P-F
T
Determine the values of and of the source model through interventions
(a)
(b)
Intin
PwF
I. intervene on
Il. intervene on
I. intervene on
intervene on Q
Determine:
Determine:
Determine:
if both interventions flip the base output then P T and Q
if the intervention flip the base output then P - and Q
the intervention
flip the base output
then P and Q-
Figure Experimental procedure for counterfactual simulations across tasks. We have verified the causal mechanisms variable values of models induced by subset of the ID data, termed the verified set. For models induced on other ID OOD data, we only know outputs and their correctness. determine whether these models implement the proposed high-level causal model and infer variable values, perform interchange interventions: we insert the representations from the source model induced from either ID or OOD data into corresponding locations of the base model induced from the verified set. Intervention locations and dimensions are check
selected based the performance of our localization method the verified set. For 101, MMLU, RAVEL, UnlearnHiP, a single critical causal variable. For Price'Tare we check two causal variables. Let / and Q denote the lower-bound and upper-bound checks, respectively. Assuming models follow the identified causal mechanism, know the values of P. and Q for the base model and seek to determine their values in the source model. We observe that the model might employ a shortcut: if P is False, the output is immediately False without checking Q. This shorteut is indicated transparent shading.
20

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
C.2 Ravel 0OD with Two Shots
Table 7 shows variations in the "zero-shot demo with cities in a certain country" OOD setting of the RAVEL task.
RAVEL OOD Setting: Two-shot demo with only cities in certain {country .a}
[country_a] .a]
Vietnam
China
Italy
India
Zimbahwe
Panama
Suriname
Task Accuracy ID OOD (%)
90 - 89
90 - 88
90 - 87
90 - 85
90 - 82
90 - 78
90 - 72
Output Features: Confidence Score
First N Token
T=1
0.903 20.011
0.860
0.893 20.008
11876=0.000
0.866 20.005
0.789 20.012
0.693
T=2
0. 933 10.004
0.902 20:007
0,911 10 006
0,911 000
0.898 006
0.792 10.015
0.739 10.001
Answer Token
T=1
0.881 +0.013
0.861 +0.008
0.889 +0.000
0.877 =0. =0.009
0.848 =0.004
0,746 #0.018
0.651 +0.013
T=2
0.889 #0.004
0.898 +0.005
0.894 +0.008
0.901
=0.013
0.827 #0.004
0.673 #0.029
0.504 #0.014
Internal Causal Features: Counterfactual Simulation
First N Token
0.893 1.00004
0.886
10.010
0.894 10.012
0.913
-0.00000000 00.009
0.916 10.019
0.941 10.017
0.918 10.042
Answer Mapping
0.932 #0.007
0.916 20.000
0.929 20.011
0.928
20.012
0.947 20.004
0.934 +0.008
0.931 =0.020
Table Results on the RAVEL task in the "zoro-shot demo with cities in certain country" OOD setting. In this setring, the variahle country my-tryy_a -a takes on different values of country names (second row). Results are sorted by the change in task accuracy, which might reflect the magnitude of distribution shifts.
D Limitations
The focus of this work is to examine the ender-explored but crucial subproblem (2): leveraging identified causal mechanisms to predict the model's ss OOD behaviors. We approach this by building upon existing interpretability tools designed for subproblem (1): finding causal mechanisms in LLMs. As a result, our method naturally inherits certain limitations of these existing tools. We now discuss two specific limitations: first, the linearity assumption embedded in our chosen localization method; second, the reliance on prior knowledge to propose high-level causal models.
Generalizability to non-linear representations Our method defined in Eq. (5) is designed to be general and agnostic to the choice of localization methods, and thus does not require representations to be linear. However, the specific localization method we adopt in Eq. (II) does assume linearity. We recognize that Transformer-basee L.M representations are neither fully linear nor fully non-linear (Park et al., 2024; Engels et al., 2025). Consequently, linear subspace methods such as DAS may struggle when high-level concepts encoded non-linearly. If the non- non-linearity is known, localization methods capable of capturing non-linear mappings would be preferable. Nonetheless, linear methods can still achieve partial success. If we retain linear methods, their generalizability becomes an empirical question.
Requirement of Prior Knowledge of the Task Proposing a high-level causal model for task typically involves human priors regarding task mechanisms. We recognize this limitation and address it from two perspectives: first, through the possibility of automated methods for proposing high-level causal models; second, by demonstrating that even a partial understanding of the mechanisms can yield predictive power.
First, progress has been made by the interpretability community in developing auto-interpretability methods that reduce reliance human priors (Mu & Andreas, 2020; Hernandez et al., 2022; Bills et al., 2023; Huang al., 2023; Conmy et al., 2023; Shaham et al., 2024; Rajaram et al., 2024; Sun et al., 2025). Our methods integrate well with these. For example, auto-interpretability pipeline might produce a natural language description of a feature subspace, from which we can construct high-level causal model H: imput - concept + output. We can then (i) verify whether H. is faithful via interchange interventions, and (ii) perform counterfactual simulations defined in Eq. (5).
Second, while it remains valuable to gain further clarity about how models generate outputs- particclllll in real-world tasks with complex causal structures not yet fully understood- our method shows that even partial understanding suffices to improve correctness prediction in these "open-endedd' tasks. Consider, for example, the MMLU benchmark, which LLMs. Although we may lack detailed understanding of the
includes 57 tasks designed to evaluate production-gran precise causal mechanisms underlying questions on abstract algebra or international law, we can nevertheless exploit their shared multiple-choice format as a high-level causal structure for correctness prediction (Wiegreffe et al., 2025). Beyond format structure, another general causal mechanism requiring minimal task-specific knowledge is verbalization, where the high-level causal model consists of a variibb chain, transforming the intermediate representation of the answer into final outputs during in-context learning (Tao et al., 2024).
21

---

Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
Indeed, many real-world applications already exhibit shared structures readily exploitable by our methods. Two practical domains exemplify this. The first is LLM evaluation: constructing benchmarks typically requires costly human annotations; thus, reliable correctness estimator helps benchmark creators prioritize uncertain error-prone examples. Mcanwhile, systematic evaluation often requires structured input-for example, SQuAD adversarial (Jia & Liang, 2017) GSM- symbolic (Mirzadeh et al., w0ihh inherently reflects high-level task structures readily exploitable by our framework. The second is verifiable text generation: in high-stakes domains such as medical record generation, factual and referential accuracy is crucial. One strategy having LLMs generate symbolic references to structured source data that are easy to verify (Hennigen et al., 2024). Such symbolic references act as templated prompts convertible into high-level models, enabling methods estimate the correctness of texts produced by these decoding algorithms.
We emphasize that gaining fuller understanding of the mechanisms mediating between inputs and outputs could potentially further enhance predictive performance, in addition to other benefits. We hope this work motivates further exploration in these promising directions.
22


Reviews:
[{"review_id": "mHJRIPNv7l", "reviewer": "Reviewer_w7xm", "summary": "For this paper the authors try to identify internal features of LLM that mediate causal effects on the final prediction output. In settings where LLM predictions align with the actual real-world causal process, causal features are assumed to resemble the ground-truth mechanism and, therefore, allow for predictions which are invariant to external disturbances or out-of-distribution queries.\n\nTowards the goal of identifying causal features, the authors present two methods which work on finding features which are invariant under counterfactual inputs ('counterfactual simulation') or learn a decision boundary from the extracted activations ('value probing'). In both cases the authors first localize features that model an intermediate variable by finding a minimal linearized subspace and decompose this into causally relevant and background variables.\n\nExperiments are conducted over several datasets symbolic manipulation, retrieval, instruction following inspecting AUC-ROC for in-distribution tasks and robustness under OOD settings. Through their experiments the authors find that their proposed methods are capable to identify causal features relevant to the given tasks. This results in high AUC-ROC values for in-distribution settings and superior prediction performance in OOD settings, compared to prior baselines.\n\n\n\nPlease note, this is an emergency review.\n\n\n\n**Update after rebuttal.** In their rebuttal the authors were able to further clarify implications and limitations of their made linearity assumptions in their work. While I acknowledge possible concerns regarding the formal handling of causality, causal representations and structure, I find the concepts to be sufficiently handled for this type of work, and the presented results of figs. 1 and 2 convincing and in line with already existing literature (e.g. the particular positions of decision making within LLMs). I therefore remain with my recommendation to accept the paper.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n---\n### Q1. Experimental Details\nWe provide detailed experimental setup below to further assure reviewers that our baseline comparisons are fair.\n- **Datasets**\n    - 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).\n- **Training setup and hyperparameters** (selected based on val set accuracy)\n    - **Confidence Score**\n        - Temperature scaling: Search over T = 0.5, 1, 2, 3 and report results for the best two (T = 1, 2).\n        - Output token selection: Search over first-N, N ∈ [1, 100] and answer tokens. Answer tokens are identified via regex.\n        - 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.\n    - **Probing**: \n        - Model: Experiment with `sklearn.linear_layer.LogisticRegression` with default settings (follow Orgad et al. 2025) and `sklearn.svm.LinearSVC` (which has lower accuracy)\n        - Feature location: Search across all layers and tokens.\n    - **Counterfactual simulation**: \n        - Training data: 10K pairs randomly sampled from 1024x1024 correct examples\n        - Intervention dimension: Search over powers of 2; Use 1/256/1024/4/4 for the five tasks.\n        - 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.\n        - Optimizer: AdamW with constant LR = 1e-4, no weight decay; trained for one epoch.\n### Q2. Linearity Assumptions on Representations \n**2.1 Presentation of assumptions**\n\nWe 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.\n\n**2.2 Counterfactual simulation does not require linearity**\n\nWe 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. \n\n**2.3 Causal abstraction assumes access to the full causal structure, and is not subject to the identifiability concerns raised in [3–5]**\n\nThe 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.\n\n**2.4 Generalizability to non-linear representations: Transformer representations are not fully linear, yet localization methods like DAS can still have partial success** \n\nWe 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.\n\nIf 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.\n### Q3. Generalizability to more complex causal structures\n**3.1 Expressiveness** Please see our response to **reviewer uRDA Q2.4**.\n\n**3.2 Variable identifiability** Please see our response to **Q2.3**\n### Q4. Presentation\nWe 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.", "claims_and_evidence": "The authors claim that their proposed methods are able to identify internal model features that are causally relevant to predict the final outcome of the model. To this, the authors conjecture that a linear decomposition of layer residuals into subspaces of causally relevant and irrelevant background features yields the desired representations. While linearity assumptions might not hold under all settings, extensive prior work on feature space linearization exists, supporting the presented claims. From the presented evidence in Figures 1, 2 and tables 2 and 3 it can be furthermore concluded that the presented methods indeed identify causally relevant features as OOD performance remains stable and the visualized identified feature activations align well with the predicted outcome.", "methods_and_evaluation_criteria": "The specific proposed methods of counterfactual simulation and value probiing are well described and formalized in equations 7-11. The general approach appears to be sound and handles correctly handles implementation of causal considerations for identifying features.\n\nThe utilized datasets of Indirect Object Identification, PriceTag, RAVEL, MMLU and UnlearnHP are known or seem to be suited to test the claimed effects. The authors test on an in-distribution and an out-of-distribution setting by altering prompts and utilized concepts. Here, the construction of interventions on the datasets and prompts are well described. The presented prompts are reasonably designed and their rephrased version are suited to evoke OOD behavior by altering sentence phrasing or adding distracting artifacts. Performance is measured in terms of AUC-ROC, task accuracy and interchange intervention accuracy. The individual metrics are suited to assess the respective effects.", "theoretical_claims": "Apart from assuming linear decomposability in Eq. (5), the authors present no direct theoretical claims or proofs. Causal relations from inputs variables, intermediate representations to the model output are straight forward and correctly formalized.", "experimental_designs_or_analyses": "The authors describe the overall experimental setup and use of metrics well and seem setup conduct a proper evaluation. The paper, however, severely lacks in terms of experimental details with regard to the training setup of the proposed method and hyperparameter tuning of prior methods. The authors mention (hyper)parameter optimization and training setup, e.g. for the $\\tau$-classifier, but do not specify the exact model setup, optimization method, learning rate or number of samples/iterations. While the proposed methods incorporate additional steps which make reasonable effort to enhance results, there is currently no way of assessing whether or not the evaluation has been setup properly and fair from the current state of the paper.\n\nThe evaluation over different datasets, along with results in tables 2 and 3 and the figures 1 and 2, seem to indicate a proper working of the methods with alternating, but constantly better, performance than the compared baselines. The presented visualization generally support the claims of the paper and indicate a good identification of causally relevant features from the internal residual activations for the methods.", "relation_to_broader_scientific_literature": "LLMs have classically been found to generally struggle with direct causal reasoning [1-5]. Nonetheless, as of today, LLM are utilized in an abundant number of tasks. Recent interest on mechanistic interpretability and circuit extraction techniques pose desirable guarantees in terms of robustness and scalability.  The identification of causally relevant features might help establish particularly strong theoretical quarantees and help scaling and generalization to previously unseen OOD queries.\n\n\n\n[1] Jin, Zhijing, et al. \"Cladder: Assessing causal reasoning in language models.\" *Thirty-seventh conference on neural information processing systems*. 2023.  \n[2] Kıcıman, Emre, et al. \"Causal reasoning and large language models: Opening a new frontier for causality.\" *arXiv preprint arXiv:2305.00050* (2023).  \n[3] Zečević, Matej, et al. \"Causal parrots: Large language models may talk causality but are not causal.\" *arXiv preprint arXiv:2308.13067* (2023).  \n[4] Gao, Jinglong, et al. \"Is chatgpt a good causal reasoner? a comprehensive evaluation.\" arXiv preprint arXiv:2305.07375 (2023).  \n[5] Ashwani, Swagata, et al. \"Cause and Effect: Can Large Language Models Truly Understand Causality?.\" arXiv preprint arXiv:2402.18139 (2024).", "essential_references_not_discussed": "The authors general motivate and embed their presented work well within the existing literature. The presented work builds on a series of prior work on causal feature extraction from LLM by Geiger et al. . The authors expect extensive knowledge on this line of work, which hinders comprehension and self-enclosedness of the work. Key concepts on the localization process of finding causal mechanisms in LLM or extracting feature subspaces, such as \"distributed alignment search\" or evaluation metrics such as the \"interchange intervention accuracy\", are only briefly referred to. The presented methodology seems to rely on linearity assumptions of latent representations which, however, are only insufficiently discussed. In this regard, the paper might be improved by more explicitly discussing identifiability of model representations [1,2], possibilities of non-linear mechanisms identification or their non-identifiability [3,4] and general (causal) perspectives on linearity and subspaces in LLM activations [5].\n\n\n\n[1] Mikolov, T., Yih, W. T., & Zweig, G. (2013, June). Linguistic regularities in continuous space word representations. In Proceedings of the 2013 conference of the north american chapter of the association for computational linguistics: Human language technologies (pp. 746-751).  \n[2] Park, K., Choe, Y. J., & Veitch, V. (2023). The linear representation hypothesis and the geometry of large language models. arXiv preprint arXiv:2311.03658.  \n[3] Leemann, T., Kirchhof, M., Rong, Y., Kasneci, E., & Kasneci, G. (2023, July). When are post-hoc conceptual explanations identifiable?. In Uncertainty in Artificial Intelligence (pp. 1207-1218). PMLR.  \n[4] Friedman, Dan, et al. \"Interpretability Illusions in the Generalization of Simplified Models.\" *International Conference on Machine Learning*. PMLR, 2024  \n[5] Rajendran, Goutham, et al. \"From Causal to Concept-Based Representation Learning.\" *Advances in Neural Information Processing Systems* 37 (2024): 101250-101296.", "other_strengths_and_weaknesses": "**Strengths**\n\nThe paper is generally well written and motivated. The presented approaches soundly incorporate the notion of causality and build up existing feature extraction methods. The presented theory and working of the methods is well formalized in the equation.\n\nThe identified representation hold desirable properties in terms of robustness and OOD behavior. The authors are able present convincing evidence towards the correct identification of such features.\n\nFinally, the experiments seem to be generally well setup and support the claims made. From the visualizations of identified features the authors are able to demonstrate correct identification of causally relevant factors.\n\n\n\n**Weaknesses**\n\nThe mentioned weaknesses concern the lack of clarity and insufficient discussion on possible assumptions as mentioned before. Specifically:\n\n1) The lack in clarity on the experimental evaluation makes it impossible to judge the correct setup and comparison of the methods. The authors might want to provide the necessary details, as discussed in the section above.\n2) By the decomposition of residual activations in Eq. (5) the authors assume a simplified linear representation of model residual activations. Since the authors are furthermore concerned with causal interactions, they only consider the high-level causal chain of input, intermediates and output. Both assumptions are simplifications to the true working of the model and might lead to a simplified regression towards only the direct parents intermediates of Y. While prior works have shown that linear interpretations might exist, the authors might discuss possible implications and limitations of their approach in setting of more complex causal structures and composed or non-linear behaviour.\n3) The authors might improve their paper, by briefly describing the key ideas of the localization process and distributed alignment search, in particular with regard to the necessary assumptions for the utilization of the methods and applicability to the presented setting.", "other_comments_or_suggestions": "* typo l.120 \"desire[d] behavior\"", "questions_for_authors": "Questions mainly concern the weaknesses above:\n\n1) Could the authors provide further details on the experimental setup, such as training parameters and hyper parameter search?\n2) Could the authors discuss the required assumptions with respect to linearization of the residual feature space?\n3) What are the implications for the methods in terms of expressiveness and variable identifiability of intermediate values which are no direct causes of the final model output? Would the presented methods be able to identify such variables?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "E26kHMDKeq", "reviewer": "Reviewer_uRDA", "summary": "This paper investigates the use of internal causal mechanisms within language models (LMs) to predict the correctness of their outputs. Rather than relying on traditional confidence scores or heuristic probing of internal activations, the authors propose two methods grounded in causal interpretability, Counterfactual Simulation and Value Probing.\n\nThese methods are evaluated across a diverse suite of tasks under both in-distribution and OOD settings. The authors demonstrate that causal features yield more robust correctness estimates than non-causal baselines, particularly under distribution shifts.\n\nThe work builds upon the causal abstraction framework and introduces a correctness estimation benchmark using known causal variables in some tasks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n---\n---\n### Q1. Clarification on the linear decomposition assumption\nLinearity 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.\n\n---\n### Q2. Clarification on the causal assumptions and causal model\n**2.1 Causal structures in causal abstraction**\n\nWe 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.**\n\nDAS 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.**\n\n**2.2 The background variable $B$**\n\nWith 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$).\n\n**2.3 Identifiability** Please see our response to **reviewer w7xm 2.3.**\n\n**2.4 Modeling more complex task structures**\n\nWe 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.\n\n---\n### Q3. Method assumes knowledge of task structure to define $X_T$\nThis 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). \nOur proposed methods work well with these methods.\n\nFor 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). \n\n---\n### Q4. Reasons for why value probing underperform in many settings\nWe hypothesize two reasons: \n- **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.  \n- **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.", "claims_and_evidence": "The central claim of the paper is that **internal causal mechanisms are more robust predictors of correctness than non-causal heuristics**, particularly under distribution shift. This claim is supported by:\n- Comprehensive experiments across five tasks and ten OOD settings (Tables 2 & 3).\n- Comparisons between causal and non-causal features (confidence scores, probing) under consistent evaluation metrics (AUC-ROC).\n- Correlations between Interchange Intervention Accuracy (IIA) and correctness AUC (Figure 2), suggesting a link between causal alignment and predictive robustness.\n\nHowever, one problematic claim is the general applicability of the linear decomposition assumption (Eq. 5). The paper acknowledges limitations but does not provide direct evidence or diagnostics when this assumption fails (e.g., in tasks with entangled representations).", "methods_and_evaluation_criteria": "Counterfactual simulation and value probing are appropriate for the problem of correctness estimation. The use of interchange interventions to isolate causal variables is well-motivated by prior interpretability literature (e.g., Geiger et al., 2021, 2024).\n\nEach task is evaluated under both in-distribution and OOD prompts, with OOD shifts chosen to stress test causal robustness.", "theoretical_claims": "The key theoretical claim is that internal model representations can be linearly decomposed into task-relevant (causal) and background (nuisance) components, enabling causal inference via projection.\n\nThis is not formally proven in the paper. The decomposition in Eq. (5) is derived from prior work on causal abstraction and DAS, but:\n- No formal conditions for identifiability are given.\n- No proofs of convergence or uniqueness of the causal basis Q.\n- The assumption that task behavior is mediated by a single variable X_T is strong and underspecified.\n- No formal causal graph given.", "experimental_designs_or_analyses": "Overall solid, tasks cover different scenarios.\nWhy value probing not performing well in a lot scenarios?", "relation_to_broader_scientific_literature": "The paper is directly based on the causal abstraction and mechanistic interpretability literature:\n- Geiger et al. (2021, 2024a, 2024b) on DAS and interchange interventions\n- Wu et al. (2023), Huang et al. (2024), and others on task-specific circuit discovery\n- Work on LLM trustworthiness via internal probes (e.g., Azaria & Mitchell, 2023; Ferrando et al., 2024)\n\nNovelties including:\n- Reframing causal representations as predictors of correctness, not just explanatory tools\n- Demonstrating improved robustness under shift over confidence scores and surface-level probes", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Strengths\n- Strong empirical benchmark with real OOD variation\n- Introduces novel causal predictors (simulation + value probing)\n- Bridges interpretability and evaluation\n\nWeaknesses\n- Causal assumptions (e.g., disentanglement, uniqueness of $X_T$) not deeply validated\n- Theoretical work directly based on prior work (DAS)\n- Method assumes knowledge of task structure to define $X_T$", "other_comments_or_suggestions": "Page 5: \n\n> “We empirically study whether these low-confidence predictions correspond to incorrect predictions in Section 4.3.”\n\n“Section 4.3” may be a typo. Based on context, the relevant content likely appears in Section 5 (Experiments). Please confirm and update accordingly.\n\nPage 5:\n> “We evaluate four correctness predictions methods over a suit of five language modeling tasks.”\n\n“We evaluate four correctness prediction methods over a **suite** of five language modeling tasks.”\n\n---\n\nExplicitly present the assumed structural causal model (SCM) as a diagram or formal figure.\nThe implicit model appears to be:\n$X \\rightarrow X_T \\rightarrow Y, X \\rightarrow B \\rightarrow Y$", "questions_for_authors": "See comments above", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Pd6ENKEeVx", "reviewer": "Reviewer_D6Cz", "summary": "This paper focuses on correctness prediction of large language models. It separates internal features into causal features and background features and suggests two approaches for predicting the correctness of model outputs. In one, permutations are used to determine whether predictions are robust against changes in non-causal features. Another one learns a linear model and checks how close predictions are to the decision boundary.\n## update after rebuttal\nThe authors addressed my concerns and clarified some misunderstandings. Therefore, I have raised my score to accept.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n---\n---\n### Q1. Comparison with strong baselines in correctness estimation\nA major concern is that \n> there is no comparison to any other approach for correctness prediction. \n\n**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.\n\nTo 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.\n\n---\n### Q2. Presentation: Clarification on background variables\n**2.1 Where does the background variable fit into the causal graph**\n\nPlease see our responses to **reviewer uRDA 2.1 and 2.2**\n\n**2.2 Are causal and background variables understandable for humans**\n\nThis 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.\n\n---\n### Q3. Details on dataset construction\n- For details on **split generation**, please see our response to **reviewer w7xm Q1**.\n- 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.\n\n---\n### Q4. Clarification: when does the causal relation $X_T → Y$ hold and when does it not hold\nAs discussed in our response to **reviewer uRDA 2.1**, the causal relation $X_T → Y$ is an abstraction of low-level neural networks.\n\n**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.\n\n---\n### Q5. Clarification: Why does the counterfactual dataset consist only of correctly predicted samples\nThis 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.\n\nThe 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?\n\nIntuitively, 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.\n\n---\n### Q6. Presentation\nWe appreciate the reviewer’s suggestion to clarify the definitions of \"Distributed Alignment Search\" and \"Interchange Intervention Accuracy.\"\n\nWe will include preliminary explanations of these terms in the revised manuscript!", "claims_and_evidence": "The paper claims that causal representations are beneficial for correctness prediction. There is experimental evidence to support this claim.\nIt is not fully clear how much these features model meaningful causal relations; however, they improve performance in the experimental evaluation. There is no comparison to any other approach for correctness prediction.", "methods_and_evaluation_criteria": "The evaluation makes sense. However, additional experiments would be beneficial to further support the claims made in this paper (see \"Claims and Evidence\"). There is also no comparison to any other methods. Further information on the experimental setup would be helpful, such as the number of samples considered for each experiment.", "theoretical_claims": "There are no theoretical claims.", "experimental_designs_or_analyses": "The experimental design is sound and valid.", "relation_to_broader_scientific_literature": "To the best of my knowledge, this is the first paper that considers correctness prediction in large language models under a causal perspective. Both the idea and implementation make sense, and the results look promising. However, there is no experimental comparison of other approaches for correctness prediction and how this approach would compare to them.", "essential_references_not_discussed": "There are no essential related works missing that I am aware of.", "other_strengths_and_weaknesses": "**Strengths**\n\nThe idea of using causal features for correctness prediction is promising since predictions based on non-causal features should be less reliable. This should, in particular, help when making predictions on out-of-distribution data, when predictions based on non-causal features that were informative before fail to be useful. The mathematical formalization of the problem is good. Considering output features, last prompt token features, and internal background features as alternatives to the causal features in the experimental evaluation helps highlight the benefit of the proposed approach.\n\n**Weaknesses**\n\nFor weaknesses regarding the experimental evaluation, see the fields above.\n\nClarity and Presentation: The paper illustrates the methodology sufficiently well to understand their approach. However, there is also much room for improvement to make understanding easier. For one, the paper would benefit from a figure on the methodology, outlining the approach in a visually understandable manner. Adding intuition on steps in the methodology would also be helpful. In particular, what do the background variables represent? And what is the causal graph that is assumed by the authors? I understand that $X \\rightarrow X_\\mathcal{T} \\rightarrow Y$, but where does $\\mathcal{B}$ fit in? To me, the most sense considering the methodology would be that $X \\leftarrow \\mathcal{B} \\rightarrow Y$ (B acting as a kind of confounding feature), but this is not entirely clear from the paper. Concepts such as distributed alignment search and interchange intervention accuracy would also benefit from more details in this paper to reduce the reliance of readers to be familiar with the corresponding papers.", "other_comments_or_suggestions": "- There is no reference to Figure 1 in the text.\n- The evaluation (Section 3.2) should be placed after the methodology and at the start of section 5 since they are not essential to understand the methodology but only matter for the experimental evaluation.\n- Should $x$ in Equation 3 not be $x_L$?\n- \"Since the model’s behavior can vary widely as the input distribution shifts, a reliable correctness prediction method should be able to robust under behavior changes.\" There is a small grammar mistake in here (Section 3.2).", "questions_for_authors": "1. What is the intuition and what are the causal assumptions regarding the background variables? What are examples of what should be causal and what background?\n2. Can the causal and background variables found by the method be analyzed such that they are understandable for humans?\n3. How exactly is the balanced dataset constructed?\n4. Under \"Value probing\", the paper states \"...when the causal relation between $\\mathcal{X}_\\mathcal{T}$ and $\\mathcal{Y}$ holds.\" What does this mean, i.e., when does the causal relation hold and when does it not hold?\n5. Why does the counterfactual dataset only consist of samples where the model prediction is correct? Is this sufficient to learn to predict correctness? Why?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EQ4AFd0ILM", "reviewer": "Reviewer_55c3", "summary": "This paper presents a method for estimating whether a language model’s output is correct by examining “causal” internal representations. It studies both symbolic tasks and open-ended tasks, finding that features which directly mediate model behavior are more reliable than simpler methods. The authors claim that causally grounded features remain robust under prompt changes and distribution shifts.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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)**.\n\nWe see this as an opportunity to further clarify the nature of our contribution using the additional space we would get in our next version. \n\n---\n### Q1. Clarification on the definition of  “causal”\n\nOur 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.\n\nNonetheless, 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**).\n\nAs 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.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "It is well-aligned with work on correctness estimation and the causal mechanism in language models.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths\n\n1. The authors provide a thorough evaluation across multiple datasets.\n\n2. The paper covers a broad range of tasks under both ID and OOD settings.\n\nWeaknesses\n\n1. The rationale for calling the method “causal” is unclear. The paper never formally defines “causality,” despite it being central to the proposed approach.\n\n2. Several core notions—causal features, causal mechanisms, causal interpretability, and causal role—are mentioned but not defined. The term “causal” appears throughout without a clear theoretical foundation.\n\n3. Some terms, such as “causal representation” and “causal structure,” have precise definitions in existing research. However, the paper neither references nor appears to adopt those definitions. Specifically, the conclusion highlights the assessment and use of “causal representation” and “causal structure” as central contributions, yet neither term is introduced in the main text. In the literature of causal representation learning, a “causal representation” refers to latent variables and the structure among them, while in causal discovery, a “causal structure” is the causal graph whose edges denote structural equations between variables.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. What are the precise definitions of these terms related to causality?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JJQ461Xw1k", "reviewer": "Reviewer_oMv7", "summary": "Understanding the behavior of complex black-box models has always been a challenge since the inception of deep neural networks (DNNs). This problem has worsened after the introduction of large language models (LLMs). In this paper, the authors investigate whether understanding the internal causal mechanisms of LLMs can improve the prediction of output correctness. In particular, they show that the most robust features for correctness prediction are those that play a causal role in the model's behavior using two proposed methods -- \nCounterfactual Simulation and Value Probing. The methods are evaluated on diverse tasks like symbol manipulation, knowledge retrieval, and instruction following, where causal features consistently perform better, particularly under out-of-distribution settings. In addition, the methods improve model safety and reliability by providing more accurate correctness estimates.\n\n## update after rebuttal\n\nThe authors addressed all my concerns and I vote for the acceptance of this paper.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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!\n\n---\n### Q1. Generalizability to complex real-world tasks\n\nWe 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.\n\n**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).\n\nTherefore, 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.\n\nWith this context in mind, we would like to further address limitations of these interpretability tools we used!\n\n**1.1 Reliance on human priors to identify causal mechanisms**\n\nThis is indeed a recognized limitation of concept localization methods. We discussed solutions in our response to reviewer uRDA Q3. \n\n**1.2. Scale to real-world complex tasks**\n\nWe 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.**\n\nIn 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.\n\nBelow, we highlight two additional practical domains where our methods are directly applicable:\n- **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.\n- **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.\n### Q2. Computation cost of counterfactual simulation\nWe 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.\n\nThe 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.\n### Q3. Which components contribute most to correctness prediction\nThanks for the suggestion! It indeed aligns with our discussion in Section 6.4 on using correctness estimation as an evaluation task for interpretability methods.\n\nWe 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.\n### Q4. Risk of overfitting causal subspaces\nAs 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.**", "claims_and_evidence": "While the paper provides substantial evidence to support its claims, it lacks some key details that raise questions about the effectiveness of the proposed method.\n\ni) While the methods perform well on the selected tasks, the paper does not extensively address how well these findings generalize to other, more complex real-world language model applications.\n\nii) The success of the methods relies on identifying causal variables, which may require domain-specific knowledge. It would be great if the authors could comment on the applicability of the approach in scenarios where causal structures are not well understood.", "methods_and_evaluation_criteria": "While the proposed method and evaluation criteria make sense within the context of understanding the internal causal mechanisms of LLMs, there are several questions concerning the methodology and evaluation of the proposed methodology.\n\ni) The benchmark tasks are diverse but they are relatively controlled and lack the full complexity of real-world scenarios (e.g., complex multi-hop reasoning or dialogue systems). This raises concerns about generalizability.\n\nii) The methods rely on identifying causal subspaces, which may require domain expertise or prior knowledge. This dependency limits the approach’s scalability to more opaque or less understood tasks.\n\niii) It would be great if the authors could comment on the computational bottleneck of counterfactual simulation, which involves multiple forward passes to assess causal stability. Would this prohibit the use of the proposed method for large-scale language models like GPT-4?\n\niv) While the paper compares causal vs. non-causal features in Sec. 5.2, a more granular ablation study (e.g., the impact of individual layers or attention heads) could reveal which specific components contribute most to correctness prediction.\n\nv) Will the selection of hyperparameters and causal subspaces for a given task inadvertently lead to overfitting on specific benchmarks?", "theoretical_claims": "NA", "experimental_designs_or_analyses": "Yes, I read the experimental setup and thoroughly reviewed the analysis.", "relation_to_broader_scientific_literature": "The paper presents a novel perspective to understand the correctness of language model predictions.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "By leveraging internal causal features rather than relying on output probabilities, one key strength of the proposed method is that it can effectively handle distribution shifts and hallucinations.", "other_comments_or_suggestions": "NA", "questions_for_authors": "Please refer to the \"Claims And Evidence\" and \"Methods And Evaluation Criteria\" for more details.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This work is nicely summarized by Reviewer w7xm and I quote, "The authors try to identify internal features of LLM that mediate causal effects on the final prediction output. In settings where LLM predictions align with the actual real-world causal process, causal features are assumed to resemble the ground-truth mechanism and, therefore, allow for predictions that are invariant to external disturbances or out-of-distribution queries." Two methods are proposed, which are then evaluated across a diverse suite of tasks under both in-distribution and OOD settings.

Overall, the reviewers were very happy with the submission and positive regarding the contributions. The paper nicely balances the diverse tasks within language modeling along with the complexity of presentation and clarity of writing. The rebuttal further clarified the issues of the reviewers and strengthened the overall story of the work. Even the most critical reviewer explicitly stated that they are not opposed to acceptance. I would still advise the authors to carefully take into account all the comments by the reviewers and make the necessary changes in the final version of the manuscript.

The AC recommends acceptance.
Decision: Accept (poster)