Counterfactual reasoning: an analysis of in-context emergence
Authors: Moritz Miller
Bernhard Schölkopf
Siyuan Guo
Venue: NeurIPS 2025 poster ()
Keywords: In-context learning; transformers; reasoning; causality
Primary area: general_machine_learning
TL;DR: whether and how transformers perform in-context counterfactual reasoning

Abstract: Large-scale neural language models exhibit remarkable performance in in-context learning: the ability to learn and reason about the input context on the fly. This work studies in-context counterfactual reasoning in language models, that is, the ability to predict consequences of a hypothetical scenario. We focus on a well-defined, synthetic linear regression task that requires noise abduction. Accurate prediction is based on (1) inferring an unobserved latent concept and (2) copying contextual noise from factual observations. We show that language models are capable of counterfactual reasoning. Further, we enhance existing identifiability results and reduce counterfactual reasoning for a broad class of functions to a transformation on in-context observations. In Transformers, we find that self-attention, model depth and pre-training data diversity drive performance. Moreover, we provide mechanistic evidence that the latent concept is linearly represented in the residual stream and we introduce designated *noise abduction heads* central to performing counterfactual reasoning. Lastly, our findings extend to counterfactual reasoning under SDE dynamics and reflect that Transformers can perform noise abduction on sequential data, providing preliminary evidence on the potential for counterfactual story generation. Our code is available under *https://github.com/mrtzmllr/iccr*.

Counterfactual reasoning: an analysis of in-context emergence
Moritz Miller Bernhard Schölkopf
Siyuan Guo 13
Maa Planck Institute for Intelligent Systems
2ETH Zurich
University of Cambridge
Abstract
Large-scale neural language models exhibit remarkable performance in in-context learning: the ability to learn and reason about the imput context on the fly. This work studies in-context counterfactual reasoning in language models, that is, the ability to predict consequences of a hypothetical scenario. We focus on well-defined, synthetic linear regression task that requires noise abduction. Accurate prediction is based on (1) inferring an unobserved latent concept and (2) copying contextual noise from factual observations. We show that language models are capable of counterfactual reasoning. Further, we enhance existing identifiability results and re- duce counterfactual reasoning for broad class of functions to transformation in-context observations. In Transformers, we find that self-attention, model depth and pre-training data diversity drive performance. Moreover, tic evidence that the latent concept is linearly represented in the residual stream and we introduce designated noise abduction heads central to performing counterfac- tual reasoning. Lastly, our findings extend to counterfactual reasoning under SDE dynamics and reflect that Transformers can perform noise abduction on sequential data, providing preliminary evidence on the potential for counterfactual story gen- eration. Our code available under sttps://githur gighgu......../ com http///tttmlrrireeeeeee
provide mechanis-
Introduction
Thinking is acting in an imagined space. Konrad Lorenz
Lorenz, 1973
Large language models demonstrate remarkable capability in task-agnostic, few-shot performance, such as in-context learning and algorithmic reasoning [Brown et al., 2020). Despite their success, one of the grand challenges of artificial general intelligence remains the ability to unearth novel knowledge. This requires principled reasoning over factual observations instead inventing information when uncertain, phenomenon termed hallucination AAhiam et al. 2023).
Thinking, in Konrad Lorenz's words, is acting in an imagined space. Counterfactual imagina- tion/reasoning, studied in early childhood cognition development (Southgate and Vernetti 2014] and formalized in causality (Schölkopf, 2022 Pearl 2009 Hernan 2024]. predicts the consequences of changes in hypothetical scenarios. It answers whhat if questions and predicts potential outcomes that could have occurred had different actions been taken. Such principled thinking enables fast and responsible learning, e.g., efficient reinforcement learning [Mesnard et al. 2021 Lu et al. 2020), counterfactual generative networks [Sauer and Geiger] 2021 responsible decision- makinn in complex systems |Wachter et al., 2017), and fairness [Kusner et al. 2017].
In precision healthcare, one may ask what would have happened had a patient not received treatment to assess the individualized treatment effect. One step further, in-context counterfactual reasoning per- tains the hypothetical consequence for individual after observing the effect of different treatments
"Corresondencce moritz. mi wwwereteeiingen...... www.eiinnnngg..........m de
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

S was given 75- doxorubicin.
J received 58 of doxorubicin.
B took 69 of doxorubicin.
Had_ | received tee what would be predicted tumor size reduction?
Consider patient
Her tumor has shrunk 32%..
The tumor has by 3%..
The tumor size decreased by 49%.
11
yj
En
Y/
2
CF I
Noise abduction Infer noise from patient
CF y
f(x ; (
CF
The predicted reduction would have been 38%
Intervention dosage
Intervene
pre-training
Prediction
in-context
Predict counteriactual
CF
Figure I: In-context counterfactual reasoning. Training on corpus of sequences that come from mixture of distributions (each e on the far left represents a single sequence from distinct distribution parameterized by 0). Suppose each observation satisfies = fæ, Myy for some noise uy. An in-context sequence takes the form examples. This is concatenated with index token 2 referring back to observed factual observation (zz, B5) when r (x), M....... Tn Mr x²) In-context counterfactual reasoning can be measured via accurate prediction on yf Accurate prediction requires noise abduction from factual observation, that is, to infer Uy consistent with (Iyz Vj), and prediction based the intervention xF and inferred Uy.
and the hypothetical new information
on multiple patients. Figure displays a simplified setting where n patients receive neoadjuvant chemotherapy treat breast cancer. Given different dosages of doxorubicin, the tumor size reduction recorded as output. A counterfactual reasoner now imagines the tumor size reduction of patient j. had she received 45 me of doxorubicin instead of the prescribed dosage of 65 mg Augmented with mn agentic verification, such principled counterfactual reasoning would be valuable tool for automatic scientific discovery and guardrail for safe Al deployment.
To concretely understand to what extent language models perform in-context counterfactual reasoning. use the counterfactual framework of Pearl [2009) and study controlled synthetic setup similar to Garg et al. [2022]. That is, let y = f(z, uy) for some function f E F, for function class F. The model predicts target 1f given counterfactual query z conditioned on a prompt sequence (x1, f1.............................. . Tk. yk+ 2, _CC)) where is index token indicating the position of the factual observation that such counterfactual query is based on. Given factual observation (I, y). counterfactual reasoning requires three steps:
noise abduction: infer uy that is consistent with the factual observation (ax v).
intervention do X = x. (F); intervene on X consider the hypothetical value CF
prediction ~ _ [T, uy): predict the effect of changes x by inheriting noise uy from the factual observation.
For a complete introduction on the basics of causality, see Section 2.1 In contrast to natural language. under this controlled setup, we can formally ask
Can language models perform in-context counterfactual reasoning?
with concrete metrics EE
CF Upred)
for model prediction speed and squared error function
This work studies counterfactual reasoning in exchangeable data, as we note language models are pre-trained on mixture sequences coming from different distributions. The training dynamics allow random permutations over different input sequences, i.e., implicitly assume sequences are exchangeable (see Def. [T. [De Finetti [193T] shows that any exchangeable sequence can be modeled mixture of conditionally i.i.d. sequences. In other words, there exists a latent variable 0 with

---

probability measure π such that
n III p(sa 0)dz(0)
P(s
(I)
for 8; the 7th sequence. Pre-training thus amounts to learning mutual information 0 between sequences and their prior (0). At inference time, the model computes the posterior predictive distribution: given sequence with tokens zi, " F.j, the next token prediction writes
,..,,,.
p(z)(='e)) - f ()
Guo et al. 2024a bJ have established a theoretical foundation to study observational and interventional causality exchangeable data. We extend this framework counterfactuals. Building upon previous work, our contributions are:
We relate in-context counterfactual reasoning to transformations on the factual in-context observations (Lemma D. Extending known identifiability results (Nasr-Esfahany al. 2023]. we provide theoretical guarantees for our exchangeabee setup (Theorem I).
We empirically show that in-context counterfactual reasoning emerges Transformers in form of designated noise abduction head (Section 4.4). Causal probing reveals that the residual stream linearly encodes the latent 0 (Section
We find that data diversity in pre-training, self-attention and model depth are key for Transformers' performance (Sections 4.2 and 4.3). More interestingly. our findings transfer cyclic sequential data (Section 55, demonstrating concrete preliminary evidence that language models can perform counterfactual story generation in sequential data.
Background
2.1 Exchangeability and causality
Definition (Fxchangeable sequence). A sequence of random variables is exchangeable if for any finite permutation a of its indices,
P(S1, ..., S.) (I) So(n))
Here we treat each random variable S as a sequence with bounded context length t.
Causality fundamentals. When X causes Y, the structural causal model (SCM) is determined by functional assignments fx, Jy and exogenous variables Ux, i.e., X = fx Y = (X,UU)) An intervention performed as do X = x). Thus, one replaces X with value and allows to inherit the replaced value, i.e., fy ). A counterfactual statement of the form "Y would be y had X been in situation = u" is often denoted as (u) = y. For explicit representation, we use y CF for values intended for counterfactual predictions. In contrast the interventional setting, the counterfactual entity inherits noise consistent with the factual observation.
variable X is represented
X := =,
Transformer architecture
Transformers Vaswani al.) 2017 operate sequence of input embeddings by passing them to blocks consisting attention and multi-layer perceptron (MLP). Mimicking present-day large language models, this work focuses the decoder-only, autoregressive GPT-2 [Radford al. 2019 architecture for next-token prediction. The softmax operation with causal masking is denoted by softmax.. Given sequence of input embeddings E E T of context length T and embedding dimension E, the model first projects each token into D-dimensional hidden embeddings via X, EWE pos(E) Wr, for embedding matrix WE E REDD positional embedding Wp R with absolute positional encoding. Given an input sequence of embeddings, multi-head attention at layer İ passes the embeddings to query, key, value weight matrices Whh Wh W E R DxD
and computes attention per head as A = softmax,
(
Then the model

---

concatenates the multi-head output and transforms it via output weight matrix Wo E R Rx as M, ) ..... A," The output is added into the residual stream as R; = - M₁ and passed to an MLP X = MR.) + R, After the last Transformer layer, the embeddings are mapped back to logits through the unemhedding matrix: O Xx We ignore layer normalization.
3 In-context counterfactual reasoning
We study a linear regression task that requires noise abduction. The structural causal model (SCM) considered is fx(Uxx Ux. (x, 3Xx+UX variables. We are interested in accurate counterfactual prediction on y given new information c. and factual observation (ar, :) (U,, Bu, U,y), where y 3r Vyy
+ Uy, where Ux. Uy are independent exogenous
Pre-training sequence generation. Our pre-training corpus consists of a mixture of sequences coming from different distributions that share the same causal structure but different functional classes. Throughout, we use the terms training and pre-training interchangeably. To generate a single sequence. randomly draw latent variable 0 and parameterize both the regression coefficient / and noise variables Ux. Uy based on f. Each sequence takes the form (I1, V1..... .w, Un,, 2, x .F where ny in-context examples are sampled from the SCM parameterized via I and denotes the index token for that factual observation which the counterfactual query would be based on. We vary the number of in-context examples nı observed each sequence to avoid prediction based on information from positional encoding only. In particular, we respect the causal attention mask by ordering observed variables in topological order in a sequence format.
CF Training objective. We minimize EE Ely - for Upred the model's predicted outputs and y the true counterfactual completion. We choose l to be the mean squared error (MSE). In-context counterfactual reasoning can be represented via the posterior predictive distribution
CF [11, Vi.............................. En: Vns x CF == [ 5(y Br CF +y² (Bz,)Pa(B1 (f)x (2) p(y
The model is required to learn how to compute posteriors Pp. given a query at pre-training, In- context completion reduces to identifying z, inserting query elements (xz, Va. (f and weighting the resulting CF with the posterior of 3 given x.
We first present Lemma and show how accurate prediction on counterfactual values y can be reduced to transformation on in-context observations. Lemma [] incorporates broad class of functions, including additive noise models (ANMs) [Peters et al. 2014), multiplicative noise models and exponential noise models.
Lemma 1 (Counterfactual reasoning as transformation on observed values). Suppose y =T)(2)) ") for some function T y xl - y. Assume for any fixed f((x) V, the inverse T-1(((x), exists for all y, i.e., u T-1 (f(), y). Then, counterfactual reasoning reduces to learning a transformation h observed factual observations (x, v) and counterfactual information 2: with
y h (f(x v(
where -UU
v)) operates on elements of y only.
Lemma T shows that given the existing information in the query, in-context counterfactual reasoning on the fly task Transformers are known to be capable of performing [Garg et al. 2022 Akyürek et al., 2023 von Oswald et al. 2023 The assumptions in Lemma I cover a broad class of functions, for example.
no more than estimating the transformed function T, f. -
Additive noise models (ANMs). Functions taking the form Y = f(x) + for arbitrary function can be equivalently represented as y = (x), u) f(xx + u. For fixed f(x).
the inverse T (f(x), exists and reads u = (f( v) -y-f(xx))
Multiplicative noise models. Functions taking the form = f(X)-U for arbitrary function f. For fixed fz), the inverse T-1 exists and reads T (f(x). v) T(
Exponential noise models. Functions taking the form Y =ex([[X] + U) with inverse u = TT = log(y) f(x).

---

On top of that. Nasr- Esfahany et al. [2023] provide three cases under which the class of bijective generation mechanisms (BGMs) is counterfactually identifiable. Under a of assumptions, one can disentangle the causal mechanisms solely by knowing the causal graph and observational data. Our setup can be subsumed under that framework of BGMs. We therefore extend the result in the Markovian case [Nasr-Esfahany et al. [223)] to our exchangeable setup and show that the transformation on observed values counterfactually identifiable. Details are deferred to Appendix B
Theorem 1 (Counterfactual identifiability under exchangeability), Let X, Y, U, 0 scalar random variables with X IL U|B. Take T Y x U - ) with (T)()) u). Assume (f E V, the inverse (f)(f(x),) exists for all y, i.e,, continuous. Suppose further that f(x) continuous Vx, and strictly monotonic in a. Then, set = TTT((XX) Given the joint law P Px,y is counterfuctually identifible.
T=1(f(x).)), and Vf(x) E y, T(f(x),-) is
Experiments
Based on the training setup described Section 3 with details in Appendix C we choose a controlled synthetic setup to allow concrete evaluation on in-context counterfactual reasoning.
4.1 Language models can in context perform counterfactual reasoning on linear functions
We assess the impact of model architecture choice on in-context counterfactual reasoning performance. In contrast to natural language tasks, the performance in task is clearly defined and can be measured as:
accuracy in prediction: MSE ,FF).
learning efficiency: number of in-context examples achieve satisfactory prediction.
We compare the STANDARD decoder-only GPT-2 Transformer architecture against several autoregres- sive baselines. We the terms STANDARD and GPT-2 interchangeably. The models considered are:
GPT-2
Radford et al.
2019] 12 layers, 8 attention heads, hidden dimension 256,
LSTM
Hochreiter and Schmidhuber
1997|: 2 layers, hidden dimension 256,
GRU Cho et al.
2014: 2 layers, hidden dimension 256,
ELMAN RNN
Elman
[990]: layers, hidden dimension 256.
We evaluate each model unseen sequences sampled in-distribution and report results averaged over 6400 sequences. Our synthetic generation yields E[yCF] = 0 and log (sg(std()(FF) = 2.56. Appendices and Dinclude data generation particularities as well experimental and model details.
Figure compares the impact of different model architectures on in-context counterfactual reasoning. We plot the log-transformed in-context MSE against the number of in-context examples observed in prompt. We hypothesize that more observed examples lead to higher counterfactual prediction accuracy, and that better models require less in-context examples for accurate prediction. For contexts of over 14 in-context examples, the Elman RNN achieves significantly higher in-context MSE than the rest. Across LSTM, GRU and STANDARD Transformer, we observe significant performance difference. In Appendix D we reference literature discussing the RNN architectures. For now, we confine ourselves to autoregressive Transformers such that all future statements refer to this setup only.
4.2 Counterfactual reasoning emerges in self-attention
Self-attention. Lemma D shows the reliance of counterfactual reasoning on copying in-context observed values. Self-atiention has been shown to perform copying by implementing induction heads OOlsson et al. 2022 Akyürek et al.) 2024). We thus hypothesize that counterfactual reasoning emerges in the attention sublayers of the Transformer.
We conduct an experiment that switches on and off attention layers in the Transformer and separately train each model on our task. Figure 3a shows the in-context MSE for MLP-Only and AO (attention- only) with 2, 4, layers and the STANDARD setup. training time, we observe that AO models mimic the loss curve of the STANDARD setup. while the MLP-Only's 's loss remains constant at
5

---

8-ayer, 1-head Elman RNN LSTM
GPT-2
In-context MSE (log)
4
30
Number of in-context examples
Figure 2: Model comparisons for in-context counterfactual reasoning. In-context counterfactual prediction accuracy measured via log-transformed MSE averaged over 6400 sequences number of in-context examples observed in prompt. We compare GPT-2 (STANDARD), LSTM, GRU, and Elman RNN. Though not significant, GPT-2 achieves lowest error and fastest convergence rate for small number of in-context examples. For more than 14 in-context examples, the Elman RNN obtains significantly [Elron 1979 higher in-context MSE than the three other architectures.
the
the level of the theoretical variance, Var(YCF). Figure 3a shows that counterfactual reasoning performance improves for all considered Transformers with increasing context length. Although at higher loss than the STANDARD Transformer, smaller models do perform better longer contexts. Oblivious of context length, MLP-Onnl does not appear to reason counterfactually.
Model depth. Elhage et al. [221]] show that self-attention performs copying task. Lemma relates counterfactual reasoning to copying multiple values and learning transformations. Such task requires the composition of attention heads to pass information across layers. We hypothesize that model depth plays an important role. Fixing the total number attention heads 8, we train 4 Transformer-basse with layers of 1 head each. We expect that depth does not influence performance, given the same number of attention heads, all models should have comparable in-context MSE. Figure represents the loss on 6400 sequences of 35 in-context examples for both Full and AO models. Additional results can be found in appendix With increasing depth, loss declines for both setups. In partieular, the Full 8-layer, 1-head Transformer yields lowest MSE, while the 4-layer, 2-head designs achieve competitive results.
models with increasing depth: from 1-layer, 8-head Transformers to models
MLP-Only
Full
2-layer AO
AO
In-content MSE (log)
4-layer AO
In-conteect MSE (log)
8-ayer
GP
2
10
20
30
40
50
1-ayer, 8-head 2-layer, 4-head 4-aaeer, 2-head Model
Number of in-oontext examples
(a) Switching attention. In-context MSE stagnates with increasing number of in-context examples observed for the MLP-Only model. We compare to attention-only (AO) Transformers of 2, 4, 8 layers with 1 head each as well as the STANDARD setup. For all Transformers. in-context MSE decreases more in-contex examples are observed.
(b) Varying depth. Keeping the number of attention heads constant
compare Full and attention- only (AO) Transformers with layer, heads; layers, heads; layers, 2 heads; 8 layers, head. We observe decrease in in-context loss as model depth increases. We evaluate 6400 sequences of 35 in-context exam- ples each.
Figure 3: Attention and model depth matter.

---

Latent parameter. A final way to confirm the relevance of self-attention is by probing the residual stream for the latent 0. As 0 governs the distribution of the in-context sequence, is natural to check for that signal. By doing so, we find that the 8-layer AO Transformer learns the latent parameter early on. Figure shows that linear model predicting 9 from the residual stream after layer 2 does so at adjusted R² z 0.9104. After this initial spike, all future layers absorb that information as the adjusted R² stays above 0.9. addition we train separate probes on the regression parameter 3 and find similar results to probing for 0. As ~ ~O, 1). this is in line with our intuition. Moreover, Figure 4bcchecks whether the difference of residual streams after and before every layer carries latent information. In fact, we have additional evidence that layers 2 and 5 are central in learning 0 (and 3), while the final layer's s contribution is marginal at adjusted R < 0.2. By running this standard experiment, find that self-attention alone suffices linearly encode the latent theme. We thus provide additional support for the Linear Representation Hypothesis (Park et al. 2024
1.0
Adjusted
Aduutedd
0.4
0.4
0.2
0.2
Layer Index
Layer Index
(a) Relevance every layer. We plot the adjusted R²
Additional relevance of this layer. We subtract
after every layer against the layer indices. Starting at the second layer, we observe adjusted R² > 0.9
residual streams before each layer from the residual stream after and report the adjusted R². Layers and are especially relevant while layer 8 does not add substantial information.
for predicting 0, indicating that the latent is encoded linearly in the residual stream.
Figure 4: The Transformer linearly encodes the latent parameter. We train linear probe on 6400 prompts from a fresh evaluation set after every layer. We evaluate on 1280 sequences. All layers after the first one encode relevant information for predicting 0 from the residual stream only.
4.3 Data diversity and out-of mistr--distribution (OOD) generalization
Causal structure identification, previously deemed impossible task from i.i.d. observational data alone [Pearl] 22009), has been shown to be feasible in multi-domain data [Guo et al. [024a]] a natural setting for exchangeable data. We hypothesize that (counterfactual) reasoning similarly only feasible under diverse pre-training data. As measure of pre-training data diversity, we use the effective support size [Grendar 2006), defined
Ess(O) exp
Σ p(U) loo)))))()) exp (H())), dEGe
where Og is the space of latent 0 sampled at pre-training, and H the Shannon entropy [Shannon) 1948). We further compare the pre-training data diversity's impact on simple OOD generalization. We pre-train on both UNIFORM and NORMAL sampling of latent 6 and evaluate on UNIFORM sampled test data. Figure Sa shows clear trend that diverse pre-training data, measured as a higher effective support size, yiclds lower in-context MSE for both in in-distribution and OOD cases. Appendix E shows that such generalization ability persists when evaluated on NORMAL, and that in-context loss declines with increasing data diversity across a changing number of in in-context examples. We also discuss how this relates to existing research.
4.4 Emergent model behavior
The ordinary least squares estimator for the regression parameter 3 writes
B Σn= (Ix () 8)
> (xi 8)2
7

---

Training distribution
Saquence
Normal (000)
loxs (log)
5-dirensional 4-dimnntional 2-dimensiiona
2-dmeeniioal
Uniform (in-diat)
1-dmensional
Inconntt MSE (log)
ioconct loxs
2
Optinal MSE 30 examples
b
Effective support dize (log)
Training ieerations (x10")
(a) Data diversity in pre-training. In-context MSE (log-scaled) averaged over 6400 prompts at 35 in-
(b) Phase transitions during training. In-context training MSE against the number of training steps. When training on 100' 000 iterations, observe that the model with F 5 reaches phase transition after more than 50'000 steps. Depending on embedding di- mension E, the phase transition occurs sooner or later.
context examples against log-scaled effective support size. Fach point represents one fully pre-trained model on either the UNIFORM sampled or NORMAL sampled 0. Models are evaluated on UNIFORM.
Figure 5: Data diversity and emergence. In-context counterfactual reasoning emerges at training. To generalize to unknown 0, the model is to be trained on a sufficiently diverse pre-training corpus.
which leads to the fitted value o = Bxx xz) + Yz. and the baseline MSE between the linear
regression implementation U and the ground truth counterfactual, MSF V yF We refer to this as the (estimated) optimal MSE and note that lies between 0.6132 and 0.4739 for 30 to 50 in-context examples. We therefore compare models trained on data with embedding dimension E E 21, 4,5} to this estimated optimal MSE. Here, we fix = 14 to enable faster learning relative to the setup with variable index token. Figure 5b displays the training progression for Full GPT-2 Transformers trained for 100'000 steps. We observe that the ability to attain the estimated optimal MSE emerges training proceeds (Wei et al. 2022 Nanda et al. 2023 Given that primarily train on 50'000 training steps, we observe that at E 5, the counterfactual reasoning behavior yet emerge. For sufficient training length, therefore, the model implements solution performing better than the estimated optimal MSE. In Appendix [E] we confirm this behavior for variable by introducing a designated noise abduction head which emerges at the 7th layer of the Full GPT-2 architecture. We also move beyond linear regression, and we compare our results to models trained on the setup in Garg et al. 2002
5 Cyclic sequential dynamical systems
I Language is sequential. Sentences within story depend on each other. Counterfactual story generation naturally occurs by prompting the model with a factual story and querying it to complete the story under hypothetical scenario. This is done while keeping the noise unchanged. Although difficult to evaluate in language, our key insight is that such behavior can be mimicked by ordinary differential equations modeling the underlying causal mechanism [Mooij et al.) 2013 Peters et al.) 2022 Lorch et al. 2024]. Given an imitial condition (æo. Mo),, a dynamical system determines the state (2t, Vt), for t 2 O. We model causal dependencies using Itō stochastic differential equations (SDEs),
dY: f(x,, Y,)di + oy- dW,
dX,- == g(X, dt +ax dU,
(4)
Xo Eo and Yo : vo:
for initial condition (Eo, un). Here, W., Ur denote independent Brownian motions with constant diffusion by oy,
Ty, and fig represent the drift coefficients for X, Y. respectively.
Autoregressive models are discrete. To address the challenge of adapting continuous SDE data generation to discrete observations, we draw n event times from a uniform distribution over bounded interval. We evaluate the above SDE at time l, for m S n and input observations to the model. Instead of deterministically slicing the interval into equivalent regions, we can capture the complete bounded interval as the randomly sampled set of event times, independent across iterations. Noise abduction is thus required over both the realization of and of event times tm < ty
8

---

In summary, we provide context (It, ,.,,... 2t, Mr..), counterfactual token indicating the start of counterfactual story generation, and the start of the counterfactual sequence (2 Vt, Then, we train the model on completing the full counterfactual story. We thus ask the model to generate (a12 1/2 ....,-,, ). Models are evaluated on the MSF across all predicted tokens. Contrary to above, we assess performance not just on the final token, but over the entire counterfactual continuation consisting of (n (n--) 1) 2 tokens. Additionally, the full context corresponds to a single instance of factual story. Given this story, the model must infer the noise associated with all observational examples, rather than just single example pair. We thus have deterministic functional assignments f. and copy the noise We, Ug from the observational sequence. With constant diffusion, we invoke Lemma] which connects the eyclic extension to the regression setup.
y -f,[ F(X,F y. )ds - f((x Y,) Y,,
-
(5)
We conduct preliminary experiments on data generated from the Lotka-Volle Volterra model describing predator-prey dynamics,
Lotka 1910
f(X,,Y) == ax,
(6)
g(X,, Yi) -
(7)
with a: 8, 2 0. Appendix FF specifies training and data generation details and analyzes the attention behavior in the Transformer. Figure 6a shows that the Full GPT-2 Transformer obtains the lowest in-context MSE across the four studied models. Error bars [Erron| 1979) point competitive performance across RNN-type architectures. We evaluate the models on contexts of n 20 realizations of in-context pairs. Recall that these are not conditionally i.i.d. in-context examples as above, but rather follow a sequential pattern. Figure 6b) illustrates this sequentiality. Given query consisting of interleaved observations of the prey, predator concentrations (I, y). the model infers the underlying latent governing the system. It then infers the full counterfactual trajectory beyond (x ) Notably, while the observed y [1.1, 1.9]. the model correctly infers that the counterfactual realized a lower level, y CF E [0.1, 0.7]. This indicates that the model internalizes the underlying dynamics and uses them to generate coherent predictions.
0.02
Lalka-Vellerra concentration
In-pontaxt MSE
0.01
Elman
LSTM
GRU
GPT-2
0.3
Model
Time
(a) In-context evaluation on Lotka- Volterra SDEs.
(b) Prediction of Votke-Volterra SDEs in-context.
In-context MSE of four models queried with data on
The query consists of the observational interwoven (z, y) data. In-context, the model learns the dynamics governed by the latent 0 and completes the counterfac- tual trajectory for all tm. > O. Although observational spans [1.1, 1.9). the model infers that counterfactual constrained [0.1, 0.7].
20 distinct time points. The MSE is competitive for Elman RNN, LSTM, and GRU at over 0.1. For the Full GPT-2 Transformer, we observe MSE 0.0025, which is significantly below the performance of the former three architectures.
Figure 6: Cyclic causal relationship.
Discussion and related work
Counterfactual reasoning in natural language. Prior investigation into counterfuctual story generation has been mainly focused on empirical evaluation on off-the-shelf language models. Tandon et al. (201)) [i et al. (2023) study lexical associations between factual statements and the hypothetical "what if" scenario; Qin et al. [019 provide datasct for counterfactual rewriting

---

evaluation. [Bottou and Schölkopf [2025] link the ability to hypothesize about counterfactuals to generating meaningful text beyond the training distribution. Raviogel et al. [2025] study language
directly as structural causal acyelic models and Yan et al. [223] focus on representation identification for adaptive counterfactual generation. This work respects the sequential nature of language and
performs controlled pre-training and concrete evaluations that demonstrate preliminary and promising
evidence that language models are capable of counterfactual story generation.
In-context learning. Language models [[rownn et al. 2020] have shown surprising in-context
reasoning ability demonstrating the potential for practical user interactions in the form of chatbots.
Transformers, the backbone architecture for modern language models, have been shown to in-context
learn different function classes [Garg et al., 2022]. Some studies investigate the models' ability to
perform gradient descent over in-context examples (Dai et al. 2023 DDutch et al. [024]] von Oswald
al.] 2023 while others discuss implicit Bayesian inference Xie et al. | [022 Falck et al.) [024] Ye
and Namkoong, 2024] and induction heads [Olsson et al. 2022 Akyürek et al.
2024). J. A recent line of
work considers guarantees for in-context learning algorithms on linear models [Akyürek et al. 2023).
Causality and exchangeability. Pearl 2009]. Hernan 2024] formalize the notion of counterfactuals
in causality. |Peters et al.
[022 Lorch al. [2024] study cyclic causal models via differential
equations. Mixture of i.i.d. data or multi-domain data have been the foundation of modern pre-
training corpora. Guo et al. [2024a] demonstrate that inferring causality is feasible in such a mixture of i.i.d. data, providing the potential for language models to exhibit causal reasoning capabilities. Recent work has continued to investigate connections between causality and exchangeability from intervention [Guo et al. 2024b) to representation learning [Reizinger et al.) [025]. Concurrently,
interest in exchangeability bit and language models has resurfaced. [Zhhag et al. [2023] study topic
recovery from language models and Falck et al. [2024). Ye and Namkoong (2024) Xie et al. [2022]
discuss in-context learning as performing implicit Bayesian inference in the manner of de Finetti. This provides potential for causal de Finetti, which allows structured in-context learning with the potential for principled reasoning and fast inference. This work serves as a preliminary study on in-context counterfactual reasoning in the bivariate case and demonstrates its potential.
Limitations. Our study focused on controlled synthetic datasets to circumvent evaluation difficul- ties inherent to natural language. We therefore assess reasoning from a causal perspective (Peters
al. 2017 Mondorf and Plank] 2024) as algebraic manipulation of existing knowledge [Bottou We only train on small-scale GPT-2 type models to demonstrate preliminary evidence on counterfactual story generation. Furthermore, we assume no observed confounders, which difficult to satisfy in the real world, as observations are finite and contain missing information which requires var-abl- generalization [Guo et al. 2023). Realistic benchmarks on natural language evaluation and larger-scale pre-training that provide principled reasoning for real-world impact, such as scientific reasoning, represent potential avenues for future research.
2014
Broader impacts. Since our models are not trained on natural language, the typical societal concerns associated with counterfactual story generation do not directly apply. We view our current work as a scientific exploration of desirable properties for the next generation of Al models.
7 Conclusion
We discuss in-context counterfactual reasoning in both linear regression tasks and cyclic sequential data. We observe that counterfactual reasoning can emerge in language models. We find that Transformers, in particular, rely on self-attention and model depth for in-context counterfactual reasoning. Data diversity in pre-training is essential for the emergence of reasoning and OOD generalization. We provide insights on how counterfactual reasoning can be interpreted as a copying and function estimation task, with theoretical and empirical evidence over broad class of invertible functions II. We mechanistically support our findings by introducing a designated noise abduction head and probing the residual stream for the latent concept. We develop a sequential extension grounded in the Dotka-Volteraa model of predator-pre dynamics [Lotka| 1910]. We show that both Transformers and RNNs can in context learn such complex dynamics, and thus provide preliminary but concrete evidence on the potential for counterfuctual story generation that is important to scientific discovery and AI safety.
10

---

Acknowledgements
MM thanks Jonas Peters for valuable discussions the concept of counterfactual reasoning and the role noise abduction. MM also thanks Hsiao-Ru Pan and Florent Draye for their input on identifia- bility and causal probing. MM acknowledges financial support from the Seerrad-ddenuer-SSiee Stiftung.
Bibliography
J. Achiam, S. Adler, S. Agarwal, L. Ahmad, I. Akkaya, L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, S. Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023. URL
https://arxil: org abs/ 2203 08774
E. Akyürek, D. Schuurmans, J. Andreas, T. Ma, and D. Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023, URL
E. Akyürek, B. Wang, Y. Kim, and J. Andreas. In-context language learning: Architectures and algorithms. In Forty-first International Conference on Machine Learning, 2024. URL https wtwpe//wwwwweeeetottiiiiieieeee www.torumm 12239m92rr
//openreview.o
M. Arjovsky, A. Shah, and Y. Bengio. Unitary evolution recurrent neural networks. In Proceedings of the 33rd International Conference on International Conference on Machine Learning Volume 48, ICML' 16, page 1120-1128. JML..org. 2016.
L. Bottou. From machine learning to machine reasoning. Machine Learning, 94(2):133 2014 doi: 10.10077/ URL
L. Bottou and B. Schölkopf. The fiction machine, April 2025. URL https://www. org/ publications ssssmmiammmaa htt.aaaeeeeaaatttiieeeeeeeeeiiie machinne/
T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert- Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S, Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, Sutskever, and D. Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877-1901. Curran As- sociates, Inc., 2020. URL URU/attps://pprocediagg neurips wttpcccppppppaaceeee feaper_filese files/ /
pdf
Charleux, E. Roux, Goyallon, G. Feverati, and P. Nagorny. Lotka-volterra equations. 2018. 3_
readthedocs ht//e///atts////ttbooo..oo rst/ htw.:/a.i.............a.rarteeae http://LLotaaaoleeeeeeeeeooeeoo Volterra_modell html
K. Cho, B. van Merriënboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Bengio. Learning phrase representations using RNN e for statistical machine translation. In A. Moschitti, B. Pang, and W. Daelemans, editors, Proceedings of the 2014 Conference Empirical Methods in Natural Language Processing (EMNLP), pages 1724-1734, Doha, Qatar, Oct. 2014. Association for Computational Linguistics. doi: 10.31 1//1//////111/1/9 URL btap:::/aacaantholeee org HTrgDD1--1799
D. Dai, Y. Sun, L. Dong, Y. Hao, Ma. Z. Sui, and F. Wei. Why can GPT learn in-context? language models implicitly perform gradient descent as meta-optimizers. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023. URL https /ooeeneview www.:oourr/oideeeeooeooe.. fabHRd
B. de Finetti. Funzione caratteristica di un fenomeno aleatorio. Atti della R. Academia Nazionale dei Lincei, Serie 6. Memorie, Classe Scienze Fisiche, Mathematice e Naturale, 4:251-299, 1931. in
URL //www, inetti ortppeeeiitt..../eeeee/mmmoaaaaa

---

G. Deutch, N. Magar, T. Natan, and G. Dar. In-context learning and gradient descent revisited. In K. Duh, H. Gomez. and S. Bethard, editors, Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pages 1017-1028, Mexico City, Mexico, June 2024. Association for Computational Linguistics. doi: 12118653/1/204.mmaca/// //aclanthology.o org/ 2024. naacl-long. 58/
URL https
B. Efron. Bootstrap methods: Another look the jackknife. The Annals of Statistics, 7(1):1-26, 1979. ISSN 00905364, 21 21688966. http://w:// jstor
eettable29958800
N. Elhage, N. Nanda, C. Olsson, T. Henighan, N. Joseph, B. Mann, A. Askell, Y. Bai, A. Chen, T. Conerly, N. DasSarma, D. Drain, D. Ganguli, Z. Hatfield-Dodds, D. Hernandez, A. Jones, Kermion, L. Lovitt, K. Ndousse, D. Amodei, T. Brown, J. Clark, J. Kaplan, S. McCandlish, and C. Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. wtwp:/te.ppp//2222///aaaeeee/ indw/xiiixx..... html
URL https://transfoone
L.. Elman. Finding structure in time. Cognitive Science, 14-21179-211, 1990. URL https //oniinelibrrry wiley. cwm/ ht////di/////////// 10. 1
F. Falck, Z. Wang, and C. C. Holmes. Is in-context learning in large language models bayesian? A martingale perspective. In R. Salakhutdinov, Z. Kolter, K. Heller, A. Weller, N. Oliver, J. Scar- lett, and F. Berkenkamp, editors, Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 12784-12805. PMLR, 21-27 Jul 2024. URL http::t/procoeeeeeeeeee
mlr. press/ressvve ffl//2aakk2.aa. html
L. Feng, F. Tung, M. O. Ahmed, Y. Bengio, and H. Hajiminsadeghi. Were rnns all we needed?, 2024. URL https: //arxiv wtrpaiiwww..a.aaiaaee.eooeoooeee abs/ 242.000200000
S. Garg, D. Tsipras, P. Liang, and G. Valiant. What can transformers learn in-context? a case study of simple function classes. In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS '22. Red Hook, NY, USA, 2022. Curran Associates Inc.
ISBN 0781713871088. URL httc:////ppparreeeeeeeeee//aaaap
M. Grendar. Entropy and effective support size. Entopy, 8(3))169-1114444 2006. ISSN 1099-4300. doi: 10.3390/e8030169
httt.hhtps://// //ww www.:.dp..... http:///////99440///////////
A. Gu and T. Dao. Mamba: Linear- sequence modeling with selective state spaces, 2024. URL
petps://aariis org. abs/ 2212 00752
S. Guo, J. B. Wildberger, and B. Schölkopf. Out-of-variable generalisation for discriminative models. In The Twelfth International Conference Learning Representations, 2023. URL
S. Guo, V. Tóth, B. Schölkopf, and F. Huszár. Causal de finetti: On the identification of invariant causal structure in exchangeable data, 2024a. URL /tttps://arxiin org/oogg///o/oe als/aa/3../55//5//
S. Guo, C. Zhang, K. Mohan, F. Huszár, and B. Schölkopf. Do finetti: On causal effects for exchangeable data, 2024b. URL https://arxiv http:///aaaxiivorrra/////2000/ 18836
M. Henaff, A. Szlam, and Y. LeCun. Recurrent orthogonal networks and long-memory tasks. In
M. F. Balcan and K. Q. Weinberger, editors, Proceedings of The 33rd International Conference Machine Learning. volume 48 of Proceedings of Machine Learning Research, pages 2034- 2042, New York, New York, USA, 20-22 Jun 2016. PMLR. URL htRL:hhps:///pppt htttpp:///proceedings. mlr
press/v48/henaff16. html
M. A. Hernan. Causal Inference: What If. Taylor & Francis, Boca Raton, 2024. URL https //mi mmiguelhernan.co org whotifbook/
S. Hochreiter and J. Schmidhuber. Long short-term memory. Neural Computation, 98):
1997. URLLhttps: h/www.///www... bioinf. jku. http://bbaaaaatooooo////////////
pdf
D. Kingma and J. Ba. Adam: A method for stochastic optimization. In Y. Bengio and Y. LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. http::t::://arre //rivvv.///iv.///////// www.wwwwwww....eaaae..e.ee org htgggabs//4122 6980
12

---

A. Klenke. Probability Theory: A Comprehensive Course. Springer, 2008. URL https://limg springer. wwom/boo///// 10000007797777 84800-048--3
M. J. Kusner, J. Loftus, C. Russell, and R. Silva. Counterfactual faimess. Advances in neural information processing systems., 30, 2017. URL http///pps////ooccoonn....o pdf
cc/paper
J. Li, L. Yu, and A. Etringer. Counterfactual reasoning: Testing language models' understanding of hypothetical scenarios. In A. Rogers. J. Boyd-Graber, and N. Okazaki, editors, Proceedings of the 61 Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pages 804-815, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10.18653/v1/2023. acl-short. URL
acl- short. 70/
X. Li, T.-K. I.. Wong, R. T. Q. Chen, and D. Duvenaud. Scalable gradients for stochastic differential equations. International Conference Artificial Intelligence and Statistics, 2020. URL https påf
L. Lorch, A. Krause, and B. Schölkopf. Causal modeling with stationary diffusions. In S. Dasgupta, S. Mandt, and Y. Li, editors, Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, volume 238 of Proceedings of Machine Learning Research, pages 1927- 1935. PMLR, 02-04 May 2024. URL httppss://poccceddnnnggeo/// mlr. weal..prees///2/1orcrhe html
K. Lorenz. Die Rückseite des Spiegels Versuch einer Naturgeschichte menschlichen Erkennens. Piper, München [u.a, 2. aufl. edition, 1973. ISBN 3492020305.
I. Loshehilov and F. Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019. http:psss//ppp.r.eeeee.eeeeeeeee
A. J. Lotka. istry, 14(3):271-274, 03 1910. doi: 10.1021/j150111 115000111000,,..0000000000000410 URL https://s//di...... otg//...102220 1021/ j150111a004
Contribution to the theory of periodic reactions. The Journal of Physical Chem-
C. Lu, B. Huang, K. Wang, J. M. Hernández-LLobato, K. Zhang, and B. Schölkopf. Sample-efficient reinforcement learning via counterfactual-baseee data augmentation. In Offline Reinforcement Learning- Workshop at the 34th Conference on Neural Information Processing Systems (NeurIPS), 2020. BRLLttts://ooflneeert -e-uuriip gwt.hb..tuuuu.... io/pdf /34 pdf
G. Maruyama. Continuous markov processes and stochastic equations. Rendiconti del Circolo Matematico di Palermo, 4(1):48 1955. doi: 10. 10.11007BFF2844602 httpstss::///o.. org/
10. 1007/8F02846028
T. Mesnard, T. Weber, F. Viola, S. Thakoor, A. Saade, A. Harutyunyan, W. Dabney, T. S. Stepleton, N. Heess, A. Guez, E. Moulines, M. Hutter, L. Buesing, and R. Munos. Counterfactual credit assignment in model-free reinforcement learning. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 7654-7664. PMLR, 18-24 Jul 2021. URL https:://ppoceeedingg
mlr. prress/v19// /tt////ennnnnaaaaaeee html
P. Mondorf and B. Plank. Beyond accuracy: Evaluating the reasoning behavior of large language models - a survey, In First Conference on Language Modeling, 2024. /tRLltttps///operevviie
net/forumiidd-mi Lm
Mooij, D. Janzing, and B. Schölkopf. From ordinary differential equations to structural causal models: the deterministic case. In A. Nicholson and P. Smyth, editors, Proceedings of the Twenty- Ninth Conference Annual Conference on Uncertainty in Artificial Intelligence, pages 440-448, Corvallis, OR 2013. AUAI Press. URL http:://ww: is tuebingen. dpw..d//ileddd.i.ii.... user
uai uai.pdf
N. Nanda, L. L Lieberum, J. Smith, and Steinhardt. Progress measures for grokking via mech- anistic interpretability. In The Eleventh International Conference on Learning Representations, 2023. http://////ppeeeeeeeve/ee///////
13

---

A. Nast- Hasr-Esfahany. M. Alizadeh, and D. Shah. Counterfactual identifiability of bijective causal models. In A. Krause, E. Brunskill, K. Cho, B. Engelhard,, S. Sabato, and J. Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 25733-25754, PMIR, 23-29 Jul 2023. JJRl.https:/prrceeddig
mlr. press/ https///v22nnnassrrree esannayy2aaaaaaaaa htmml html
C. Olsson, N. Elhage, N. Nanda, N. Joseph, N. DasSarma, T. Henighan, B. Mann. A. Askell, Y. Bai, A. Chen, T. Conerly, D. Drain, D. Ganguli, Z. Haffield-Dodds, D. Hernandez, S. Johnston, A. Jones, J. Kernion, L. Lovitt, K. Ndousse, D. Amodei, T. Brown, J. Clark, J. Kaplan, S. McCandlish, and
C. Olah, In-context learning and induction heads, 2022. URL BtL/http::/// http://xxivv.oooo hwt.://g///////////// abs/2299 11895
K. Park, J. Choe, and V. Veitch. The linear representation hypothesis and the geometry of large language models. In Proceedings of the 41st International Conference on Machine Learning, ICML'24. JMLR.org, 2024.
A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Köpf, E. Yang. Z. DeVito, M. Raison, A. Te- jani, S. Chilamkurthy, B. Steiner, L. Fang, Bai, and S. Chintala. PyTorch: an impera- tive style, high-ppeformance deep learning library. Curran Associates Inc., Red Hook, NY. USA. 2019. URL httpps: neurips....p........ wwtpcppperfttieeeeeeeeeeeeeeeeee files/ wwwesfiapeeeeeeee /aper//0199/f1le/ 0bbba2/8fe111122225557777777---0
pdf
J. Pearl. Causality: Models, Reasoning and Inference. Cambridge University Press, 2nd edition, 2009. Utt.httss:////// pttttps:/bbyes...e www.. edui hidu/B0O--2-5
J. Peters, D. Janzing, and B. Scholkopf. Causal inference on discrete data using additive noise models. IEEE Transactions on Pattern Analysis and Machine Intelligence, 33(12):2233-2444 2450, 2011. URL httts://ieeexpl.ooo
ieee. org/ wtrg/document/5700 574mm228
J. Peters, J. M. Mooij, D. Janzing, and B. Schölkopf. Causal discovery with continuous additive noise models. The Journal of Machine Learning Research, 15(1):2009-200 2053, 2014. URL https //jmlr. otw/parerrsaee...... caapperr//oluuuue wolume 15/ 1
Peters, D. Janzing, and B. Schölkopf. Elements of Causal Inference: Foundations and Learning Algorithms. The MIT Press, 2017. ISBN 0262037319.
J. Peters, S. Bauer, and N. Pfister. Causal Models for Dynamical Systems, page 671-690. Association for Computing Machinery, New York, NY, USA, edition, 2022. ISBN 9781450395861. URL http:::/ddo..org// 10. /145/3501714.3500
Petty, S. Steenkiste, I. Dasgupta, F. Sha, D. Garrette, and T. Linzen. The impact of depth on compositional generalization in transformer language models. In K. Duh, H. Gomez, and S. Bethard, editors, Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pages 7239-7252, Mexico City, Mexico, June 2024. Association for Computational Linguistics. doi: 10. longl-long.e 402/
URL https: http::://aacaaattoooo org/
L. Qin, A. Bosselut, A. Holtzman, C. Bhagavatula, E. Clark, and Y. Choi. Counterfactual story reasoning and generation. In K. Inui, J. Jiang, V. Ng, and X. Wan, editors, Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP) pages 5043-5053, Hong Kong, China, Nov. 2019. Association for Computational Linguistics. doi: 10.18653/w1/111-155000 URL https: aclthooooyy wrrg/orrgg.../ -
A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language mod- els are unsupervised multitask learners. arXiv, 2019. URL https://cdn.opr: http//bbtttrrllnuggggggeee learners. pdf OpenAl Technical Report.
are_ unsupervised_mulle lecddddduuttttttteeee
S. Ravfogel, A. Svete, V. Snæbjarnarson, and R. Cotterell. Gumbel counterfactual generation from language models. The Thirteenth International Conference on Learning Representations, 2025.
URL httpp:://opennreriees. http://www.......eoooootoootora.
14

---

P. Reizinger. S. Guo, F. Huszár, B. Schölkopf, and W. Brendel. Identifiable exchangeable mechanisms for causal structure and representation learning. In The Thirteenth International Conference on Learning Representations, 2025. URL
A. Sauer and A. Geiger. Counterfactual genirtttieaggeeeeeeieeeeee networks. In International Conference Learning Representations, 2021. URL http://////opeevveeeeeeee/////o/
B. Schölkopf. Causality for machine learning. In Probabilistic and causal inference: The works of Judea Pearl, pages 405-804. Association for Computing Machinery, New York, NY, United States, 2022. URL httpsttss:/// //dl. acm. wrg//oi/1/..ooo 1145/3501714. 3501755
C. E. Shamnon. A mathematical theory of communication. The Bell System Technical Journal, 27 (3):3794423, 1948. UtL//attss:/ppeopleeee math. harvard edu/ ctm//home/teex///hhrrx
shannon/ aatmmon/entrooy/entzroo.o pdf
V. Southgate and A. Vernetti. Belief-based action prediction in preverbal infants. Cogni- tion, 130(1):1-10, 2014. https//s:://ww htt./tttps//w.sssimmeerr..... http://ww..aaaeeeeeee.ec//////as
N. Tandon, B. Dalvi, K. Sakaguchi. P. Clark, and A. Bosselut. WIQA: A dataset for "what if..." reasoning over procedural text. In K. Inui, J. Jiang, Ng. and X. Wan, editors, Proceedings of the 2019 Conference Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLL) pages 6076-6085, Hong Kong, China, Nov. 2019. Association for Computational Linguistics. doi: (0.18653/1//19911660 URL bectps://aclanthot
org.oog////... 169-
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polo- sukhin. Attention is all you need. In Proceedings of the 31st International Conference on Neural Information Processing Systems, NIPS' 17, page Associates Inc. ISBN 9781510860964. URL
Red Hook, NY, USA, 2017. Curran www.iiee..ppp... cc/paper_ Paperr
htles//aaprr/201/ffttaa 0tt.://0/00011112022666666611160
J. von Oswald, E. Niklasson, E. Randazzo, J. Sacramento, A. Mordvintsev, A. Zhmoginov, and M. Vladymyrov. Transformers learn in-context by gradient descent. In Proceedings of the 40th International Conference Machine Learning, ICML'23. JMLR. org, 2023, URL mlr.
oatpa//asaaaaaaaaaaaooooaaaaaaaa oww.ld23a. pdf
S. Wachter, B. Mittelstadt, and C. Russell. Counterfactual explanations without opening the black box: Automated decisions and the gdpr. Harv. JL & Tech., 31:841, 2017. URL https://arxiv 17399909991119999
J. Wei, Y. Tay, R. Bommasani, C. Raffel, B. Zoph, S. Borgeaud, D. Yogatama, M. Bosma, D. Zhou, D. Metzler, E. H. Chi, T. Hashimoto, O. Vinyals, P. Liang, J. Dean, and W. Fedus. Emergent abilities of large language models. Transactions on Machine Learning Research, 2022. ISSN Survey Certification.
8836.
W. Xiao, H. Zhao, and L.. Huang. The role of diversity in in-context learning for large language models, 2025. URL https://atxxiii www./og.//////////oo//..oe..ee// abs/2b055 /1////555.11116
S. M. Xie, A. Raghumathan, P. Liang, and T. Ma. An explanation of in-context learning as implicit bayesian inference, 2022. URL htt.:httts////aat /ttpp:::////xi....//// wwwpsiiivvrrraeoooooe wtt.a//../.....e.eee.e...e..eeee
H. Yan, L. Kong, L. Gui, Y. Chi, E. Xing, Y. He, and K. Zhang. Counterfactual generation with identifiability guarantees. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
N. Ye and H. Namkoong, Exchangeable sequence models quantify uncertainty over latent concepts, 2024. TRLp:ttts:///rx/ wwtexivi.rrg/aae a20/ 03307
L. Zhang, R. T. McCoy, T. R. Sumers, J.-Q. Zhu, and T. L. Griffiths. Deep de finetti: Recovering topic distributions from large language models, 2023. htLpst :/arxii////xvv.. wrg./oor/////aaeeoooo 14226
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The introduction provides bullet point list of contributions with direct reference to the corresponding claims the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
Limitations are discussed in the final section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model sell-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
16

---

Answer: [Yes]
Justification: All statements in the main text are proven in the appendix. Lemmata in the appendix either proven directly accompanied by reference proof.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proof' in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Multiple appendix sections discuss model setup, training details and the synthetic dataset.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
17

---

reproducible, they Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
Anonymized code is provided in the supplemental material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https /nips
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Technical details and experimental setup for both settings are discussed in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Where applicable, all plots supported by basic bootstrap errorbars. For readability, some plots containing error bars are deferred to the appendix, with the figure in the main text missing error bars. The reader is explicitly referred the appendix these cases.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
18

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: This can be found under the training and model details in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics httt:s///heurrps..cc//olllbtcttt
Answer: [Yes]
Justification: This study does not involve human research subjects and experiments are run synthetically generated dataset. Therefore, data security, copyright biases not imperiled here.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: This paper contains Broader impacts paragraph discussing the consequences of potentially malicious use.
Guidelines:
19

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymentt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pre-trained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper uses a synthetically generated dataset.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The existing asset used is an established open-source codebase. The paper recognizes the authors of that codebase.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
20

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
If assets are released, the license, copyright information, and terms use in the datasets.paperssiii paperswi thcode. www.com//aaaa datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
This is an experimental study on a synthetic dataset.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
21

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: No LLMs are used for research-related purposes.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should not be described.
22

---

A Posterior predictive distribution
Proof of posterior predictive distribution. Let 6 C I B R and B E B. Define the observational
data x (21, 1/, ... Bnz yo) for fixed n. Then, with query == (4) M1 m. Za ya 2, x: CF).
Pvg (v=
CF qqurry
Pyg ( Iquery,
x, 0, CF Pyg (13
=)d0
PYg Vy=
x, 0, CF (2)
(8)
PYg (3 Cr x.) + Uz x, ,, or 2, (0 x. (2)
e B
dF
Pycr (3(x CF xx) --xx))
+ yı T.. U2. 0, CF (3)3)
(9)
8(Y -3x ys e
where 8 follows from independence, as (X Z). and ( by de Finetti (de Finetti Klenke 2008 each sequence {(z1, v1),..., (En, Un)) forms exchangeable unit.
1931
Mean and variance linear additive frumework. Both follow immediately by the tower law (H) of iterated expectations.
yy
B[3x
Uv
E E 38
E X
E E [Uy 0]]
[[
Var(ycr)
[F E (8X + Uy)
E ((X 2 + U3
" E [39] [P) 2E [30y|] [X [U}]
[ [ Var (x + E [U3
= E [1 + 02 (Var (x +
(1 Var (0) 1) 132 169.
The proof for the variance under the multiplicative extension, Vy) analogous.
B Transformation lemma and identifiability
B.1 Transformation lemma
Proof of Lemma As T is invertible in given f(x),
= = (
We require injectivity in to uniquely determine from (f(z), y). Else, y T(f u) is ambiguous and counterfactuals are ill-defined. Thus, the counterfactual completion writes
yCF h (,, f(x),,))
for h x> x y.
23

---

B.2 Counterfactual identifiability of bijective causal models
To prove Theorem we restate Definition 6.1, Proposition 6.2, Lemma B.2 and Theorem 5.1 from Wasr-Estahany et al. [203].. For brevity, the proof' are omitted. We begin by introducing the Bijective Generation Mechanism (BGM) our notation. Let (X.,... be sequence of endogenous random variables. Each jj. J E [d] is generated by
Xj = ʃg (Pa(X)).Ug))
Importantly, fy is a bijective mapping from U, to X, for each realization of Pa(X,) with Pa(X;) denoting the causal parents of Xj. In the bivariate case, write
Y = / (Pa(Y), Uv) Uv)
for f bijective, as described in Lemma Thus, with causal graph X - Y. we have Pa(Y) {X), Pa(X)= Definition 2 (Definition 6.1, Equivalence, Nasr-Esfahany et al. (002]). BGMs fi and f2 are equivalent iff there exists an invertible function g(-) such that
Va, y fT (x, (J2'(zx (a. v) or alternatively
Vx, u (i = f2 x.g (u)
Proposition (Proposition 6.2, Easr-aafa et al. counterfactuals > f1, f2 are equivalent BGMs.
2023 BGMs fi, J2 produce the same
Lemma 2 (Lemma B.2, Nasr-Esfahany et al. (2023
BGMs f and f that produce the same
distribution PD Y) are equivalent if
(Markovian) U T1 X and Ú IL X.
2. for all x. (xx and fax, :) are either strictly increasing or strictly decreasing functions. Theorem 2 (Theorem 5.1, Nast-Essfhan et al. (2023). BGM f is counterfactually identifiable given Fx,y if
1. (Markovian) U IL X.
2. for all x, f(x,,)) i) is either a strictly increasing or a strictly decreasing function
B.3 Prove counterfactual identifiability under exchangeability
Theorem (Counterfactual identifiability under exchangeabilityy. Theorem Let X, Y,U, 0 scalar
random variables with X IL U8.. Take T y x l y with y (Tf()) Assume Vf(x) E V. the inverse T (f()),) exists for all y. i.e., = T (f y), and vf(x) E Y, T(f(x),) is continuous. Suppose further that [(x) continuous Vx, and strictly monotonic æ. Then, set y= TTT((XX) Given the joint law Px,y T is counterfactually identifiable.
Proof of Theorem We apply Lemma 2 and Theorem 2 East-asfa al., 2023 the case of conditional independence. The proof is analogous by taking Qx == Fx,y Clearly, the joint satisfies Markov conditional on the latent 0 by design. Next by inverse function theorem, (T)(),,)) is strictly monotonic, vz(x), satisfying condition 2 in Theorem 2 With f(x) strictly monotonic, continuous, T' == T(f(x)) :) satisfies condition for r.
C Training details
To start, we introduce notation [ : ().....)) for instance, [n] := (1...., n). We construct synthetic datasets with fixed noise. Conditional a uniformly distributed latent parameter 0 θ we draw n, ~ U(12, ..., 50}) observational data points with noise U 00, U 00 E R E [n.] and weights 000 E R E [N B] to enforce non-i..dd data. We choose E 5. Next, we sample N. B data points with
6(
3 N( C ~ u([66, 6
(10)
U UJ0~N((011E
24

---

for N = 50'000 training steps of batch size B = 64. We thus have sample size N B. We set
Y - B + Ux + Uy
with @ denoting element-wise products to have input and output embedding dimension agree without zero-padding. Including the indicator token Z/ = 26 E 26 E {1,..., N.). the corpus of queries then consists of N batches of the form
{(x6:1 Vb;1, Xb:ns Vo Zb. X
CF
on which the model has to predict y Bizy E [B]. Our target yCF is zero-mean with Var(YCF) 169 and log((orr)]g (Var(((CF)) 5.1299.
CF
We train minimizing the per-batch mean squared error (MSE) between the counterfactual predic- tion y(e) and the ground truth y [18)
MSE
1
B
yor
CF bisb
|
[B]
B
E
b=1
and evaluate on an unseen test set following (TO).
Our code based on the repository by Garg et al. 2022]. We therefore adopt the learning rate of 10 for all function classes and models but use the AdamW optimizer [Loshchilov and Hutter 2019 instead of Adam [Kingma and Ba 2015). We implement the experiments in pytorch (Paszke al. 2019] and use one NVIDIA GeForee RTX 3090 GPU for training. The conducted experiments require between minutes and hours training depending on model setup and task complexity.
D Model details
The Transformer architecture is described in Section 2.2 Here lay out the details the three recurrent models used and investigate the relevance of model depth in Transformers more closely.
D.1 Model details on RNN architectures
The Elman RNN [Elan 1990 a foundational recurrent neural network architecture. It consists of an input layer, a hidden layer with recurrent connections, and an output layer. The hidden state at time step I, denoted by h, computed as
h; Maan(WWnar + Wh + bn))
y/: Whoh, + bo.
where is the input time i, and tanh denotes the tangens hyperbolicus, non-linear activation function. The output is given by y/1 and W., b. represent learnable weight matrices and biases.
The Long Short-Term Memory (LSTM) network [Hochreiter and Schmidhuber 1997] addresses the vanishing gradient problem by incorporating memory cell and three gates: input (ie), forget (ft), and output (ot). These gates regulate the flow of information and enable the network to retain long-term dependencies. The LSTM cell is defined by the following equations,
ft a (Wyat + Uphi-1 1 + by)
4 = (Wex, + U,h,-1 + b,)
0t (Wort + U,ht-1 + bb)
St Lanh(W g Ught- bg)
C: ft O c-- + i @ gr
h @ tanh(c.).
In this formulation, is the cell state and h, the hidden output state. Activation vectors of the three gates are given by 4,, fe, or and g the cell input activation vector. By a, we denote the sigmoid function and element-win wise multiplication by @. U. represents another weight matrix.
25

---

The Gated Recurrent Unit (GRU) [Cho et al. 2014) introduces gating mechanisms to better control the flow of information. It simplifies the LSTM architecture by combining the forget and input gates into a single update gate, žς. The GRU computes its hidden state using the following equations,
F1 (W,rz + - bz)
2t o o 24 + U-h4-1 by) tanh (W, ( 4 + (r: @ 1)) + bm)
h, (1 21) o n 2 ① h1-1-
Here, rt is the reset gate, 21 the update gate, n4 the new gate, the candidate update vector, Note that the dependence of the hidden state on the update gate varies across documentations. We follow the definition in the pytorch package.
We perform a hyperparameter sweep on hidden size D and the number of layers over the grid D (64, 128, 256), L (1,2). We select the configuration with hidden size 256 and layers for each of the three architectures.
D.2 Copying in RNN architectures
In general, we study whether neural sequence models are able perform counterfactual prediction. We observe that Transformers do so, but also RNN-type sequence models achieve low error. Figure 2 displays this finding for STANDARD GPT-2 Transformers as well as LSTMs, GRUs, and Elman RNNs. In their seminal paper, Hochreiter and Schmidhuber [1997] lay out the copy memory task for their paper,/Hochreeee developed method, the LSTM. Arjovsky et al. [2016] address vanishing and exploding gradients by modifying RNNs and Henaft al. [0166] consider LSTMs and RNNs under longer contexts. Finally, Feng et al. [024 test parallelizable modifications of LSTMs and GRUs on the Selective Copying Task [Gu and Dao 2024). Given this extensive literature, restrict most of our analysis to autoregressive Transformers, as these represent the state-of- totte--r-the-rt architecture.
D.3
Varying depth for additional values
Full
Full
AO
AO
In-context MSE (log)
In-oontext MSE (log)
MSE
1ayeer, 2-ayer, 4-ayer, 6-ayer,
1-ayer, 2-layer, 4-ayer, 8-ayer, Model
Model
(a) Varying depth for two in-context examples.
(b) Varying depth for forty-five in-context examples.
Figure 7: Varying depth more closely. Building top of the results displayed in Figure 3b compare Full and AO Transformers at a constant number of attention heads. We observe a decrease in in-context loss model depth increases across shorter and longer contexts of and 45 in-context examples, respectively. We evaluate on 6400 sequences.
Analogous to Figure 3b, note that increasing model depth leads declining in-context MSE. This holds for contexts both shorter and longer than the 35-example version considered above. Figure iillustrates the in-context MSE for varying depth, evaluated and 45 in-context examples, respectively. We note that the loss of the 1-layer, 8-head AO Transformer exhibits no significant differences between in-context examples and 45. We interpret this as indication that the model is unable to infer information on the latent 0 from the context. Reinforcing our findings that the 8-layer, 1-head Transformers achieves lowest in-context MSE, this pattern serves as evidence that in-context inference occurs across layers. A contrasting notion may be that the Transformer maps different junks of information different subspaces of the embedding space. Acting on each subspace individually, each attention head would then focus different task before the MLP would
26

---

combine the information and output the final prediction. Instead, the model appears to benefit from the interplay between layers, which supports the claim that attention heads across layers communicate through the residual stream [Elhage et al. 2021
E
Data diversity and robustness
E.1
Data diversity adjustments and out-of-distribution
In Section we discuss data diversity and generalizability to OOD. The UNIFORM setup largely follows the overall setup described in (10) with the following adjustment: instead of N B F,, we sample d < N realizations of the latent and compute the effective support size [Grendar 2006]. Note that we control the absolute number of mee-imensional realizations as our setup effectively analyzes E independent examples at every turn. Thus, we draw d realizations and allocate across iterations and embedding dimensions. The parameter d corresponds to the di diversity argument in our code.
The NORMAL setup follows the above structure. Compared TT. we sample the latent from (O, /12 Ig) such that the theoretical variance is equivalent across setups. Everything else analogous. Note that the effective support size is equivalent to the support size under the UNIFORM setup. This can be seen by interpreting Ess as the logarithm of the Shannon entropy [Shannon] 1948). Indeed, p(0) is uniformly distributed over discrete Oo if i denotes realization of the latent 0,
H HUNIFORM (0) = Σ JE@O
log p(0).
Therefore, we allocate the d = 10ol realizations uniformly across iterations and embedding dimen- sions. Under the NORMAL setup, p(1) depends on the distance of e to O. Here, we sample d Gaussian realizations. We then allocate them proportional to p()) across iterations and embedding dimensions. In analogy to Figure 5al we plot the in-context MSE against Ess for pre-training setting and evaluating both in-distribution and on UNIFORM. Figure Nillustrates the in-context MSE relative to effective support size contexts of two examples. Here we train all models for 50'000 training steps. Although the loss is at a higher level than for 35 in-context examples, we observe that the MSF declines with increasing effective support size. Moreover, when evaluated under the NORMAL setup, the model converges more rapidly than under the UNIFORM evaluation scheme. At d 3. the Transformer already achieves performance comparable to that obtained when training on 1000 unique realizations. For the latter case, the results remain consistent with the findings above, indicating that an effective pre-training support size of approximately 10 is sufficient.
the NORMAL
distribution
Training distribution
Normal (OOD)
Normal (in-daxt) Unitormm (OO0)
Unitorm (In-diet)
Incorted MBE (1og)
Incortae ozt MSE ([ng]
support azee (log)
Effective support (log)
(a) Data diversity for evaluation on UNIFORM.
(b) Data diversity for evaluation on NORMAL.
Figure 8: Data diversity for pre-training on two in-context examples. We measure in-context MSE scales) averaged over 6400 prompts at two in-context examples against log-scaled effective support size. Each point represents one fully pre-trained model on either the UNIFORM sampled or NORMAL sampled We train the models for 50'000 training steps. All models are evaluated on both datasets. Error bars are the 95% basic bootstrap confidence intervals.
E.2 Non-linear and non-additive extensions
On top of the linear regression setting, we extend in-context counterfactual reasoning to non-linear, non-additive functions which are invertible in The general dataset setup remains the same with
27

---

adjustments made only to the computation of sF We evaluate the 8-layer, 1-head Full and AO Transformers, (Talle reports the in-context MSE on 6400 queries with 35 in-context examples each. Depending on the exact configuration of architecture and extension, the in-context MSF is at least one tenth of the empirical variance of the test set.
[wx u)
Empirical variance
MSE of AO Transformer
MSE of Full Transformer
Tanh
tanh (T(Bx u))
0.3404
0.0132
0.0058
Sigmoid (a)
0.0387
0.0019
0.0009
1-exp(-7 (Bx++u)
Multiplicative
Bru
0.9274
0.0668
0.0364
Var(BXGFUy
Table Robustness different function classes for regression tasks. We report in-context MSE averaged over 6400 sequences for the 8-layer AO and Full Transformers under invertible nonlinear activation functions and multipiccative noise. We compute the empirical variance of target Y on the test set. We observe for both architectures that the MSE is 10 to 60 times lower than the empirical variance, suggesting both can in-context learn more complex function classes beyond linear regression and complete the counterfactual query. All functions are applied element- in one dimension. lanh and a are scaled by r = I3 to counteract congestion around (-1,1} and (0, 1). respectively. To guarantee theoretical variance of 1, we divide the multiplicative term by v /Var(8XFFyy)+ /3410.4.
Note that under multiplicative noise, counterfactual reasoning reduces pure copying task as VVar(YCr)) excr Y VVr((F)
YCF
1
3x
X
Y
Var(Y
8X
X
where estimation of 3 not required for counterfactual completion.
E.3
Data diversity in the literature
Recent theoretical research on data diversity aligns closely with our observations, Xiao et al. 2025 develop formal framework showing that diversity in example selection improves in-contexi inference. With respect to model depth, Petty et al. [2024) show that at least two attention layers are required to support compositional generalization beyond memorization. This is in line with our findings. Figures abdand 7b ib] that the additional benefit of deeper models decreases after having trained models of depth 2. Improvements are especially high for layers relative
Intuitively, we relate this to the multi-step process of performing counterfactual prediction. The model retrieves the relevant pair, (x.5.0) Uj) for = j. infers noise, and predicts the counterfactual CF given CC Figure
E.4 Emergent learning of variable token
We confirm in Figure that the model is capable of learning the latent parameter context. We therefore find evidence that the model learn the prior over the latent information hypothesized in (2). The final part toward in-context counterfactual reasoning pertains noise abduction. Our setup special
include embedding token z which indicates the relevant in-context pair (xz, Vz) to abduct noise from. Only the model can learn this connection on the fly, we can guarantee the three-step process of noise abduction, intervention, prediction. Indeed, we find that the Full GPT-2 Transformer with 12 layers and 8 heads trained for 1'000'000 steps on E = 1 implements noise abduction head at the 7 layer. We define noise abduction head as one which attends the relevant input-outpu pair (æz, yz) when processing the 2 token. This noise abduction head emerges at the sixth head of the seventh layer, and attends to ys for variable realization of z. Visually inspecting the attention behavior of 100 batches averaged over 64 in-context sequences, we find that this arises every single batch. Note that the token is constant within one batch. In Figure 9 we show the behavior for four different indices. For swift comparability, each prompt consists of 50 in-context examples. We conclude that noise abduction, central to counterfactual reasoning, emerges in designated attention head.
28

---

wa
how
Attended-to keys
vou Voiinaa via whe whe wir sas
Attended-to keys
baavie tarii vi dá has the va loe las via vi nhe vie van has vis vos vis voe Vis rie ier Held Inic Attended-to keys
Attended-to keys
0.0
0.2
0.4
0.6
0.8
1.0
Figure 9: Noise abduction head emerges in GPT-2 Transformer. We plot 4 representative noise abduction heads evaluated sequences with 8, z 34, z 15, 42, respectively. Each graph represents the average over 64 sequences with identical index and 50 in-context pairs. We find that when processing embedding token Z, the model adapts to different values of
F,5
Comparison non-counterfacttual regression
After reviewer feedback, we include a brief comparison of our results with training models on an observational continuation of the regression setup. Linear regression, as studied by Garg et al. [002), does not require the ability to perform noise abduction. We take one step further by requiring the model learn the functional form f in context and to abduct the unobserved noise. To underscore 1'000'000 training steps. Figure 10a represents
our point, train a Full GPT-2 architecture the relative performance improvement of the counterfactual sequence relative to the observational setup. We find that for over 5 in-context examples the counterfactual MSE lies significantly below estimated MSE the observational sequence. In Figure 10b we observe the phase transition to significantly lower in-context training MSE occurs after 300'000 steps. This is in line with hypothesis that in-context counterfactual reasoning, in principle, leads to a lower prediction error than the standard regression setup Gaag et al. 2022 the noise is fixed and can be inferred directly.
F Model details for cyclic sequential dynamical systems
F1
Lotke-Volterra model
The Lotka- Voltera framework models the temporal evolution of the concentrations of two interacting species: predators and prey. the absence of predators, the prey concentration grows continuously over time. Conversely, the predator concentration declines unless sustained by consuming prey. Predation enables predator growth, which simultaneously reduces prey concentration. We model the four individual dynamics by real, nonnegative parameters. The dynamics of the prey concentration (6) parameterized by the intrinsic growth rate a and the rate at which the prey concentration decreases with predator concentration, B. For predator concentration 7 we capture the species-inherent decreasing concentration by 7 and the rate at which predators feed off prey by In the context of biological species, concentration can be thought of as the population density.
29

---

counterfactual (1m)
Sequence
In-corloxt training lots (log)
GPT-2 continuation (1m)
In-context MSE (log)
-20
2
10
20
30
40
50
200
400
800
800
1003
Number in coccoattox examples
Training ilerations (x10*)
(a) GPT-2 performance.
(b) Phase transitions.
Figure 10: Distinction between noise abduction and observational continuation. We measure in- context MSE (log-scaled) averaged over 6400 prompts for the Full GPT-2 Transformer. We observe sigmificant improvements for contexts longer than examples. This we understand as indication ef- fective noise abduction in the model. We train the model for 1'000'000 training steps and observe that the phase transition for the model trained on counterfactuals occurs after over 300'000 steps. The log- arithmic in-context training loss for the continuation framework remains higher level of above O.
F.2
Bounds on Voltera parameters
Similar to the regression setup, we draw the model parameters conditional on the latent 6. In order overcome exploding concentrations, we adopt the criiatio-oof-contaae method to bound the parameters by controlling the maximum and minimum concentrations. This ODE-based approach we utilize as guiding principle throughout. As stochasticity may lead to a case where bounds are violated, however, we posteriori ensure evaluation positive concentrations only. In the following, we adopt the ODE notation and write 6 and 7 with lower-case variables as
dx:
= a32 Brett.
(11)
dt
dy:
Y + 6z
(12)
dt
and observe that
(a - -Bye)ar = O, dæ.
dt
dy: (-y + 8xz )(r 0.
dt
For now, we confine ourselves to (II). Dividing by Tz, we obtain
dx./dt (a By))
(13)
Zt
with dæ./dt d log(x.). By the fundamental theorem of calculus, integrating (13 over [0, t] yields 21
d S (a By.) ds = S/¹ log(æ,,) )ds = log(x;) log(2o) (2o).
ds
By similar reasoning the result follows for 12 such that we can reorganize
at - B J y.ds)
Ft = To exp
at
(14)
- + So xzds))
yt = yo exp
(15)
Now, we set (I, (y, Ø) to be lower and upper bounds of It, Hr. respectively, & E [0, TJ. Then,
Jt
at B3yt 2 log
2 at Bitt
20
30

---

such that
It
Et
a
2
log
+
fy,
log
+ 5g 2 a.
20
As log
i 20
2 log
log
vt,
20
1
1 By
log
log
8y
t
20
J0
For E [0, TJ. the bound for lim 0 [2, 7 such that
trivial. We therefore decide to bound the final value at T in
T By.
1
T
log
+
log
+
3y
T
20
T
ID
To guarantee that the bounds are in fact lower and upper bound, we bound 3,
1
T
1
log
By S
log
Bÿ
T
T
1
log
S
(ū- 5)
T
log
8
T
y)
Note that stating an upper bound for 3 violates the model assumption 2 0. Similarly, we write for (15)
Y/
7t + Egt S
log
I -t
Szt
3/0
y
1
7 ōx <
log
log
I -y + 83
yo
yo
such that for T>0 O.
2
E 8τ
log
log
T
We, too, establish the bound by enforcing that lower and upper bounds are in correct order, i.e.,
y
1
y
8x
log
I 5x
log
T
V/u
T
Vuu
1
7
y
log
log
<
8 (5
x)
T
yo
34
1
log
<
8.
T
[
②)
y
Put together, have the bounds
1
log
+
By
< a
<
log
3g
ā.
370
T
30
1
E
0 =
log
S 3
T
y)
x
1
y
2 = dx
log
S
87
log
=: 7,
T
yo
T
yo
1
log
<
6.
T
(
y
31

---

F.3 SDE Parameterization
In accordance with our exchangeable setting, we construct the dataset with
f 2(11
U|0 ~ ~
X Yo ~((112)
(16)
and bounds on the variables (g. = (y. () (0.5, 2). We parameterize the Lotka-Voltera model,
β - _ ~ Exp(
8 δ Exp(0)
~ u al)
7 u ((7, 71)
Finally, we sample the initial condition (Xo, Yo) from Beta 2-8 2) distribution. Recall that the one-dimensional Beta (K, A) on the measurable space ([0, 1]. [ 1))) attains its theoretical mode -+--- for K, À > 1, B the Borel o-field. Here we deviate from the conventional parameterization Beta(o, (3) to avoid confiit with the Lotka-Volteree parameters. To enforce an initial condition depending on the latent concept, we choose (K. A) such that the theoretical mode equals Equip the measurable space (θ, (()) with probability measure with supp(r)) = [1, 2). This is analogous (T). To align the support of the Beta law with that of æ,, we require
at _--__
0 - k
K K X 2
(0 - 1xxx+>>>+>+- -- 2) K-1
0(A 2) +3 (2 0)
ER >>> 1. 0(A (A --3)
2-8 8
Now, À 1 holds automatically such that we can set arbitrary X and parameterize K A(> (A-(33) For straightforward implementation, we choose À = 2 such that k. 2-6 To con- clude, the initial value problem as well as the VolterVoltera parameterization depend on the latent 6.
2
F.4 Training details
We sample 20 distinct time steps Im ~ 44 (10, 0.5)), E [n]. In addition, we parameterize the Lotka- Volterra bounds defined above with = We again consider F independent sequences at once by stacking independent one-diment one-dimensional SDEs. Within one example, all SDEs are evaluated at the same distinct time steps. At any i, we then provide the observational sequence as
(IL, Vty. Yt, Z, X
for z counterfactual delimiter token, identical across i. Given the query, we ask the model to predict the full completion of the counterfactual sequence, comp, == (x y ..... x.., XCF autoregressively. We again evaluate the model on the per-batch MSE of the predicted completion, comp and the underlying ground truth sequence, comp)e).
y,.)
MSE
(CmmpI comp1e))
1
Σ
comp, comp,
B
and evaluate on an unseen test set following 16
We train for N 50'000 training steps at batch size B = 8 and numerically approximate the solutions to the SDEs in using the Euler-Maruyama [Maruyama 1955) scheme. The counterfactual sequence is generated analogously but by recycling the Brownian motion used for the observational sequence. The manifestation of this can be seen in Figure 6b for observational and counterfactual y. We inherit the torchsde package from Li et al. 2020 and follow Charleux al. (018) for the implementation
the Lotka- Volterra equations. All other specifications are similar to the regression setup described in inpppendix C For fast approximation of the SDEs, we generate 20 batches of size 2500 8 at once. Across one batch, (t,,.., tx } is constant. At N = 50'000 training steps and a resulting dataset sample size of 50'000 8, we hence train on 50'000 20 distinct realizations of the 2500
time sequence (.., ). this can be seen as curriculum learning, we do not fully adhere the aforementioned notion of exchangeability. Due to the complexity of approximating SDEs, however,
32

---

we give priority to an efficient implementation and acknowledge the resulting shortcoming, Last, the Muler-Maruyama method is an SDE upproximation on N equal subintervals of width At > 0. Therefore, we also train models with equidistant time steps {0, (1,..., (n) for constant = 20. For this setup, the observational sequence starts lo =0
(xo. Yo.
Z. X No
with completion comp == (xi Yt CT. ytt) and evaluation on n 2 tokens. Throughout our experiments, we state the differences between these two setups.
E55
Attention behavior
Figure 12ssswws that the 8-layer, 1-head AO Transformer trained N equal subintervals performs counterfactual reasoning by repeatedly moving information forward. Averaging over 8 test sequences, observe that the tokens of the counterfactual sequence aftend to the position of the delimiter token and the initial condition X at the first layer. This is intensified by layers 2 and 4 with attention to the initial x Note that we only plot attention values between tokens of the counterfactual sequence. The Transformer at the fifth layer implements an induction head [Olsson et al. [0222 that copies forward the information one step. Afterwards, attention head 6 implements decaying attention over prior posi- tions. Maximum weight is here put on the current token, and progressively lower weights on all previ- ous tokens relative token distance. This design facilitates forward information propagation across to- kens. Similar behavior can be observed for the first two layers with attention to tokens {x6] ..... x
In accordance with the hypothesis that induction heads drive in-context learning [Elhage et al. 2021 Olsson et al., 2022]. we notice this pattern for larger models, with multiple heads per layer, and including the MT.P. For instance, the AO STANDARD Transformer contains five heads across the first six layers which mostly attend to the delimiter token z. At the eighth and ninth layer. the model imple- ments each one induction head and one 2-gram [Akyürek et al. 2024 head in parallel. For the Full STANDARD Transformer, we similarly observe five copying heads which move information forward layers and When trained on variable sequence length, the 8-layer, 1-head AO Transformer also exhibits this pattern with an induction head at layer 7. Taken together, this provides more evidence that counterfactual reasoning emerges, at least partly, in the self-attention layer of the Transformer.
G Connection between synthetic setup and natural language
G.1 Discrete regression setup
Due to the synthetic nature of our setup, the connection to natural language is not immediately apparent. To illustrate how unobserved noise u can be interpreted linguistically, we return to the introductory example. Each pair (x,, M/) can be thought of as a pair of sentences describing a patient's 's treatment and effect. For j [m)] 4y captures the treatment patient j receives, and yj provides information on the effect of that medication. After presenting the model with n such in-context examples, we prompt it to predict the counterfactual effect had patient j received different treatment. The model therefore needs to disentangle the subject-specific noise uj from the observational pair (z)) Figure visualizes this process natural language.
In addition to the one one-prompt understanding, we can interpret the exchangeable setup such that each se- quence resembles different health issue. While we analyze breast cancer above, this enables training on a diverse set of treatment-effeett pairs in order to perform inference over an unseen clinical picture. Establishing such principled machine can assist informed decisions in personalized medicine.
G.2 Continuous SDE setup
The eyclic causal dependencies modeled by the SDEs in 44) align with the sequential structure of natural language. Each query can be interpreted as factual narrative of length fx. Given a hypothetical initial condition (= 36 (CF), we task the model with counterfactually completing the story while keeping the noise fixed. This noise may represent semantic variability as well as narrative- internal uncertainty. By doing so, the model can generate plausible counterfactual continuations conditioned on the observed prompt. This approach resonates with the work of Qin et al.) [2019], who explore how language models can reason about alternative story outcomes through counterfactual perturbations to character actions events.
33

---

107
layer
2nd layer
kif
query token
vir-
eer-
sure
ruse
Attended-to keys
Attended-to keys
3rd layer
4th layer
Queryy token
Attended-to keys
Attended-to keys
5th layer
utt-
6th layer
Token
sur-
Yacr
vite
Altended-to keys
Attended-to keys
34

---

7th layer
8th layer
10
query token
197-
ent-
4. nar
C. ran
Attended-te keys
Attended-lo keys
0.0
0.2
0.4
0.6
0.8
1.0
Figure 12: Cyclic causal relationship. The 8-layer AO Transformer, trained on counterfactual reasoning in Lotka- Lotka-Volterra SDEs, includes two dedicated copying heads. At layers and 7, induction heads shift residual stream information from the previous token forward one position. Earlier layers 1, 2 and 4 focus on the initial condition x and implement decaying attention over prior positions.
35


Reviews:
[{"review_id": "rmulTl0u6V", "reviewer": "Reviewer_3LJY", "summary": "The paper analyzes LLMs’ capabilities to perform in-context counterfactual reasoning. To this end, the authors prompt a model with what-if questions based on a numeric synthetic dataset using a function-learning probe. The authors focus on a linear SCM (Structural Causal Model) setting, working in the latent space, where they add noise to various stages of token representations.\nIn a GPT-2–sized Transformer, they show that counterfactual prediction copies the noise from the factual pair and applies a learned transformation. They empirically demonstrate that it uses fewer in-context examples than other architectures (e.g., LSTMs). The authors further show that performance degrades when self-attention is removed (using only the MLP), and that the mixture and number of SCMs influence pre-training. They argue that their findings also generalize to data mimicking counterfactual story generation.", "questions": "Have you conducted any preliminary experiments on real-world tasks to assess the generalizability of your approach?\n\nBeyond analyzing attention scores, have you explored other interpretability techniques, such as methods from mechanistic interpretability (e.g., probing the residual stream)?", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strenghts:\n\nThe authors focus on a very timely topic: counterfactual reasoning. The paper is generally well-structured and well-written.\n\nThe task and synthetic dataset are well designed for the research question and thoroughly developed: the task is methodologically well-formalized, and the experiments follow this design closely, answering relevant questions and presenting meaningful ablation studies. In this way, the authors present a sound theoretical framework to test SCMs in an in-context counterfactual reasoning setting.\n\nThrough this approach, the authors provide several insights—such as the importance of self-attention—which constitute a meaningful contribution to the growing body of work on reasoning in LLMs.\n\nWeaknesses:\n\nSince the approach is only tested on a synthetic dataset, its general applicability to real-world tasks remains uncertain. It would have been valuable to include an evaluation on at least one real-world dataset.\n\nFurthermore, the method is evaluated only on a comparatively small model.\n\nThe approach primarily builds on existing function-learning probing techniques, applying them to in-context counterfactual reasoning. Moreover, the task of answering \"what-if\" questions has been addressed in prior work. While the paper tackles a meaningful problem in a well-controlled setting, I find that the methodological novelty is rather modest.", "was_revised": true, "final_justification": "The authors provided detailed explanation on the theoretical approaches and also the comparison to existing work. While the contribution is mainly theoretical, the approach would benefit from evaluation on real-world examples and also on a larger scale. E.g., papers focusing on induction heads evaluate on real-world examples, such as on the IOI task [1], or [2] on factual knowledge; (I find the distinction between theoretical and method-oriented approaches not to be exclusive in recent interpretability research). Therefore, I will keep my original score that is leaning to accept the paper.\n\n[1] https://arxiv.org/pdf/2211.00593 [2] https://arxiv.org/pdf/2405.17969", "rebuttal": "We thank the reviewer for their thoughtful feedback on our project. We are happy to hear that the reviewer finds our method methodologically well-formalized and the topic timely. The concerns raised are mostly related to model size and dataset setup. Below, we provide answers to the questions raised and discuss some of the mentioned weaknesses.\n\nW1 *Since the approach is only tested on a synthetic dataset, its general applicability to real-world tasks remains uncertain. It would have been valuable to include an evaluation on at least one real-world dataset. Have you conducted any preliminary experiments on real-world tasks to assess the generalizability of your approach?*  \n**Our work  shows that noise abduction in neural sequence models is in principle possible.**  \nThe main purpose of our paper is to connect the theoretical result of performing counterfactual prediction to known properties of Transformers. Primarily, we are interested in confirming that Transformers are capable of performing noise abduction while learning a certain functional form in-context. Showing this theoretically through Lemma 1, we check whether this is in principle possible on a small scale. Note that is in contrast to a methods paper that proposes new algorithms.\n\nW2 *The method is evaluated only on a comparatively small model.*  \n**We adopt the standard approach to design a small-scale experiment to underscore our theoretical findings.**  \nFor a paper trying to provide theoretical understanding, we think it is preferable to train as small a model as possible while capable of exhibiting the desired performance (in our case, in-context counterfactual reasoning). This is similar to previous literature where the analysis is restricted to small models. For instance, [1] implement a two-layer attention-only Transformer with neither biases nor layer normalization. Similarly, [2] extrapolate their findings from an architecture with linear self-attention. In addition, our work is grounded in earlier findings on induction heads [3;4] and the in-context learning of functions $f$ [5], which, too, are derived from small models. As we are interested in showing that the discovered mechanism arises in-context, we do not intend to build a counterfactual reasoning machine ourselves.\n\nW3 *The approach primarily builds on existing function-learning probing techniques, applying them to in-context counterfactual reasoning. Moreover, the task of answering \"what-if\" questions has been addressed in prior work. While the paper tackles a meaningful problem in a well-controlled setting, I find that the methodological novelty is rather modest.*  \n**We introduce results from counterfactual identifiability to in-context learning and propose the idea to model natural language by stochastic differential equations.**  \nTo the best of our knowledge, existing work that addresses “What-if” questions focuses on direct evaluation of pre-trained large language models through question-answer pairs.  In lines 263-269, we refer to different lines of work that tackle the natural language continuation of hypothetical scenarios. However, such direct evaluation on pre-trained large models creates ambiguities: first, it is unclear how to evaluate the correct response in natural language; second, the uncontrolled sources in pre-training data leads to ambiguities in attributions of success. We thus opt to train neural sequence models from scratch, i.e., we do not rely on off-the-shelf foundation models to perform counterfactual reasoning. Further, we are the first who rigorously attempt to bridge the gap between recent work on counterfactual identifiability [6] and well-studied properties of Transformers.\n\nWith respect to the theoretical analysis of Transformers, the closest to our work is [5]. We adopt a similar setting for our regression setup, but explicitly focus on the noise abduction process. On top of that, we study a cyclic SDE-based setup which, to the best of our knowledge, previous work did not perform. As language is inherently sequential, counterfactual story generation pertains to the completion of a hypothetical scenario after observing the factual narrative. Within one story, consecutive sentences depend on earlier ones. To model this sequential relationship but still maintain the benefits of evaluating on a synthetic setup, we utilize an SDE extension of the popular Lotka-Volterra [7] framework. That way we combine counterfactuals with the previous literature on in-context learning while drawing connections to the work on counterfactual story generation [8]. We are not aware of any other work that tested in-context learning in a counterfactual setup, where it is essential to abduct the noise. In that sense, our work closely aligns with the literature on counterfactual identifiability. To the best of our knowledge, we are the first ones to extend this theoretical finding on invertible, i.e., bijective, functions to Transformers and RNN-type models. In essence, we therefore lay the theoretical foundation for counterfactual reasoning in neural sequence models.\n\nAll that being said, we highly value the feedback of our reviewer. That is, if our response has not addressed your concern, we would be happy if you could provide some references adjacent to our work. As we are not aware of any literature which rigorously tackles counterfactual reasoning in an in-context learning setting, we would be grateful to hear your thoughts.\n\nQ2 *Beyond analyzing attention scores, have you explored other interpretability techniques, such as methods from mechanistic interpretability (e.g., probing the residual stream)?*  \n**We experimented with different techniques (including causal probing) and decided on attention maps due to their interpretable results and prevalence in the literature.**  \nWe thank the reviewer for suggesting probing techniques. When sketching the outline of this paper, we considered different techniques from mechanistic interpretability to include in our work. Apart from attention visualization, we thought about gradient-based explainability methods and probing. For several reasons, we have decided on presenting attention maps as our interpretability method of choice.\nTo start off, the literature on in-context learning is centered around induction heads [1;3;4]. The standard way to analyse which head copies what information in attention-only Transformers is through attention maps. As we understand noise abduction as an implicit mechanism to copy the unobserved noise forward, we connect our findings to that line of work.\nOn top of that, relative to other mechanistic interpretability methods, attention maps are straightforward to grasp. In an otherwise technical paper, we understand attention maps to be a refreshing oasis for readers from different communities.\nMost importantly, causal probing techniques have recently been put under scrutiny. First, several authors argue that seeing probes as evidence for causal properties may be misleading [9;10]. Others underline this idea by providing concrete examples [11]. Second, it is unclear to what degree probes can assist the engineer in explaining the underlying structure. One desideratum of probing is its low complexity [12]. If, for instance, we use nonlinear MLP-based probes, we might be able to predict some latent parameter [13]. Now, the MLP itself is not interpretable in the desired way. In this sense, we are able to “explain” the model by some mechanism. That mechanism itself we do not understand, however. Our initial goal of understanding the model shifts to a mechanism-understanding problem. \nAll that being said, we performed linear probing on our baseline architecture: the 8-layer, 1-head attention-only Transformer. To do so, we fit a linear model to predict $u_y$ from the residual stream, $r$, after every sublayer. For each of the sublayers, we fit a weight matrix and perform all of the following experiments on the residual streams of all sublayers. For clarity, we have dropped indices. Based on the fitted weight matrix $\\hat{W}$, we construct the projection matrix $P$ and predict $r’ = (I - P)r$. This idea of nullifying the linear information carried by the residual stream did not yield relevant results. Similarly, computing the Moore-Penrose pseudoinverse, $W^\\dagger$, and injecting a fresh $u_y^\\text{new}$ into the residual stream through $r’’ = r + W^\\dagger \\left( u_y^\\text{new} - \\hat{u}_y \\right)$ for the linearly fitted noise $\\hat{u}_y$ did not allow us to causally steer the prediction of $y^\\text{CF}$. Taking a step back, this is in line with our intuition as the self-attention sublayer includes nonlinearities through the $\\mathrm{softmax}$ and layer normalization. Therefore, we reaffirm our choice of interpretability method as attention maps are adopted by the community and straightforwardly decodable.\n\n**References**  \n[1] Elhage et al. (2021), A Mathematical Framework for Transformer Circuits  \n[2] von Oswald et al. (2023), Transformers Learn In-Context by Gradient Descent  \n[3] Olsson et al. (2022), In-Context Learning and Induction Heads  \n[4] Akyürek et al. (2024), In-Context Language Learning: Architectures and Algorithms  \n[5] Garg et al. (2022), What Can Transformers Learn In-Context?  \n[6] Nasr-Esfahany et al. (2023), \"On the Identifiability of Counterfactuals from Invertible Generative Models.\"  \n[7] Lotka (1910), Contribution to the Theory of Periodic Reactions  \n[8] Qin et al. (2019), Counterfactual Story Reasoning and Generation  \n[9] Elazar et al. (2020), Amnesic Probing: Behavioral Explanation with Amnesic Counterfactuals  \n[10] Ravichander et al. (2020), Probing the probing paradigm: Does probing accuracy entail task relevance  \n[11] Geiger et al. (2021), Causal Abstractions of Neural Networks  \n[12] Alain et al. (2016), Understanding intermediate layers using linear classifier probes  \n[13] Hewitt et al. (2019), Designing and Interpreting Probes with Control Tasks"}, {"review_id": "UVRZJAxRoe", "reviewer": "Reviewer_oXZV", "summary": "This paper studies in-context counterfactual reasoning using sequential models. It first introduces an in-context learning setup for counterfactual inference, aiming to answer queries of the form $P(y'_{x'} \\mid x, y)$ given a causal graph $X \\rightarrow Y$. The paper shows that this counterfactual query can be reduced to transformations on in-context observations. It then empirically demonstrates that transformers are capable of inferring such counterfactual queries in both linear and exchangeable sequence settings. Extensive empirical studies are also conducted to identify which components of training transformers or sequential models are critical for successful counterfactual inference. Finally, the experiments are extended to nonlinear and more dynamic sequential settings.", "questions": "**Q1.**\n Is there a more closely aligned example that directly illustrates the in-context counterfactual reasoning problem as formalized in the paper? (Please see the discussion in **W1**.)\n\n**Q2.**\n Could the authors briefly connect the identifiability results in invertible structural causal models to Lemma 1? This connection seems highly relevant and, in my opinion, could strengthen the theoretical foundation of the paper.\n (See **W2** for further discussion.)", "limitations": "**L1.** The paper suggests that data diversity in pre-training, self-attention, and model depth are key for the performance of ctf reasoning with experiments. However, it lacks theoretical explanations or intuitive insights to support these observations. Including such a discussion would strengthen the paper and help guide future research.\n\n**L2.** Additionally, as acknowledged by the authors, the experiments are conducted on small-scale synthetic datasets, which may limit the generalizability of the results. There could be a gap between these controlled settings and real-world applications.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "# Strengths\nS1. The paper is well-structured and clearly written. The problem statement and conclusions are presented clearly. \n\nS2. The problem setting is well formalized and original, which may provide valuable insights for future work on in-context learning and counterfactual reasoning. \n\nS3. The main conclusion, that language models can perform in-context counterfactual reasoning, is interesting and well supported by the experiments. This finding could guide future research in both causal inference and language modeling.\n\n\n# Weaknesses\n\n**W1. The example about John in the introduction (lines 30–35) is not well aligned with the actual problem setting of in-context counterfactual reasoning.** In the formal problem setup, the pairs $(x_i, y_i)$ are sampled from a family of SCMs sharing the same causal graph. However, the example, \"John got the flu and is recovering in bed. Had he not attended the get-together, he would now visit his grandmother for cake.\", is just a counterfactual statement and does not reflect the in-context formulation. It would be more helpful if the example directly illustrated the actual problem setting.\n\n**W2. The theoretical insight (Lemma 1) it provides does not seem directly relevant to the core empirical claim of the paper.** The paper states that it \"provides insights into how in-context counterfactual reasoning is equivalent to transformations on in-context observations (Lemma 1),\" but this does not clearly support the conclusion that transformers are capable of performing counterfactual reasoning.\n\nIn particular, Lemma 1 suggests that $y^{\\mathrm{CF}}$ can be expressed as a transformation of $x,y, x^{\\mathrm{CF}}$ (Equation 2). However, this transformation may not be unique. If the true transformation underlying the data does not match the learned model's transformation, then the counterfactual prediction could be entirely incorrect. This corresponds to the identification issue mentioned later in the paper. Similar identifiability issues for invertible SCMs are discussed in prior work [1]. I believe combining such identification results with Lemma 1 would significantly strengthen the theoretical support for the paper’s main claim.\n\n[1] Nasr-Esfahany, Arash, Mohammad Alizadeh, and Devavrat Shah. \"Counterfactual identifiability of bijective causal models.\" International Conference on Machine Learning. PMLR, 2023.\n\nPS: I am happy to raise the score if the concerns are solved but may lower score if the concerns are not solved.", "was_revised": true, "final_justification": "I will maintain my justification as previously communicated to the authors. I find the paper interesting, particularly in its modeling approach. It also validates certain known properties of transformers through experiments on a synthetic dataset. The authors’ rebuttal addressed my questions Q1 and Q2. However, the concerns raised in L_1 and L_2 remain unresolved. The paper neither provides theoretical insights into these properties nor validates them on larger datasets. Therefore, I recommend a borderline accept.", "rebuttal": "We thank the reviewer for their constructive feedback and for appreciating our experimental setup. The concerns are mostly of technical nature, which we address below.  \nQ1 *Is there a more closely aligned example that directly illustrates the in-context counterfactual reasoning problem?*  \n**We provide a more flexible instantiation of our variables that is aligned with the problem setup.**  \nWe agree that the conceptual leap from the example we use in the introduction to our problem setup is large. Consider the following.\nWe define “noise,” $u_y^k$, as the speaker’s personality or mood. The input-output pairs consist of a neutral event description, $x_k$, and the speaker’s subjective framing of that fact, $y_k$. Given a sequence of $n$ in-context examples sharing a common theme, $\\theta$, we query the model to extract the noise from one of the examples, $(x_j, y_j)$ for $z = j$. Next, we supply a counterfactual query, $x^\\text{CF}$ that is not based on the common theme, $\\theta$. The model is then to copy the sentiment of the input-output pair at position $z$, and state the counterfactual continuation, $y^\\text{CF}$. We give an instantiation below.\nSuppose the shared theme $\\theta$ to be disruption events. Three input-output examples of neutral description and subjective continuation read:  \n$x_1$ = [The train was late.], $y_1$ = [Well, I guess I’m walking again.]  \n$x_2$ = [It started raining.], $y_2$ = [Yay, finally I have the chance to try out my new coat.]  \n$x_3$ = [The meeting was canceled.], $y_3$ = [What a relief! I can prepare dinner early today.]  \nNote the shared theme of disruption events. In addition, we observe three different moods ranging from (1) annoyed over (2) cheerful to (3) relieved. Consider now an intervention, $x^\\text{CF}$, unrelated to $\\theta$.  \n$x^\\text{CF}$ = [I found a slice of cake in the fridge.]  \nClearly, this is not a disruption event. Yet, depending on the index token $z \\in \\{1, 2, 3\\}$ and the associated personality, we expect different continuations,  \n$y^\\text{CF} | z = 1$ = [Yet another person who didn’t finish their food. It’s still a shared space here.]  \n$y^\\text{CF} | z = 2$ = [Sweet! Today just keeps getting better.]  \n$y^\\text{CF} | z = 3$ = [Thank god, I don’t have to go buy dessert now.]  \nImmediately, we see how the three personalities from above relate to the different continuations. In this bivariate setting, the “what-if” question is implicit. Instead of explicitly intervening on the examples in question format, we provide a new out-of-distribution example and the index the model should base its personality on. Whilst this may represent a carefully concocted illustration, we view in-context counterfactual reasoning as a general framework for latent attribute abduction and application. Beyond mood, this includes stylistic continuations across dialects, registers, or musical genres.  \nQ2 *Could the authors briefly connect the identifiability results in invertible structural causal models to Lemma 1?*  \n**We introduce a new theorem that links the identifiability results under Markov to our exchangeable setup.**  \nWe understand the reviewer’s point that the connection between the counterfactual prediction involving a transformation on the observed values $x, y, x^\\text{CF}$ and the Transformer effectively learning the causal structure is not immediate. First, we point out that the class of functions we consider can be subsumed under the class of bijective generation mechanisms [1]. Second, each function we test our hypothesis on (App. E.2) is monotonic and continuous on $\\mathbb{R}$. Hence, by the inverse function theorem, they are invertible ($\\Longleftrightarrow$ bijective). Last, invoking the exchangeable setup we use throughout, we have conditional independence of the form $X \\perp U | \\theta$. Our mechanism $T(f(x),\\cdot)$ invertible, continuous (for $f(x)$ invertible, continuous) is then monotonic for all $f(x)$. Given the law $\\mathbb{Q}{X,Y} := \\mathbb{P}{X,Y,\\theta}$, we invoke Lemma B.2 and Theorem 5.1 from [1] and claim:\n\nLet $X, Y, U, \\theta$ scalar random variables with $X \\perp U | \\theta$. Take $T: \\mathcal{Y} \\times \\mathcal{N} \\longrightarrow \\mathcal{Y}$ with $y = T(f(x),u)$. Assume $\\forall f(x) \\in \\mathcal{Y}$, the inverse $T^{-1}(f(x),\\cdot)$ exists for all $y$, i.e., $u = T^{-1}(f(x),y)$, and $\\forall f(x) \\in \\mathcal{Y}$, $T(f(x), \\cdot)$ is continuous. Suppose further that $f(x)$ continuous $\\forall x$, and strictly monotonic in $x$. Then, set $Y = T(f(X),U)$. Given the joint law $\\mathbb{P}_{X,Y|\\theta}$, $T$ is counterfactually identifiable.\n\nThe proof to this statement is analogous to [1]. Regarding efficient learning of the underlying causal structure, we invoke previous literature on the Transformer’s ability to learn various function classes in-context [2]. We will include this connection in the revised version.  \nL1 *The paper suggests that data diversity in pre-training, self-attention, and model depth are key for the performance of ctf reasoning with experiments. [...] Including such a discussion would strengthen the paper and help guide future research.*  \n**We provide prior literature to underscore our Transformer-related findings.**  \nBelow, we establish an intuitive connection. In the next version, we are happy to include a more detailed discussion.\nRecent theoretical research on data diversity aligns closely with our observations. [3] develop a formal framework showing that diversity in example selection improves in‑context inference. With respect to model depth, [4] show that at least two or three attention layers are required to support compositional generalization beyond memorization. This is in line with our findings. Figure 3(b) indicates that the additional benefit of deeper models decreases after having trained models of depth 2. Improvements are especially high for 2 layers relative to 1. Intuitively, we relate this to the multi-step process of performing counterfactual prediction. The model retrieves the relevant pair, $(x_j,y_j)$ for $z = j$, infers noise, and predicts the counterfactual $y^\\text{CF}$ given $x^\\text{CF}$ (Figure 1). Adopting the residual stream view introduced by [5], we argue that we require multiple layers to copy the information forward. We observe distinct attention patterns [6;7]. While a typical head in the first layer attends to the $z$-token, we also find heads copying the information of the counterfactual intervention forward. As noise abduction requires two steps - copying $x_j$ and $y_j$ can happen in parallel in different heads and transforming the latter to obtain $u_j^y$ occurs at the next layer - we intuitively relate our results to established findings on self-attention.  \nL2 *There could be a gap between these controlled settings and real-world applications.*  \n**We follow the literature in designing controlled experiments to confirm our theoretical hypothesis on a small scale.**  \nWe agree with the reviewer that real-world applications like natural language are likely to be more complex than the synthetic task we study. The main purpose of our paper is to connect the theoretical result of performing counterfactual prediction to known properties of Transformers.  \n1. We are interested in confirming that Transformers are capable of performing noise abduction while learning a certain functional form in-context. Showing this theoretically through Lemma 1 in combination with the above statement, we check whether this is in principle possible on a small scale.  \n2. We envision that the mechanism behind in-context counterfactual reasoning remains identical for larger models. Extrapolating results from smaller models is in line with existing literature on the theoretical understanding of large models. For instance, [9] implement a two-layer attention-only Transformer with neither biases nor layer normalization. Similarly, [8] extrapolate their findings from an architecture with linear self-attention. Even more so, our work is grounded in earlier findings on induction heads [6;7] and the in-context learning of functions $f$ [2], which, too, are derived from small models.  \n3. Synthetic tasks play a central role in causal inference. Controlled setups often form the backbone of key theoretical insights. [9] benchmark causal discovery methods on synthetic data before moving to real-world examples like the canonical temperature-altitude case. Similarly, [10] evaluate their Do Finetti algorithm exclusively on synthetic data. In our work, we aim to demonstrate that a causal mechanism can emerge in-context, without constructing a large, standalone model ourselves. Instead, we hope this work lays the groundwork for in-context counterfactual reasoning by integrating causality with in-context learning.\n\n**Conclusion** We thank the reviewer again for their thoughtful feedback and their refreshing input on counterfactual identifiability. We have added a more aligned example and formally connected the cited paper. If questions remain, feel free to reach out. Otherwise, we appreciate your consideration in raising the score as suggested.  \n**References**  \n[1] Nasr-Esfahany et al. (2023), On the Identifiability of Counterfactuals from Invertible Generative Models  \n[2] Garg et al. (2022), What Can Transformers Learn In-Context?  \n[3] Xiao et al. (2025), The Role of Diversity in In-Context Learning for Large Language Models  \n[4] Petty et al. (2023), The Impact of Depth on Compositional Generalization in Transformer Language Models  \n[5] Elhage et al. (2021), A Mathematical Framework for Transformer Circuits  \n[6] Olsson et al. (2022), In-Context Learning and Induction Heads  \n[7] Akyürek et al. (2024), In-Context Language Learning: Architectures and Algorithms  \n[8] von Oswald et al. (2023), Transformers Learn In-Context by Gradient Descent  \n[9] Peters et al. (2014), Causal Discovery with Continuous Additive Noise Models  \n[10] Guo et al. (2024), Do Finetti: on Causal Effects for Exchangeable Data"}, {"review_id": "DzdNgrvDem", "reviewer": "Reviewer_CTcD", "summary": "The paper intends to answer the question whether LLMs can take a sequence of factual examples, imagine a “what-if” change, and generate the output without gradient updates. It examines the problem with a synthetic linear regression task requiring noise abduction.  A large transformer-based model is trained to fit the data from a mixture of sequences with different distributions but the same causal structure. Other models with different architectures, such as GRU and LSTM are trained on the same data for comparison. The paper concludes that large transformer models have the capability of in-context counterfactual reasoning.", "questions": "1. The paper assumes that a sequence of random variables is exchangeable. However, the order of words in an utterance matters in language model training. What is the implication of this assumption on LLM-generated counterfactuals?\n\n2. Line 109-110 are not clear to me. How \\theta is defined and how does it parameterise $U_X$ and $U_Y$?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "S1: Examining the in-context counterfactual output generation capability of large transformers is an novel idea. This is an interesting direction considering the importance of counterfactual reasoning in decision making. \n\nS2: A controlled study with linear regression is helpful to understand the transformer behaviour on counterfactual output generation. \n\nW1: The training objective at line 116 minimises the counterfactual loss without considering the factual loss. Does the training have negative impact on factual output prediction? With such a single objective optimisation, the formulation is not much different to approaches studying emergence phenomenon of in-context learning itself, such as the following:\n\nGarg, S., Tsipras, D., Liang, P.S. and Valiant, G., 2022. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35, pp.30583-30598.\n\nIn a sense, the paper reaches a similar conclusion as the above paper due to that the training overfits counterfactual examples. I might have missed something here.\n\nW2: In addition to transformers, the evaluation also shows results from other neural network architecture, e.g., GRU and LSTM also show counterfactual generation capability. I am wondering how the claim that counterfactual reasoning emerges in self-attention is justified. \n\nW3: Line 41. it is better to describe \\mu_y clearly here. In line 48, \\mu_y is referred as context, but in Fig.1, it is referred as noise. Can you please also explain why the \\mu_y is indexed by y not x?", "was_revised": true, "final_justification": "The idea proposed in this paper is novel. The rebuttal addresses my concern on its technical treatment.", "rebuttal": "To begin with, we thank the reviewer for their constructive feedback and their positive opinion on our problem setup. Most of the concerns are related to notation and imprecise wording. We address the questions below.\n\n[W1] **We run new experiments and argue that the two setups are complementary, not identical.**\n\nWe thank the reviewer for raising the concern on methodological innovation. We agree that the setup we study is comparable to that of [1]. This we openly acknowledge in our introduction (lines 40-41). We point out how our conclusions differ from that framework. We focus on noise abduction through transformations (Lemma 1). Linear regression, as studied by [1], does not require the ability to perform noise abduction. Thus, we design a problem setup that requires noise abduction for accurate prediction. In doing so, we build upon recent development in counterfactual identifiability [2]. Thus, we take one step further than [1]. Not only does our setup require the model to learn the functional form $f$ in-context, but also to abduct the unobserved noise. To the best of our knowledge, this is the first work that explicitly studies in-context noise abduction in neural sequence models.\nTo underscore our point, we also trained an attention-only (AO) 8-layer, 1-head Transformer on the setup from [1] and found that its in-context MSE is 38% above the loss recorded for counterfactual prediction on the same model. We observe similar trends for Transformers of different configurations as well as RNN-type models. This is in line with our hypothesis that in-context counterfactual learning, in principle, leads to a lower prediction error than the standard regression setup [1] as the noise is fixed and can be inferred directly.\n\nIn addition, we discuss potential negative impacts on factual output prediction. To do so, we perform a copying task. We take the 8-layer, 1-head attention-only Transformer and ask it to predict the factual completion $y_1$ of the input-output pair $(x_1, y_1)$ after seeing only that in-context example. As can be seen in Figure 3(a), the 8-layer AO architecture performs poorly on predicting counterfactuals  from one example only. Intuitively, this is in accordance with [1] as the functional form $f$ is to be learned in-context. The experiment we conduct here, however, is function-agnostic. We simply test whether the noise abduction as indexed by $z$ works as planned. We find that the in-context prediction of $y_1$ is 71% below the counterfactual prediction $y^\\text{CF}$ with independently sampled $x^\\text{CF}$. Taken together, we understand the two above results (on the setup from [1] relative to our design and on directly predicting $y_1$) as evidence that the model does not simply overfit on counterfactual examples. Instead, it appears to maintain an internal understanding of the three-step process: noise abduction through index token $z$, intervention based on $x^\\text{CF}$, and prediction of $y^\\text{CF}$.  \nTo summarize, our work is building on top of [1]. As it has been shown that prediction in a linear regression setting is possible [1], one can build a machine that combines the two paradigms. Our framework is flexible to include an index token $z = 0$ encoding linear regression. In that sense, we can understand $z$ as a switch for counterfactual prediction.\n\n[W2] **Both classes can perform counterfactual reasoning, self-attention is the mechanism why Transformers can.** \n\nFirst, we study whether neural sequence models are able to perform counterfactual prediction. We observe that Transformers do so, but also RNN-type sequence models achieve low error (Figure 2). We stress that our RNN architectures have been tested on the (selective) copying task [3;4;5;6]. Given the literature, we confine ourselves to the analysis of autoregressive Transformers, as those represent the state-of-the-art architecture. Hence, all future statements only refer to the Transformer. Within the group of trained Transformers, we find that self-attention is relevant to perform counterfactual prediction. Given previous work on induction heads [7;8], we argue that this ability to abduct noise is attributed to counterfactual prediction [2] essentially involving the inference of unobserved noise from the observational pair $(x, y)$. In that sense, we deem self-attention relevant to counterfactual reasoning.\n\n[W3] **We will unify the vocabulary in the next version.**\n\nWe agree that the usage of the word “context” may be ambiguous here. We point the reviewer’s attention to the introduction. We understand counterfactual reasoning as the ability to predict the consequence of a hypothetical scenario, while keeping everything else constant. Put differently, we analyze the counterfactual output while fixing the context. In Pearl’s [9] causality framework, this is made explicit by defining a counterfactual statement as the prediction given fixed noise. Hence, by context, we refer to the noise in the former setting of counterfactual reasoning in natural language. It is correct that this may be misleading and we acknowledge that this subtlety is not made explicit in our work. Taking the reviewer’s advice seriously, we will unify the usage of “noise” and “context” throughout in the next version of the paper.\n\nRegarding your second point, we follow the standard notation for structural causal models [11]. Note that in the bivariate case, we define $X = U_X, Y = \\beta X + U_Y$. Hence, we implicitly refer to the variable $X$ itself as the noise associated with variable $X$. This becomes clear if we consider more than two variables. Suppose we have $(X, Y, Z)$ with mutually independent noise $U = (U_X, U_Y, U_Z)$. Consider the case where both Y and Z have one causal parent, $X$. Here, $Y$ and $Z$ are independent when conditioning on $X$. Writing $Y = \\beta X + U_X$ and $Z = \\gamma X + U_X$, it is unclear how $X$ should be defined. More strikingly, $Y$ and $Z$ are not independent given $X$. We arrive at a contradiction. We hope that this example explains our choice of notation.\n\n[Q1] **We introduce the cyclic setup to capture sequentiality.**\n\nWe apologise for the confusion and clarify. As laid out in lines 55-57, we assume exchangeability over input sequences, not within sequences. This is in line with current research [10] on relating in-context learning to De Finetti’s theorem. Within one input sequence, we assume that in-context examples are conditionally independent given the latent parameter, $\\theta$. Having hopefully clarified this, we completely agree with the reviewer that “the order of words in an utterance matters.” This marks the very reason for our extensive analysis in section 5. Exactly because natural language is sequential and deeply interdependent (see lines 214-216), we introduce cyclic causality to our setup. Counterfactual story generation pertains to the completion of a hypothetical scenario after observing the factual narrative. Within one story, consecutive sentences depend on earlier ones. To model this sequential relationship but still maintain the benefits of evaluating on a synthetic setup, we utilize an SDE extension of the popular Lotka-Volterra [11] framework. That way we unite counterfactuals with the previous literature on in-context learning while drawing connections to the work on counterfactual story generation [12]. In doing so, we are able to circumvent the strong assumption of conditional independence across in-context examples.\n\n[Q2] **All noise is parameterised given the latent $\\theta$.**\n\nThank you for raising this point, we refer the reviewer to the more detailed setup in Appendix C (see line 143). In particular, we hope that equation (9) in the appendix.pdf file of the supplementary material can resolve the unclarity. To be explicit, we sample $\\theta$ from a uniform distribution over the interval $[-6,6]$, and draw conditionally independent noise variables $U_X$ and $U_Y$ from a unit-variance Gaussian parameterised by mean $\\theta$, i.e., $\\mathcal{N}(\\theta, 1)$. This way, we enforce conditional independence given the latent parameter $\\theta$ both between $U_X$ and $U_Y$ and across in-context examples. Related to the reviewer’s previous question, we hope that this also clarifies the distinction between exchangeable input sequences and conditionally independent in-context examples.\n\n**Conclusion** To conclude, we thank the reviewer for their thoughtful feedback. We discuss the novelty of moving beyond i.i.d. data and our focus on noise abduction in neural sequence models. Our experimental comparisons further support that the model performs counterfactual reasoning rather than simple overfitting. We also spell out the reason for analyzing Transformers and explain the role of self-attention in abducting noise. Ambiguities in terminology (i.e., \"context\" vs. \"noise\") have been acknowledged, and we will revise the manuscript accordingly. Lastly, we have restated our assumptions about exchangeability and our data setup in Appendix C.\n\n**References**  \n[1] Garg et al. (2022), What Can Transformers Learn In-Context?  \n[2] Nasr-Esfahany et al. (2023), On the Identifiability of Counterfactuals from Invertible Generative Models  \n[3] Hochreiter et al. (1997), Long short-term memory  \n[4] Arjovsky et al. (2015), Unitary Evolution Recurrent Neural Networks  \n[5] Henaff et al. (2016), Recurrent Orthogonal Networks and Long-Memory Tasks  \n[6] Feng et al. (2024), Were RNNs All We Needed?  \n[7] Olsson et al. (2022), In-Context Learning and Induction Heads  \n[8] Akyürek et al. (2024), In-Context Language Learning: Architectures and Algorithms  \n[9] Pearl (2009), Causality  \n[10] Ye et al. (2024), Exchangeable Sequence Models Quantify Uncertainty Over Latent Concepts  \n[11] Lotka (1910), Contribution to the Theory of Periodic Reactions  \n[12] Qin et al. (2019), Counterfactual Story Reasoning and Generation"}, {"review_id": "fg4I2LSJbd", "reviewer": "Reviewer_nBNw", "summary": "The authors present a study of counterfactual reasoning by in-context learning of large language models. They provide synthetic experiments that show, based on MSE performance, that language models are capable of performing the task.", "questions": "1. How does the example given in appendix G represent \"what if\" questions or explain the task of noise abduction in natural language? Please provide an instantiation of the variables used in the synthetic setup on a concrete example.\n2. What is Figure 2 supposed to show: 1) Attention is required for optimal performance (GPT-2)? 2) Attention is not required since GRUs are a close second? \n3. How would you evaluate the  \"reasoning\" capabilities of the model itself, not just the results?", "limitations": "yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 1, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- Rigorous theoretical underpinning of counterfactual reasoning\n- Controllable study on synthetic data\n\nWeaknesses: \n- Due to the synthetic setup, it is very hard to understand what the experimental hypothesis is, and how the described experiment related to natural language generation. The example given in appendix G does not help since it does not represent \"what if\" questions nor does it explain the task of noise abduction in natural language.\n- What is Figure 2 supposed to show: 1) Attention is required for optimal performance (GPT-2)? 2) Attention is not required since GRUs are a close second? \n- The evaluation of the \"reasoning\" capabilities is based on the MSE results, not on an investigation of the reasoning process itself. It is well-known that the correctness of  the final conclusion does not necessarily reflect the validity of a model's reasoning process (see https://arxiv.org/abs/2404.01869, https://arxiv.org/abs/2305.18654)", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive feedback and are happy that the theoretical rigour of our work resonates with them. The feedback primarily centers around the presentation of results and definitions. We address the concerns below.\n\nW1 *[...] it is very hard to understand what the experimental hypothesis is [...]*\n\n**Our testable hypothesis is that neural sequence models are capable of performing counterfactual reasoning.**\n\nWe agree that the conceptual leap from the synthetic setup to natural language is nontrivial. To restate our research question, we ask whether language models can perform in-context counterfactual reasoning (line 52). Today, most research on counterfactual story generation in large language models is based on prompting off-the-shelf language models [1]. To the best of our knowledge, there exists no prior work on rigorously disentangling the mechanism underlying counterfactual reasoning. An adequate toolkit to do so is causality [2]. Thus, we train a suite of neural sequence models on noise abduction. We supply a sequence of $n$ conditionally independent input-output pairs, $(x_k, y_k)$ for $k \\in \\{1,...,n\\}$, where $y_k$ depends on $x_k$ and some noise $u_y^k$. Then, we ask the model to perform counterfactual prediction based on the pair $(x_j, y_j)$ for $j \\in \\{1,...,n\\}$ by providing an index token $z = j$. The model is to infer the unobserved noise $u_y^j$ and to predict given the new intervention $x^\\text{CF}$. Thus, the optimal prediction reads $y^\\text{CF} = x^\\text{CF} + u_y^j$.\nHere, we see the benefits of causality. The desired output $y^\\text{CF}$ is unique and deterministic. In contrast to natural language, we can straightforwardly assess whether the model output coincides with the ground truth value. Below, we will argue how this also assists us in speaking about “reasoning.”\n\nQ1 *How does the example given in appendix G represent \"what if\" questions or explain the task of noise abduction in natural language?*\n\n**We provide a more flexible instantiation of our variables that is aligned with the problem setup.**\n\nWe stress that the “what if” question is implicit in the model setup. By prompting the model with the relevant index token $z$ and the intervention $x^\\text{CF}$, noise abduction occurs naturally. In the end, we will show how this idea opens up the opportunity to study noise abduction beyond natural language.\nWe define “noise,” $u_y^k$, as the speaker’s personality or mood. The input-output pairs consist of a neutral event description, $x_k$, and the speaker’s subjective framing of that fact, $y_k$. Given a sequence of $n$ in-context examples sharing a common theme, $\\theta$, we query the model to extract the noise from one of the examples, $(x_j, y_j)$ for $z = j$. Next, we supply a counterfactual query, $x^\\text{CF}$ that is not based on the common theme, $\\theta$. The model is then to copy the sentiment of the input-output pair at position $z$, and state the counterfactual continuation, $y^\\text{CF}$. We give an instantiation below.\nSuppose the shared theme $\\theta$ to be disruption events. Three input-output examples of neutral description and subjective continuation read:\n\n$x_1$ = [The train was late.], $y_1$ = [Well, I guess I’m walking again.]\n\n$x_2$ = [It started raining.], $y_2$ = [Yay, finally I have the chance to try out my new coat.]\n\n$x_3$ = [The meeting was canceled.], $y_3$ = [What a relief! I can prepare dinner early today.]\n\nNote the shared theme of disruption events. In addition, we observe three different moods ranging from (1) annoyed over (2) cheerful to (3) relieved. Consider now an intervention, $x^\\text{CF}$, unrelated to the common $\\theta$.\n\n$x^\\text{CF}$ = [I found a slice of cake in the fridge.]\n\nClearly, this is not a disruption event. Yet, depending on the index token $z \\in \\{1, 2, 3\\}$ and the associated personality, we expect different continuations, \n\n$y^\\text{CF} | z = 1$ = [Yet another person who didn’t finish their food. It’s still a shared space here.]\n\n$y^\\text{CF} | z = 2$ = [Sweet! Today just keeps getting better.]\n\n$y^\\text{CF} | z = 3$ = [Thank god, I don’t have to go buy dessert now.]\n\nImmediately, we see how the three personalities from above relate to the different continuations. In this bivariate setting, the “what-if” question is implicit. Instead of explicitly intervening on the examples in question format, we provide a new out-of-distribution example and the index the model should base its personality on. \nWhilst this may represent a carefully concocted illustration, we view in-context counterfactual reasoning as a general framework for latent attribute abduction and application. Beyond mood, this includes stylistic continuations across dialects, registers, or musical genres.\n\nQ2 *What is Figure 2 supposed to show: 1) Attention is required for optimal performance (GPT-2)? 2) Attention is not required since GRUs are a close second?* \n\n**Both classes can perform counterfactual reasoning, self-attention is the mechanism why Transformers can.**\n1. We study whether neural sequence models are able to perform counterfactual prediction. We observe that Transformers and RNN-type sequence models do so (Figure 2). We stress that our RNN architectures have been tested on the (selective) copying task [3-6]. \n\n2. Given this, we confine ourselves to the analysis of autoregressive Transformers, as those represent the state-of-the-art architecture. Within the group of trained Transformers, we find that self-attention is relevant to perform counterfactual prediction. Given previous work on induction heads [7;8], we argue that this ability to abduct noise is attributed to counterfactual prediction [9] essentially involving the inference of unobserved noise from the observational pair $(x,y)$. In that sense, we deem self-attention relevant to counterfactual reasoning. Therefore, we argue that Transformers perform noise abduction in-context due to self-attention while RNN-type sequence models implement other mechanisms to copy information forward.\n\nQ3 *How would you evaluate the \"reasoning\" capabilities of the model itself, not just the results?*\n\n**We adopt the definition from causal literature and design our setup to interpret the MSE.**\n\nWe share the reviewer’s concern that correct outputs do not guarantee valid reasoning processes. We list three arguments for choosing MSE. \n1. We clarify our working definition of causal reasoning. Similar to [10], we define reasoning as “algebraically manipulating previously acquired knowledge to answer a new question.” In the context of causality, this reduces to “the process of drawing conclusions from a causal model” [11]. The latter statement also aligns with the work mentioned by the reviewer [12].\n\n2. The reviewer states that evaluating the model on the final loss is insufficient [13]. This is correct if we consider complex Transformers which are trained on natural language. Especially when evaluating on integers, it is often not immediate whether the correct prediction can be attributed to legitimate reasoning. We do not disagree with that comment. As acknowledged in lines 50-53, this represents the very reason why we train a small Transformer under this controlled setup. We are able to invoke known results on the attention mechanism [7;8] while extending it to counterfactual prediction. By turning off the MLP sublayer for the most part, we exclude another source of complexity. With the 8-layer, 1-head attention-only Transformer we can track the steps (noise abduction, intervention, prediction) in the model. Put differently, we can understand the algebraic manipulation of the factual input $(x,y)$ in order to answer the new query $x^\\text{CF}$. In accordance with the working definitions above, we deem the correct final prediction sufficient.\n\n3. Recent research [9] has found that it is possible to perform counterfactual prediction. The setup we choose agrees with this modeling paradigm. To be explicit, we will append to our paper a corresponding theorem which relates that finding to our work. This implies that it is in principle possible to perform counterfactual reasoning simply from observing factual data $(x,y)$. Turning the question on its head, therefore, our analysis starts at the fact that counterfactual reasoning is theoretically possible. We then design a simple experiment which relates that finding to modern deep learning architectures. We turn off those sublayers which are not mechanistically interpretable. Finally, we confirm the hypothesis by invoking known results on how information is copied forward in self-attention.\n\n**Conclusion** We provide an instantiation of the variables used on a concrete example. Further, we clarify how our results on Transformers relate to other neural sequence models. Finally, we state our reasoning definition and connect it to existing work. We hope this addresses the reviewer's concerns. If our feedback resonates with the reviewer's perspective, we would like you to consider raising the score.\n\n**References**\n\n[1] Qin et al. (2019), Counterfactual Story Reasoning and Generation  \n[2] Pearl (2009), Causality  \n[3] Hochreiter et al. (1997), Long short-term memory  \n[4] Arjovsky et al. (2015), Unitary Evolution Recurrent Neural Networks  \n[5] Henaff et al. (2016), Recurrent Orthogonal Networks and Long-Memory Tasks  \n[6] Feng et al. (2024), Were RNNs All We Needed?  \n[7] Olsson et al. (2022), In-Context Learning and Induction Heads  \n[8] Akyürek et al. (2024), In-Context Language Learning: Architectures and Algorithms  \n[9] Nasr-Esfahany et al. (2023), On the Identifiability of Counterfactuals from Invertible Generative Models  \n[10] Bottou (2011), From Machine Learning to Machine Reasoning  \n[11] Peters et al. (2017), Elements of Causal Inference  \n[12] Mondorf et al. (2024), Beyond Accuracy: Evaluating the Reasoning Behavior of Large Language Models  \n[13] Dziri et al. (2023), Faith and Fate: Limits of Transformers on Compositionality"}]

Author rebuttal:

Meta-review: This paper evaluates the ability of LLMs to perform in-context counterfactual reasoning. The experiments explore the internal mechanisms supporting counterfactual reasoning in transformers (and find the architectural factors that support / hinder in-context counterfactual reasoning). All experiments are performed on a synthetic dataset. The experiments seem sound and comprehensive, but as mentioned by several reviewers, the framing (of, e.g., natural language counterfactual reasoning; the examples given in the rebuttal also don't seem to exhibit counterfactual reasoning, which is more like what-if-instead (i.e., conditioned on alternatives), but simply what-if reasoning) is very far from the actual problem being studied, and the discussion of "emergence" of reasoning is not implied by their ability to perform such reasoning when being trained on a task that implicitly requires it. I would suggest the authors significantly revise the framing and clarity in any final version of the paper.
Decision: Accept (poster)