CCL: Causal-aware In-context Learning for Out-of-Distribution Generalization
Authors: Hoyoon Byun
Gyeongdeok Seo
Joonseong Kang
Taero Kim
Jihee Kim
Kyungwoo Song
Venue: NeurIPS 2025 poster ()
Keywords: In-context learning
Causal representation
Out-of-Distribution Generalization
Primary area: deep_learning

Abstract: In-context learning (ICL), a nonparametric learning method based on the knowledge of demonstration sets, has become a de facto standard for large language models (LLMs). The primary goal of ICL is to select valuable demonstration sets to enhance the performance of LLMs. Traditional ICL methods choose demonstration sets that share similar features with a given query. However, we have found that the performance of these traditional ICL approaches is limited on out-of-distribution (OOD) datasets, where the demonstration set and the query originate from different distributions. To ensure robust performance in OOD datasets, it is essential to learn causal representations that remain invariant between the source and target datasets. Inspired by causal representation learning, we propose causal-aware in-context learning (CCL). CCL captures the causal representations of a given dataset and selects demonstration sets that share similar causal features with the query. To achieve this, CCL employs a novel VAE-based causal representation learning technique. We demonstrate that CCL improves the OOD generalization performance of LLMs both theoretically and empirically. Code is available at: \url{https://github.com/MLAI-Yonsei/causal-context-learning}

CCL: Causal-aware In-context Learning for Out-of-Distribution Generalization
Hoyoon Byun, Gyeongdeok Seo, Joonseong Kang, Taero Kim, Jihee Kim, Kyungwoo Song" Department of Statistics and Data Science, Yonsei University (hoyun. byun, seo, doongsae, taero kim, jihee_stee kyungwoo song) Syonsei. kc.kr
Abstract
In-context learning (ICL), a nonparametric learning method based on the knowledge of demonstration sets, has become de facto standard for large language models (LLMs). The primary goal of ICL. is to select valuable demonstration sets to enhance the performance of L.LMs. Traditional ICL methods choose demonstration sets that share similar features with given query. However, our experiments reveal that these traditional ICL approaches perform poorly on out-of-distribtiion (OOD) datasets, where the demonstration set and the query originate from different distributions. To ensure robust performance in OOD datasets, is essential learn causal representations that remain invariant between the source and target datasets. Inspired by causal representation learning. we propose causal-awan in-context learning (CCL). CCL captures the causal representations of given dataset and selects demonstration sets that share similar causal features with the query. To achieve this, CCL employs novel VAF-based causal representation learning technique. We demonstrate that CCL improves the OOI) generalization performance of LLMs both theoretically and empirically. Code is available at: hattps://githin cOmMAAII Yonsei context waa---ootttet-eeeninne
1 Introduction
While large language models (LLMs) excel general- purpose pre-trained models, in-context learning (ICL) has become a key approach for aligning them to target tasks. ICL [1] enables L.MMs to adapt to new tasks with few demonstrations and without parameter updates, making it applicable in various fields. While ICL has shown significant promise, it still faces difficulties in achieving robust generalization ( A primary challenge is that LLMs rely superficial patterns in demonstration sets, which restrict their capability in unseen environments 3 Recent studies indicate that distribution shifts between demonstration
and target queries in out-of-distribution (OOD) scenarios impede the ability of LLMs to generalize effectively 4 5
6.. To fully unlock the potential of LLMs and enable reliable deployment in real-world applications, ensuring robustness in OOD scenarios plays pivotal role.
The pursuit of ensuring generalization beyond observed data naturally leads to the question of how the data was generated. Drawing on insights from causality (7), the structural knowledge of data is expressed using causal language. In causal representation learning (CRL) 88, observed data reflect underlying latent causal variables that drive the data- santa-geneatiing process (DGP). CRL aims to model the causal mechanisms among these variables [9] For example, if causal variables independent, one remains invariant even when the other, acting as an environmental factor, changes 10 [1]. The assumption about causal mechanisms suggests that learning invariant causal variables is an effective approach for models robust to distribution shifts. Consequently, CRL lays the groundwork
"Corresonniig author
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Demonstration Candidates
Selection
Inference
X1: Tom, who loves fruit, ate of the 10 bananas, How do you think he feels?
Demonstration Set
t
C1: Sentiment
y1: Positive
c Llama
X2: this game, 6 survivors from the group of 17 will be eliminated. How many will remain standing?
S
GPT
Space of
LLM
Cz:
Yz: 11
Target Query
x':
Tom ate 5 out of 12 bananas. How many bananas are left?
c²: M
y":
Space
ICL
CCL
Figure 1: To enhance the OOD performance of L.M,, causally-related demonstration set is important. Current ICL methods compare the non-causal representation I; and x², and they might choose a worthless demonstration set (Candidate 1). However, method, CCL, compares the causal representation G and c* to construct a demonstration. Because CCL leverages the causal-related demonstration (Candidate 2), CCL shows superior performance the OOD dataset.
for research on OOD generalization by learning invariant representations from training sets collected across multiple environments. 12 14).
Considering causal mechanisms allows for constructing more suitable demonstration from demonstration candidates when their environments differ from that the target query. In Figure [ the target contextual problem x= is highly similar to æ1 (highlighted in green), making æ1 strong But what exactly is the problem embedded in the target query? For LLMS successfully generalize the target query, the demonstration should be constructed to reflect the fundamental context rather than relying on superficial patterns, such as frequently occurring words or characteristics of the data collection environment [17]
candidate in ICL. 15 16
Therefore, even the superficial context differs, a candidate Z2 that addresses the same problem (N- M=K) should be included in the demonstration set (highlighted in red). It ensures the demonstration set captures problem-leve invariance even when generalizing to OOD targets from given candidates. Since causal variables which generate the contextual problem, not observable objects, it is necessary to model e under the assumption of causal mechanisms that remain invariant across environments.
In this study, we focus on constructing robust demonstration set to enhance the generalization of LLMs in OOD scenarios. Inspired by CRL, propose novel demonstration selection method, causal-aware in-context learning (CCL), which learns causal representations that remain invariant across environments and prioritizes candidates by assigning higher ranks to those with causal representations similar to the target query. Under the causal mechanism, we theoretically demonstrate that the demonstration set selected by CCL comprises candidates that the underlying problem addressed by the target query, rather than merely matching its context. The problem-level invariance of CCL ensures generalization performance for the target query even in unseen environments. We empirically validate that CCL operates robustly in OOD scenarios and demonstrates superior generalization performance both synthetic and real datasets.
closely related to
2 Related Works
2.1 In-context learning
ICL is a method where LLMs perform tasks by leveraging examples from the input context without updating model parameters (1. This approach enhances computational efficiency and achieves competitive performance in various natural language tasks without the need for model fine-ture fine-tuning

---

However, the performance of ICL is sensitive to demonstration organization, including demonstration selection [99 20]. Various approaches aim to optimize demonstration selection in ICL,
2 18.
including unsupervised methods that use similarity metrics like k-nearest neighbors [15] as well as
supervised techniques that leverage task-specific retrievers 21 and reinforcement learning [22
Despite these advancements, LLMs depend on surface-level patterns in the demonstration set, leading to primary challenge with out-of-distribution (OOD) examples [3]. While larger models tend to reduce the performance gap between in-distribution (ID) and OOD scenarios, even transformers, which handle minor distribution shifts, face significant challenges when encountering major shifts (6.4) 4). The BOSS benchmark evaluates OOD robustness in ICL. highlighting the importance of addressing OOD generalization [5]. An approach designed to improve OOD performance involves inferring latent variables from the context using the transformer architecture. However, this method struggles to apply those variables effectively in prediction, limiting OOD) generalization [23]. We propose CCL,, drawing causal representation learning, to improve OOD performance in ICI. by focusing on task-relevant causal features and enhancing robustness distribution shifts.
2.2 Causal representation learning
Unlike statistical approaches, which describe the distributional characteristics of data, causality ? focuses the structural relationships between variables. The DGP is determined by the underlying causal relationships among variables, and a structural causal model (SCM) is a generative model
that describes the DGP [T0 24. The SCM expresses the uncertainty of exogenous factors in probabilistic manner and defines functional relationships for the variables of interest (endogenous variables), thus structurally describing the causal mechanisms of the DGP. Observed data represent one of the realizations of these causal mechanisms. A causal graph visually represents the structural relationships between the variables, induced by the SCM 77.
Recently, research in machine learning has increasingly focused on moving beyond models limited to statistical associations [25] aiming to model the underlying structural properties of the data by applying the causality framework to machine learning [26 277. CRI aims to construct latent variables that capture the underlying causal mechanisms, allowing for the discovery of causal representations within observed data [8]. It seeks to deploy robust models in OOD scenarios, ensuring reliable performance even when the data distribution shifts. For example, leveraging the stability of causal mechanisms across different environments, several studies have utilized the invariant properties of causal representations under distribution shifts to enhance model performance in OOD scenarios through invariant prediction 12 28
Furthermore, there has been ongoing research into utilizing deep generative models to explicitly represent causal variables. Notably, under the assumption of independent causal mechanisms [0O], several studies have modeled these mechanisms as separate, independent modules or have focused on
learning disentangled and interpretable representations [1 29 30 Research has evolved toward learning causal representations that maintain stable mechanisms under distribution shifts, to improve OOD generalization [13]. Inspired by CRI,, we construct a novel ICL. framework using causal knowledge for OOD generalization. To build a robust demonstration set, we utilize the invariant causal representation constructed by a Variational Autoencoder (VAE) [3 bbsed model [3 32
3 Methodology
3.1 Generative model and inference model
In CCL, we consider several key variables: the task variable t represents the specific task being performed. The latent causal variable c repressents the fundamental context of the query. It is generated from the task variable and serves as a causal factor for both the input query æ and the (ground truth) answer y. Additionally, we introduce the latent source variable which influences components of x that are unrelated to the task, such as the structure of the text. The environmental variable e acts as an observable proxy for the latent source variable It represents contextual attributes of the data, such as the dataset's 's origin or the language used.
Note that both latent variables, and 8, generate v, where represents task-specific information, and 8 represents domain-specifie information. That is, we assume that the domain shift in the observed data is induced by changes in s, while remains invariant, shown in Figure 2
3

---

Phase Causal representation learning with In-pool (In-distribution) dataset
0
è
demonstration candidate
Φ
02
Φ
Dutaset
0g
9
Causal Representation
Phase 2: Causal-context learning
space of
Instance-level sample selcction
target query
Pc
Demonstration
Figure 3: Our proposed method, Causal-aware In-Context Learning (CCL), utilizes causally related demonstration sets to enhance performance on out-of-distribution (OOD) datasets. (Phase 1) First, optimize novel VAE-based causal representation learning method to capture the causal represen- tations of a given in-distribution dataset. After optimization, we store the causal representations, produced by the optimized model for the in-distribution dataset. (Phase 2) Second, CCL captures the causal representation, c", of the target query and selects the appropriate demonstration sets by comparing c and c²
We aim to model the joint distribution observed vari- ables {æ, y, e. along with latent variables s). We assume the generative model
t
(a e, = p,,x,t,e,co ( cc,
C
S
where pe e) is an invariant causal mechanism. We let denote all parameters of the generative model. We denote the unknown true source-domain distribution pa- (2, y, t, e), and we approximate with po(x, y, tt e).
x
y
x
(a) Inference model (b) Generative model
Figure 3 illustrates the overall workflow of CCL in two phases. In Phase 1, we learn causal representations from an in-distribution (ID) dataset using our VAE-based model: the inference networks db and Φ. infer the latent variables senvironmentrellatee and (lask-related), respectively, while the decoders 6₂, 02, 0g reconstruct the observed variables. This process yields the causal embeddings for the ID data. In Phase given a target query (x* t²), we apply Φ. to obtain its causal embedding c*. Comparing c* with the stored causal embeddings c, CCL then selects the most relevant demonstration examples, those with similar causal factors, to construct the prompt context. This causal representation approach ensures that our examples align with the true causal structure of the query, thereby improving model performance even under distribution shifts.
Figure 2: Graphical model CCL. The generative model shows that influences the latent causal variable c, which in turn directly affects both and y.
3.2 Learning causal representations via variational inference
Since direct maximization of log Pog,,( y, e) é) is often intractable due to the latent variables, we employ variational inference. We introduce tractable inference model (c, s v, e), where are the variational parameters. The standard Evidence Lower BOund (ELBO) on log y, is:
Po N, (,,
log y, U, = log pa(a, k,c, = log (gole.ole. gole, 83 1,
Eg (cssl.210)))))
log go(c, x, x,
LELBO
Maximizing this ELBO with respect to both 0 and @ yields tight approximation when go(e, I, z,y,ter e) RC pe(c. x, y, e).

---

Since 0" is unknown, we instead optimize the ELBO using the observed data distribution in the source domain, Pp y, e)):
max (ew.te))pppptttttt [TLLBO] 0,0
(I)
3.2.1 Reformulating variational inference for unobserved
At test time, y is always unobserved, as it is the target variable we aim to infer. While one common approach, such in CEVAE [3), is to introduce auxiliary model to explicitly predict y, we instead modify the objective function to enable variational inference without conditioning on y. Specifically, we factorize the inference model:
qo(c,s, 2,t,e) = (o x. xx,, (,,t,e) pa [ c).
which reflects the conditional independence y L (z, i, e, s) c. This design is key, as it directly injects the generative model's s causal assumption (c y) into the inference process. It serves to constrain the inference model qo to find a that consistent with Po c). the actual causal mechanism from the generator. This formulation allows to marginalize out y. By applying this factorization and Bayes' rule to the standard ELBO, we analytically marginalize out the unobserved y, reformulating the objective to depend only on (c s xr t) (see Appendix A for the full derivation). We define @ Qulz,,, = (c.2,,)) (po (ylc) as the implicit predictive distribution of The final objective of CCL given by:
max 8,6
Epp(axxt.e) log vix,,
Beste.llo.e log p, no (o, 1,
We construct the reconstruction model pe following the generative structure outlined in Figure 2b Implementing Equation requires this model, which composed decoders (e.g., po(r c, s), Poly c). pa(e {)) that reconstruct the observed variables from the latent variables. This reconstruc- tion process, particularly the Po / c) mechanism, ensures that the learned causal representation c effectively captures task-relevant information.
3.3 Regularization and conditional prior
In practice, to prevent unintended dependencies between e and during training, we further employ Maximum Mean Discrepancy (MMD) 34 loss regularization term [9. Additionally, the task variable (the parent of e) is treated as an observed input, not a latent variable requiring posterior inference. Instead, following the iVAE K5 framework, we define conditional prior po(c () for the generative model based on this observed t. Our variational inference formulation follows the approach proposed in 322
3.4 Theoretical analysis
Prioritizing demonstrations that are causally similar the query yields provably better in- in-context learning (ICL) than prioritizing demonstrations that are merely input similar. We show that input nearest selection can induce large label discrepancies even when inputs are arbitrarily close in Theorem 3.3 Furthermore, Theorem 3.4 provides both a theoretical explanation and practical guideline: prooritzzigg causally similar examples key to robust ICI.
Our analysis begins by assuming the data-generating process is modeled using an SCM : (S) and collection S of assignment equations as follows [T]]
=
E c: fc(,
::= E :== Ee( I fa(c,s,Eu), y==ff(ce)) (3)
Here, εt, Eas Sog En E Rd are random vectors with > and E, E R a random variable. We assume Eet, Eas En, satisfies joint independence. The parents of are and s, while y has only e as its parent. The causal graph is achieved by drawing edges from RHS variables of Equation 3 to LHS variables except the noise variables
5

---

We adopt a linear setting in line with [36], who demonstrate that attention-based updates in LLMs can be approximated by steps of gradient descent with convex loss on linear parameter with respect to u: x. Although real-world LLMs are more complex, the linear approximation provides clear analytical framework.
Assumption 3.1 Linear-causal assumption). We formalize simplified data Gata-generating process via the following linear-cause assumption:
24 == Bi c, + B2 s; + Ez,t, ys = (w") + Ey Each coordinate of E, is -u--ub-Gan Gaussian, and Ey. is matrices C2 and 8t, 24, Cj2 and 8; are d-dimensional vectors and yi is scalar.
o arj-sub-Guussing B1 and B2 denote coefficient
prerequisite for ICL. to construct a demonstration set Ds = (il,, from the training dataset Dy = [(1:, (i)) where S I is the selected index set. A common strategy, forming the set Dr, selects pairs (2,, 3/)) by assessing how similar 2, is to the input query z", with the expectation that y" will be similar to 14 [15]. Our method leverages latent causal variables:
associate the query x² with a causal variable c", and each training pair (Iy, y1) with its own causal variable a, We then select pairs whose c; lie close to c', forming causally similar set De. All our theorems are based on Assumption 3.1
Definition 3.2 (Demonstration sets by strategy). Let sim(:.) be similarity measure and N the size of the demonstration set. We define:
De BCL.S|=N
argmax Σ
sim(c,, c²), D3 argmax Σ sim(x,,,")
(4)
İES
SCI, S|=N iES
Theorem 3.3 (Input proximity can lead to prediction discrepancy). Let ,y") and (z, y) be two samples potentially generated by different latent pairs (c" and (c, s). Under Assumption B. B.2 in Appendix B, for every > 0, there exists a i. > O such that
| - cll |Bplaap
K
11 - v
7
where || llop is the operator norm
|B
for some constant 7,. if || - x1| < e. In other words, one can make ||2- -xl arbitrarily small while allowing Jyy" y|l remain arbitrarily large, due to the interplay between (c" ,s)) and (c, s).
Theorem 3.3 shows that even when the distance between and x is made arbitrarily small, the distance between the corresponding y" and can still be significant, as there is no upper bound on this gap. Consequently, the predicted value based on x² may coincide with y, causing a discrepancy with the true y²
Picking demonstrations from De yields better in-context learning than picking from Dz.. The upper bound of the estimation error of the learned parameter is smaller compared to that of input-based selection. Furthermore, the upper bound the test prediction error with CCL is also smaller. The parameter update in ICL, under a transformer architecture, is approximated by gradient descent on the demonstration set, following the formulation in (3). Let w(() be the weight updated via M steps of gradient descent using the empirical risk on D... and let (M) be the corresponding weight updated from Dr.-
Theorem 3.4 (Performance the c-similarity). For sufficiently large N, M, with probability least 1 8yat, the following holds under Assumption C.1-C.4 in Appendix C:
1. Tighter upper bound on estimation error. The estimation errors admit upper bounds U Uprro and U such that
U param- Uc param (1 Amin c)) CuSN and U = (1/>min (Ta)) ()))) CuSN. Cu is a some constant and SN / (lgg(//mmu)/// Amin(A) denotes the minimum eigenvalue of a matrix A. Te and I. are the empirical second moment matrices of DD and D
(lw)(M) w² S U param
( w0 || S Uperam: and Upanan
2. Tighter upper bound on test error. For the test query (x* "."y"' yy) with y² -====wTTee Tw") + E the prediction errors admit upper bounds Vies and U yea such that
[(w[M) x²
U
[ww
v²1 Uie. and U < U
Use [U |R| and test
||x" ( puraam + RR). R (w*) T x² - y

---

Theorem 3.4 shows that, with high probability, the parameter error [lw/ w"1l and the test error (
y admit upper bounds that are tighter than the corresponding bounds obtained from Da. In essence, when D. is used for demonstrations, the underlying design matrix becomes better conditioned with respect to c, mitigating the confounding effect and leading to tighter error bounds.
l[
4 Experiments
We validate the effectiveness and validity of CCL by addressing three main points. First, Section 4.2
verify that the latent variables and inferred by CCL indeed capture domain-ing invariant and domain- varrint features, respectively, for modeling the causal factors of æ. In Section [.3] we examine whether the samples characterized by exhibit similarity to the test samples convey the same underlying intent. In Section 4.4, we evaluate how the demonstration sets constructed using CCL enhance in-context learning performance under OOD scenarios. In Section 4.5 we qualitatively analyze how the latent features c and s capture distinct features. Lastly, in Section [4.6] we investigate the capability of CCL new more intricate reasoning tasks and perform sensitivity analysis.
4.1 Experimental setup
We adopt query-dependent demonstration strategy that dynamically selects the suitable examples for each test input. After embedding test query, we compute its cosine distances to all candidates in the in-distribution training pool. In the K-nearest-neighbo (KNN) variant, the K closest instances, where K equals the predefined shot size (S2), are selected directly. We also investigate K-means-based selection method that is governed by two hyperparameters, R and P. A proportion of the shot budget is allocated to the most similar instances, obtained exactly in the KNN procedure. The remaining budget K = Ω - is filled by clustering: among the next P (with E {50, 100, 300}) most similar candidates, we run K-means clustering and, from each cluster, select the sample whose embedding is closest to the centroid. This combined strategy yields prompts that simultancously maintain high relevance the query while covering broader range of semantic regions.
4.2
Synthetic data
ID Task Comparison
Env. Comparison
OOD Task Comparison
Method
Acc.
NDCG
F1
Acc
NDCG
F1
Acc
NDCG
Fl
57.7
71.5
58.6
85.7
91.3
86.0
45.0
57.9
45,4
CVAE (z)
33.3
60.2
32.5
33.1
43.2
32.2
32.4
53.6
33.9
Oracle (c)
100.0
100.0
100.0
33.5
48.7
33.7
100.0
100.0
100.0
CCL (c)
100.0
100.0
100.0
40.8
55.0
39.9
100.0
100.0
100,0
Oracle (s)
33.3
48.9
33.9
100.0
100.0
100.0
32.7
51.3
33.0
CCL (s) 36.2
51.4
36.2
100.0
100.0
100.0
33.2
48.3
31.9
Table I: Retrieval experiments synthetic data show that CCL consistently outperforms alternatives on both in-distribution and iittof-disttibutiin task queries, confirming that e captures the underlying causal structure of the tasks. Conversely, when retrieval is conditioned on environment labels, s-based retrieval excels, highlighting their sensitivity to domain- speccifc factors. CCL's learned representation, CCL (c), tracks the ground-truth causal feature particularly closely.
We construct synthetic data with three tasks and five environments. Following Figure 2b we first define the root nodes: the task variable t and the s variable. We enforce independence among task embeddings by randomly initializing them with orthogonality constraints, applying the same approach to s. Then, we generate the c embedding using a three-layer fully connected neural network that takes t input and add random noise to its output. Other variables follow a similar process. We train the neural networks, viewed as non-linear data- sittaaanneraingg functions, using contrastive learning to ensure that is similar within the same task and is similar within the same s, while enforcing dissimilarity across different tasks or environments.
To better reflect realistic scenarios, we consider similar tasks environments. Specifically, for the root nodes / and 8, we set the cosine similarity between any two or embeddings to a value between 0 and (in our experiment, we use 0.7). During contrastive training of the generating functions,
7

---

we adjust the loss weights to reduce the penalty for similar tasks or environments, ensuring their embeddings are not pushed too far apart.
Table D presents the proportion of retrieved samples whose task or environment (Env.) matches that of the target input, across different embedding types, under both in-distribution (ID) and out- of-distribution (OOD) settings, Additional experimental results and discussions on the synthetic experiments are provided in Appendix D.
4.3
MGSM
Metric
embedding
cembedding
Total Accuracy
81.03
85.84
Method
Total
ID
OOD
ID Accuracy
97.05
99.74
ZS
87.71
89.43
84.70
OOD Accuracy
53.00
61.52
ICL (Fix.)
91.20
91.26
91.10
Total NDCG
86.00
88.73
ICL (KNN)
94.07
95.83
91.00
ID NDCG
99.12
99.89
CCL
94.55
96.11
91.80
OOD NDCG
63.03
69.21
(h) Comparison of performance. 7S denotes the zero- shot baseline, ICL (Fix.) uses a fixed demonstration set. ICI. (KNN) and CCL utilize KNN retrieval
(a) Comparison of retrieval accuracy and NDCG for and embeddings MGSM in the 5-shot setting.
Table 2: (a) compares five-shot MGSM retrieval performance between embeddings derived from the original inputs x and from the causal features, (b) reports overall, in-distribution (ID), and out-of- distribution (OOD) accuracies for four prompting regimes-7ero-sher (ZS), fixed demonstrations, KNN-based retrieval, and CCL..
As another dataset to evaluate the performance of our methodology, we employ the MGSM (Multilin- human-annotated
gual Grade School Math) dataset [37]. The MGSM dataset is human annotated translation of 250 problems from the GSM8K dataset [38 into ten different languages.
Utilizing the MGSM dataset, our goal is to evaluate the precision with which CCL deduces latent variables c, that represent the fundamental context of problems. For this purpose, we evaluate the retrieval performance by examining how correctly the model retrieves the same problem given specific question.
First, we extract embeddings for each question using OpenAl's -ewemmbeddng---mmaaee model. Based on these embeddings. we split the data into an ID and an OOD dataset. We use Swahili, Thai, Telugu, and Bengali for the OOD dataset, while the remaining languages are designated as ID. We provide detailed explanation of the classification criteria in Appendix D.
In this experiment, we define the problem category as the task L. The categories include six classes, such "Arithmetic Operations" and "Geometry and Measurements". These categories generated by labeling each question using OpenAl's ol, followed by human verification. During the labeling process, only English questions are labeled, and the same labels are directly applied corresponding questions in other languages.
Table 2a presents the retrieval performance of the embeddings and the embeddings. We eval- uate how accurately each method retrieves the same problem in a different language. The results demonstrate significant improvement in accuracy and NDCG for both ID and OOD when using approach instead of x embeddings.
Next, we perform ICL. based on the retrieval results. In the MGSM dataset, evaluate performance by measuring the model's prediction accuracy. Similarly to the retrieval process, we use 5-shot setting assess performance and compare zero-shot (ZS), ICL (Fixed sample, KNN) and CCL. Unlike ICL (KNN) and CCL, which can retrieve samples from different languages, ICL (Fix.) uses predefined samples specific each language. We use GPT-4o-mini for in-context learning, We refer to Appendix for a detailed explanation of the MGSM experiment.
Table 2b illustrates the experimental results. The results demonstrate that CCL-based retrieval for in-context samples achieves higher accuracy in both ID and OOD settings than other approaches. This aligns with the strong retrieval performance of embedding indicated in Table 2a demonstrating that selecting in-context samples based on the latent causal feature is crucial for problem solving and improves in-context learning accuracy.
8

---

4.4
Generalization across tasks and domains
Language model
Retrieval method
QNLI
PIQA
WSC273
YELP
Avg.
ZS
43.36
71.33
55.31
88.98
64.75
LLM-R
29.93
69.91
61.17
79.48
60.12
Llama-3.2-3B-IT
ICL (K-means)
68.13
69.04
49.82
75.81
65.70
CCL
75.18
70.46
61.91
95.44
75.74
7S
86.34
76.01
64.10
95.76
80.55
LLM-R
85.21
74.10
65.93
96.37
80.40
Phi-4-mini-IT
ICL (K-means)
83.18
74.81
71.06
96.25
81.33
CCL
82.26
75.73
71.43
96.33
81.44
7S 91.30
94.07
90.84
97.47
93.42
LLM-R
90.32
94.23
92.67
98.27
93.87
GPT-4o
ICL (K-means)
88.28
93.04
87.55
98.17
91.76
CCL
90.77
93.15
93.77
98.36
94.01
Table 3: Out-of-distribttion accuracy QNLI, PIQA, WSC273, and Yelp for three language Booelss-Llama-.2333B--T Phi-4-mini-IT, and GPT-4o-under four prompting regimes: zero-shot (ZS), the learned-retriever baseline (LLM-R), and two K-means-based retrieval approaches, vanilla ICL and CCL. Bold numbers denote the highest score each column, and italics denote the second highest. CCL attains the best average accuracy for every model, with particularly pronounced improvements for the smaller Llama-3.2- 3.2-3B-IT.
We evaluate whether examples selected by CCL improve performance on OOD NLP tasks. Adopting the experimental protocol of LLM-R [39]. we compare against their retrieval method but instead assess the generated outputs rather than relying on token probabilities. Our approach retrieves examples with similar embeddings via KNN, clusters them using K-means, and selects the cluster centers as final candidates. As shown in Table 3 CCL consistently yields strong performance across diverse OOD tasks. We follow the same 8-shot setting used in LLM-R R to ensure a fair comparison.
4.4.1 Sensitivity to the embedding models
Language model
Emhedding model
QNLI
PIQA
WSC273
YELP
Avg.
text-cmbbeddinng---malll mullllingull-c--lage--ins
82.26
75.73
71.43
96.33
Phi-4-mini-IT
82.26
75.25
73.99
95.72
81.81
Table 4: CCL accuracy four out-of-distribution benchmarks when the same language model and the Scores given for each task and averaged; bold indicates the highest score per column, and italics the second-highed The multilingual-e5 encoder attains the top overall score, yet the gap is small, indicating that CCL remains robust the choice of embedding model.
(Phi-4-mini-IT)) is paired with two embedding models (OpenAl's text-embedding---sime mantillnguale---lage--nnmn
To evaluate CCL's sensitivity the encoder, using
the entire pipeline across the NLP benchmarks [40], an open-source embedding model that ranks among the top performers the MTEB text-embedding leaderboard [4T]. Table Tableefexerrimenatlly demonstrates that CCL maintains comparable performance despite changes in the embedding model, highlighting its robustness in inferring causal features.
4.5 Qualitative analysis
We provide a qualitative analysis of the learned latent features to better understand how e and s are interpreted in practice. To visualize the semantics encoded in these variables, we decode sentence embeddings while zeroing out one latent dimension. Specifically, we first infer e and 8 from an input embedding x. We then set 0 to generate z'=oo which highlights the domain-invariant features represented by e. Similarly, we set 0 to generate 2'-0- which reveals the domain-variant information captured by Table lists the top-5 nearest words to each decoded embedding.

---

2==0
<==0
x
2/=0
===0
horribleappetizers
unappetizing flavorless
review reviewers critiques soggy
dvd
unusable reverb expired throw
pancakes
eject
potatos
horribleappetizers inedible
disks
cancelled film crappy
hadhorrible
unusable
review trip
bad
trashed
reviews
purchased
trashed
(a) Original negative sentence is "the red velvet part cakes were horrible and brown, and potatos were over cooked and bland.. would not recommend"
(b) Original negative sentence is "Worked for about 4 months. DVD player will not eject accept disks. Do not buy."
Table 5: Top-5 nearest words on Yelp and Amazon. The sentence embedding x captures both semantic and contextual tokens. In contrast, z'=o clusters strongly around negative sentiment expressions, while =o= clusters tokens associated with contextual metadata.
4.6
Generalization and sensitivity analysis
4.6.1 Advanced tasks
Table presents presenss the generalization capability of CCL across advanced tasks. The unseen generation task involves sentiment reversal paraphrasing: the model rewrites negative sentence to express the opposite sentiment, and automatically assess its sentiment using GPT-4o-min 4iiiii. Although CCL trains only on classification tasks, it generalizes well to this unseen generation setting. For MMLU 422, we retrieve five examples for each query without distinguishing among the 57 domains. For HotporQA 43,, we provide each query with its corresponding document and retrieve examples to form document-example pairs. This experiment provides evidence that CCL may help with hierarchical and composite ainguag--undesttande problems.
Unseen & generation Reasoning Multi-hop QA
Yelp
Amazon
MMLU
HotpotQA
7$ 86.26
86.73
60.48
82,43
ICL 87.68
85.80
61.37
84.
CCL 90.05
87.70
61.52
84.43
Table Performance comparison of 7S, ICI,
and CCL across tasks using Phi-4-mini-IT.
4.6.2
Sensitivity analysis
QNLI
PIQA
WSC273
YELP
dim(
QNLI
PIQA
WSC273
YELP Avg.
ZS
43.4 0.00) (1 0.00)
54.3 0.00)
89.0 0.00)
128
69.5
72.6
56.1
93.5
72.9
LLM_R
29.9 (# 0.00) 69.9 (= 0.00)
(1.200
(=
256
75.3
71.6
60.4
94.9
75.6
ICL
68.1 0.00)
69.0
(2.8(2
(2.000)
CCL
75.2 (# 0.45) 72.4 1.12) 58.98 2.78) 95.10 (# 0.25)
1024 (ours)
75.2
70.5
61.9
95.4
75.7
(a) Mean accuracy and std over random seeds.
(b) Accuracy variation w.r.t. dim(c).
Table 7: Performance of CCL under different training conditions using Llama-3.2-3B-IT. (a) OOD benchmark accuracy across five random seeds, showing stable results despite stochastic variation in VAE training. (b) Performance changes with respect to latent dimensions, indicating that smaller dimensions do not significantly degrade accuracy.
Table 7a shows the OOD benchmark results under different random seeds used for training the VAE within CCL. Since response generation is deterministic (non-sampling). other baselines exhibit zero variance. Table (ab reports the effect of varying the latent dimensions of e and s during VAE training. The results suggest that model performance remains stable even with smaller latent dimensions.
5 Conclusion and discussion
We propose CCL, the first framework to integrate causal representation learning into ICL, addressing key limitation of conventional ICL in OOD settings. By selecting demonstrations based on causal representation rather than surface-level similarity, CCL improves robustness, and parameter estimation, with theoretical guarantees.
Limitation and Impact statement. Since CCL employs a VAE-based latent embedding, inherent structural limitations of VAE may hinder its ability to fully capture the rich and nuanced representations of natural language. We leave the deeper integration of embedding retrieval with causal inference future work.
10

---

Acknowledgments
This work was supported by the National Research Foundation of Korea(NRF) grant funded by the
Korea govemmee((MSTTySSS-2224-00000222
References
[1] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877-1901. Curran Associates, Inc., 2020.
[2] Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, and Zhifang Sui. A survey on in-context learning, arXiv preprini arXiv:2301.00225: 2022.
[3] Aaron Mueller, Albert Webson, Jackson Petty, and Tal Linzen. In-context learning generalizes, but not always robustly: The case of syntax. arXiv preprint arXiv:2311 1..0v22 2023.
[4] Arian Hosseini, Ankit Vani, Dzmitry Bahdanau, Alessandro Sordoni, and Aaron Courville. On the compositional generalization gap of in-context learning. arXiv preprint arXiv:221 2022.
[5] Lifan Yuan, Yangyi Chen, Ganqu Cui. Hongcheng Gao. Fangyuan Zou, Xingyi Cheng, Heng Ji, Zhiyuan Liu, and Maosong Sun. Revisiting oft-s--distiibutiin robustness in nlp: Benchmarks, analysis, and Ilms evaluations. Advances in Neural Information Processing Systems, 36:58478- 58507, 2023.
[6] Kartik Ahuja and David Lopez-Paz. A closer look at in-context learning under distribution shifts. arXiv preprint arXiv: /6704, 2023.
[7] J Pearl. Causality. Cambridge university press, 2009.
[8] Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Toward causal representation learning. Proceedings of the IEEE, 109(5):612- 634, 2021.
[9] Jivat Neet Kaur, Emre Kiciman, and Amit Sharma. Modeling the data-generating process is necessary for out-of-distribution generalization. arXiv preprint arXiv:2206.0777 07837, 2022.
[10] Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: founda- tions and learning algorithms. The MIT Press, 2017.
[11] Giambatisst Parascandolo, Niki Kilbertus, Mateo Rojas-Carulla, and Bernhard Schölkopf. Learning independent causal mechanisms. In International Conference on Machine Learning, pages 4036-4044 PMLR, 2018.
[12] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk mini- mization. arXiv preprint arXiv: 1907002997 2019.
[13] Chaochao Lu, Yuhuai Wu, José Miguel Hernández-Lobat and Bernhard Schölkopf. Invariant causal representation learning for out-of-distribution generalization. In International Conference on Learning Representations, 2021.
[14] Kartik Ahuja, Amin Mansouri, and Yixin Wang. Multi-domain causal representation learning via weak distributional invariances. In International Conference on Artificial Intelligence and Statistics, pages 865-873. PMLR, 2024.
[15] Jiachang Liu, Dinghan Shen, Yizhe Zhang, Bill Dolan, Lawrence Carin, and Weizhu Chen. What makes good in-context examples for gpt-3? arXiv preprint arXi::2101.06804, 2021.
11

---

[16] Jiacheng Ye, Zhiyong Wu, Jiangtao Feng, Tao Yu, and Lingpeng Kong. Compositional ex- emplars for in-context learning. In International Conference on Machine Learning, pages 39818-39833. PMLR, 2023.
[17] Sharut Gupta, Stefanie Jegelka, David Lopez-Paz, and Kartik Ahuja. Context is environment. In The Twelfth International Conference on Learning Representations, 2023.
[18] Tianxiang Sun, Yunfan Shao, Hong Qian, Xuanjing Huang, and Xipeng Qiu. Black-box tuning for language-men Lunguage-model-asse moddel-as-a-service In International Conference on Machine Learning, pages 20841-20855. PMLR, 2022.
[19] Zihao Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. Calibrate before use: Improving few-shot performance of language models. In International conference on machine learning, pages 12697-12706. PMLR, 2021.
[20] Yao Lu, Max Bartolo, Alastair Moore, Sebastian Riedel, and Pontus Stenetorp. Fantastically ordered prompts and where find them: Overcoming few-shot prompt order sensitivity. arXiv preprint arXiv:2 2021.
[21] Ohad Rubin, Jonathan Herzig, and Jonathan Berant. Learning to retrieve prompts for in-context learning. arXiv preprint w:Xiv:2112.0863: 2021.
[22] Yiming Zhang, Shi Feng, and Chenhao Tan, Active example selection for in-context learning. arXiv preprint arXiv:2211.044880 2022.
[23] Sarthak Mittal, Eric Elmoznino, Leo Gagnon, Sangnie Bhardwaj. Dhanya Sridhar, and Guil- laume Lajoie. Does learning the right latent variables necessarily improve in-context learning? arXiv preprint arX11:2405.19162, 2024.
[24] Judea Pearl and Dana Mackenzie. The book of why: the new science of cause and effect. Basic books, 2018
[25] Judea Pearl. The seven tools of causal inference, with reflections on machine learning. Commu- nications of the ACM, 62(3))54-60, 2019.
[26] Jean Kaddour, Aengus Lynch, Qi Liu, Matt J Kusner, and Ricardo Silva. Causal machine learning: A survey and open problems. arXiv preprint arXiv:2206. 15475, 2022.
[27] Aneesh Komanduri, Xintao Wu, Yongkai Wu, and Feng Chen. From identifiable causal repre- sentations to controllable counterfactual generation: A survey on causal generative modeling. arXiv preprint arXhv:2310.10011 2023.
[28] Jovana Mitrovic, Brian McWilliams, Jacob Walker, Lars Buesing, and Charles Blundell. Repre- sentation learning via invariant causal mechanisms. arXiv preprint arXiv:2010.07922, 2020.
[29] Michel Besserve, Arash Mehrjou, Rémy Sun, and Bernhard Schölkopf. Counterfactuals uncover the modular structure of deep generative models. arXiv preprint 1rXiv:1812..0225 2018.
[30] Mengyue Yang, Furui Liu, Zhitang Chen, Xinwei Shen, Jianye Hao, and Jun Wang. Causalvae: Structured causal disentanglement in variational autoencoder. arXiv preprint 20X1v22004.08697, 2020.
[31] Diederik Kingma. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 2013.
[32] Chang Liu, Xinwei Sun, Jindong Wang, Haoyue Tang, Tao Li. Tao Qin. Wei Chen, and Tie- Yan Liu. Learning causal semantic representation for -isto--disti prediction. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 6155-6170. Curran Associates, Inc., 2021.
[33] Christos Louizos, Uri Shalit, Joris M Mooij, David Sontag, Richard Zemel, and Max Welling. Causal effect inference with deep latent-variable models. Advances in neural information processing systems, 30, 2017.
12

---

[34] Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. Journal of Machine Learning Research, 13(25):723-773, 2012.
[35] Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational au- toencoders and nonlinear ica: unifying framework. In International conference on artificial intelligence and statistics, pages 2207-2217. PMLR, 2020.
[36] Yu Bai, Fan Chen, Huan Wang. Caiming Xiong, and Song Mei. Transformers statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 2024.
[37] Freda Shi, Mirac Suzgun, Markus Freitag, Xuezhi Wang, Suraj Srivats, Soroush Vosoughi, Hyung Won Chung. Yi Tay, Sebastian Ruder, Denny Zhou, et al. Language models are multilingual chain-of- thought reasoners. arXiv preprint arXiv:2210.03057, 2022.
[38] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 14168, 2021.
[39] Liang Wang. Nan Yang. and Furu Wei. Learning to retrieve in-context examples for large language models. In Proceedings of the 18th Conference of the European Chapter of the Association for Computational Linguisties (Volume I: Long Papers), pages 1752-1767, 2024.
[40] Liang Wang, Nan Yang, Xiaolong Huang, Linjun Yang, Rangan Majumder, and Furu Wei. Multilingual e5 text embeddings: A technical report. arXiv preprint arXiv:2402.055622 2024.
[41]
Kenneth Enevoldsen, Isaac Chung, Imene Kerboua, Márton Kardos, Ashwin Mathur, David Stap, Jay Gala, Wissam Siblini, Dominik Krzemióski, Genta Indra Winata, Saba Sturua, Saiteja Utpala, Mathieu Ciancone, Marion Schaeffer, Gabriel Sequeira, Diganta Misra, Shreeya Dhakal, Jonathan Rystrøm, Roman Solomatin, Ömer Çağatan, Akash Kundu, Martin Bernstorff, Shitao Xiao, Akshita Sukhlecha, Bhavish Pahwa, Rafal Poswiata, Kranthi Kiran GV, Shawon Ashraf, Daniel Auras. Björn Plüster, Jan Philipp Harries, Loic Magne, Isabelle Mohr, Mariya Hendriksen, Dawei Zhu, Hippolyte
Tom Aarsen, Jan Kostkan, Konrad Wojtasik, Taemin Lee, Marek Suppa, Crystina Zhang, Roberta Rocca, Mohammed Hamdy, Andrianos Michail, John Yang. Manuel Faysse, Aleksei Vatolin, Nandan Thakur, Manan Dey, Dipam Vasani, Pranjal Chitale, Simone Tedeschi, Nguyen Tai, Artem Snegirev, Michael Günther, Mengzhou Xia. Weijia Shi, Xing Han Lü. Jordan Clive, Gayatri Krishnakumar, Anna Maksimova, Silvan Wehrli, Maria Tikhonova, Henil Panchal, Aleksandr Abramov, Malte Ostendorff, Zheng Liu, Simon Clematide, Lester James Miranda, Alena Fenogenova, Guangyu Song. Ruqiya Bin Safi, Wen-Ding Li, Alessia Borghini, Federico Cassano, Hongjin Su, Jimmy Lin, Howard Yen, Lasse Hansen, Sara Hooker, Chenghao Xiao, Vaibhav Adlakha, Orion Weller, Siva Reddy, and Niklas Muennighoff. Mmteb: Massive multilingual text embedding benchmark. arXiv preprint arXiv:2502. 13595, 2025.
[42] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. Proceedings of the International Conference on Learning Representations (ICLR), 2021.
[43] Zhilin Yang. Peng Qi, Saizheng Zhang, Yoshua Bengio, William W Cohen, Ruslan Salakhut- dinov, and Christopher D Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. arXiv preprint arXiv: /000::/800990000 096000 2018
13

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. or [NA]-
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1 (1-2 sentence) justificatin right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked While "[Yes] " is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] providded a proper justification is given (e.g., "error bars expensive" or "we were unable to find the license for the dataset we used"). In general, answering binary way,
the checklist one of the factors in their evaluation.
reported because it would be too computationally
"[No] " or "[NA]" is not grounds for rejection. While the questions are phrased acknowledge that the true answer often more nuanced, so please just use your best, judgment and write justification elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: In the abstract and introduction of the paper, we compare and introduce our core contribution, which is the first study to alleviate the difficulties of existing ICL methods in OOD in-context learning from causal representation perspective, and validate the significance of our methodology through theoretical validation and experiments.
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 the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long it 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?
14

---

Answer: [Yes]
Justification:
We discuss our limitations in the conclusion 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 a separate "Limitations" section 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 well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect 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 on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails 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 fairesss
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't It acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play 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 proof's
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We proceed with the theoretical analysis of our methodology in Section 3. The necessary assumptions and definitions for the theoretical analysis are mentioned in the main text, and the proofs are given in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs 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 if they appear in the supplemental material, the authors are 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 be properly referenced.
4. Experimental result reproducibility
Question: 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)?
Answer: [Yes]
15

---

Justification: We mention specific experimental setups in the text and appendix, and ensure reproducibility by providing a code repository.
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 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 can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS 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) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome 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., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: 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?
Answer: Yes]
Justification: We provide all the code and data we used in our experiments in a code repository.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
mublic/guides www.dee//ooeuuuiisseeeeeeeeeeeee Cooeeuumissioooooll eeeiiinyll for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an 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 NeurlPS code and data submission guidelines (https for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
16

---

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they 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: We describe in the text how we selected examples from ICL and which linguistic models we experimented with. For the sake brevity, the specific experimental setup can be found in the appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to a level detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: We report the mean and standard deviation of performance metrics over five random seeds to assess the consistency our method. However, we did not perform formal statistical significance tests or report confidence intervals.
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should 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).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It 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 enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
17

---

Answer: [Yes]
Justification: We mention in the appendix the resources required for the experiments due to the constraints of the paper length.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs 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 or 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
Answer: [Yes]
Justification:
We followed the EthicsGuidelines for our review. Reviewed according to the
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
We do not discuss this paper due to the lack of relevant experiments.
Guidelines:
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 no 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., deploymeet 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 partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out. For example, it is legitimate to point that 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 funetioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
18

---

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., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We didn't do anything differently because we were already using a widely used benchmark.
Guidelines:
The answer NA means that the paper poses 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 encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
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?
Question:
Answer: [Yes]
Justification:
We cite all code and data we reference.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version asset used and, possible, include URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, wwwpsswieeesseeeeeeeeee thcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset'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: Not relevant.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mooel part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
19

---

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 to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [No]
Justification: Not relevant.
Guidelines:
The answer NA means that the paper does 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 the NeurIPS 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
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?
Question:
Answer: [NA]
Justification:
Does not describe
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
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 NeurlPS 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
Question: 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.
Answer: [Yes]
Justification:
We got some help with simple grammar corrections and LaTeX syntax from
LLM.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non standard components.
Please refer to our LLM policy (https: htt.a//p.....coooeeeee/e////// for what should or should not be described.
20


Reviews:
[{"review_id": "mT45lyFt4y", "reviewer": "Reviewer_CojG", "summary": "The paper introduces Causal-aware In-context Learning (CCL), a two-phase framework that first learns task-invariant latent causal variables with a VAE-style model and then selects demonstrations whose causal embeddings are closest to the query. Experiments on synthetic data, MGSM math problems, and four NLP benchmarks (QNLI, PIQA, WSC273, Yelp) across three LLMs verifie the effectiveness of proposed methods.", "questions": "NAN", "limitations": "NAN", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths \n\nPrincipled formula. Clear probabilistic model, disentangled latent factors, and derivation of an ELBO that removes the need to observe y at inference. These Principled formulation\n gives intuition and formal guarantees (within linear assumptions) for why causal retrieval helps.\n\nEmpirical evidence. This paper conduct on both synthetic and real datasets to very the effectiveness of proposed methods.\n\nWeaknesses\n\nInterpretability Challenge. While the paper adopts causal learning, the proposed latent causal variables lack interpretability. It remains unclear whether such simplified models can consistently learn domain-specific and causally invariant representations, particularly for complex tasks.\n\n\nInsufficient Experiments. The paper lacks comparison with state-of-the-art learning-based and learning-free retrievers (e.g., [1–2]). Additionally, baseline details are vague, and the evaluation is limited to a narrow set of tasks. Experiments on generation tasks would further strengthen the work.\n\nWriting and Clarity. The presentation would benefit from a clearer description of the complete algorithmic pipeline.\n\nNeed for Deeper Explanation. The method appears to leverage output labels (y) from the demonstration set in a way that resembles TTF [1]. A comparative analysis of the advantages and limitations of both approaches would be valuable.\n\nThe proposed method appears to have limited generalizability to tasks beyond those evaluated  (train on one task and test on the other task).\n\n\n[1] Unraveling the Mechanics of Learning-Based Demonstration Selection for In-Context Learning. ACL 2025\n\n\n[2] Compositional Exemplars for In-context Learning. ICML 2023.", "was_revised": true, "final_justification": "Since I have not received a reply from the authors, I still have concerns regarding the generalization of the proposed framework, and only improved my score to 4.", "rebuttal": "We sincerely appreciate your thoughtful and constructive review. Your comments reflect a deep understanding of our work, and we have made our best effort to address your concerns comprehensively. In doing so, we carefully referred to the papers you mentioned and conducted additional experiments to support our response.\n\n> Interpretability Challenge. While the paper adopts causal learning, the proposed latent causal variables lack interpretability.\n\n- We appreciate this insightful concern. In the vision domain, there has been significant research into causal representation learning using VAE-based models for OOD generalization [3,4].\n- However, to the best of our knowledge, **such approaches have not been widely studied in the natural language domain**. Our work aims to fill this gap by rigorously designing both the methodology and the objective function, and by demonstrating the validity of causal representation learning through controlled toy experiments.\n- That said, we agree with your suggestion that further analysis is needed to evaluate **whether a simple VAE structure can effectively disentangle domain-invariant and domain-variant features in natural language tasks**. In response, we conducted a qualitative analysis to interpret the meaning of the latent causal variables in our model.\n- As described in the paper, our goal is to decompose the latent features into components that are relevant to the task objective (domain-invariant) and components that are unrelated (domain-variant).\n- For instance, in a sentiment analysis task, domain-invariant features would correspond to expressions of sentiment (positive or negative), while domain-variant features might vary depending on the data source. For example, Yelp reviews may include aspects related to restaurant service or food quality, whereas Amazon reviews might focus on product quality or usability.\n\n- To qualitatively assess the interpretability of the latent variables $c$ (invariant) and $s$ (variant) in our VAE trained under the CCL framework, we visualized their effects via decoding. Specifically, given an input sentence embedding $x$, we inferred the corresponding $c$ and $s$. We then performed two decoding operations:\n  - Zeroing out $s$ to generate $x’_{s=0}$, and\n  - Zeroing out $c$ to generate $x'_{c=0}$.\n\n- Since $x_{s=0}'$ and  $x_{c=0}'$  lie in the same embedding space as $x$, we computed nearest neighbors based on embedding similarity to examine the semantics captured by $c$ and $s$.\n\nQualitative Example 1 (Yelp dataset):\nOriginal negative sentence: “the red velvet pancakes were horrible and brown, and potatos were over cooked and bland.. would not recommend”\n\nTable 1. Top 5 nearest words\n| $x$                | $x'_{s=0}$         | $x'_{c=0}$ |\n| --------           | --------           | --------   |\n| horribleappetizers | unappetizing       | review     |\n| pancakes           | flavorless         | reviewers  |\n| potatos            | horribleappetizers | critiques  |\n| hadhorrible        | inedible           | soggy      |\n| bad                | trashed            | reviews    |\n|           |            |    |\n\n- In Table 1, $x$ retrieves words related to both sentiment and specific content (e.g., “pancakes”, “potatos”).  In contrast, $x_{s=0}'$ shows clusters of **strongly negative sentiment expressions**. Meanwhile, $x_{c=0}'$ surfaces terms such as “review” and “critiques,” suggesting that domain-variant features relate to the meta-characteristics of the dataset (e.g., user review context).\n\nQualitative Example 2 (Amazon dataset):\nOriginal negative sentence: “Worked for about 4 months. DVD player will not eject or accept disks. Do not buy.”\n\nTable 2. Top 5 nearest words\n| $x$       | $x_{s=0}'$  | $x_{c=0}'$ |\n| --------  | --------    | -------- |\n| dvd       | unusable    | reverb   |\n| eject     | expired     | throw    |\n| disks     | cancelled   | film     |\n| unusable  | crappy      | review   |\n| purchased | trashed     | trip     |\n| |  |  |\n\n- Similarly, in Table 2, $x_{s=0}'$ captures **negative sentiment** towards product usability, while $x_{c=0}'$ retrieves terms more loosely associated with contextual metadata.\n\n- Our qualitative analysis indicates that CCL may learn to separate causal factors in natural language, even with a relatively simple VAE-based architecture.\n\n> Insufficient Experiments. The paper lacks comparison with state-of-the-art learning-based and learning-free retrievers (e.g., [1–2]).\n\nThank you for pointing us to these valuable references. Both [1] and [2] offer thorough analyses of retrieval strategies under in-distribution ICL settings, and we will make sure to cite them appropriately in the revised manuscript.\n\n- We position CCL differently along two key dimensions:\\\n- First, prior work typically assumes that the selection pool and the test query share the same distribution. While cross-task experiments are considered, they still assume distributional alignment between the test queries and the selection pool.\n- In contrast, CCL explicitly tackles the OOD setting, assuming a fixed demonstration pool across tasks and domains — which we consider to be a more practical and challenging scenario in many real-world situations. For example, datasets such as QNLI, PIQA, WSC273, and Yelp used in our evaluation do not appear in the demonstration pool.\n- Second, unlike prior approaches, CCL leverages domain-invariant features to guide retrieval. By explicitly modeling a causal factor $c$ that is intended to be invariant across domains, CCL aims to retrieve examples that are more likely to remain effective under OOD conditions.\n\n- We have also extended our evaluation to the ID setting. Table 3 reports the few-shot accuracy on SST-5 following the setup of [1]. CCL achieves competitive performance and outperforms other methods in the 5-shot and 10-shot settings.\n\nTable 3. Accuarcy performance comparision\n| SST-5 | 5-shot    | 10-shot   | 20-shot   |\n| ----  | -----     | -----     | -----     |\n| BERT  | 0.318     | 0.344     | 0.327     |\n| MLSM  | 0.322     | 0.349     | 0.332     |\n| EPR   | 0.383     | 0.390     | 0.369     |\n| TTF   | 0.385     | 0.390     | **0.427** |\n| CCL   | **0.397** | **0.422** | 0.416     |\n| | | |      |\n\n> Need for Deeper Explanation. The method appears to leverage output labels (y) from the demonstration set in a way that resembles TTF [1].\n\n- Thank you for this valuable observation. TTF uses a language model as a retriever often task-specific prediction modules for each task. TTF is based on the assumption that exemplars with similar labels to the predicted output of the test query are more useful for in-context learning.\n\n- In contrast, CCL uses labels only during training of the VAE — specifically, to help model the distribution of the output $y$ as a text embedding. This module is shared across tasks and does not require task-specific adaptation, allowing CCL to generalize to unseen tasks without additional retriever training.\n\n- While the latent representation used in TTF’s retriever output might superficially resemble our $c$ (since both are used to predict $y$), the underlying design philosophy is different. CCL’s $c$ is rooted in a causal framework: it is used not only for predicting $y$ but also for reconstructing $x$ as part of a principled data-generating process. This dual role is critical in ensuring that $c$ captures causal, invariant properties across domains.\n\n> The proposed method appears to have limited generalizability to tasks beyond those evaluated\n\n- We appreciate this comment. To evaluate generalization more rigorously, we conducted an additional experiment on an unseen generation task. Specifically, we tested the ability of CCL to support sentiment reversal paraphrasing, where the task is to transform a negative sentence into one with opposite (positive) sentiment.\n\n- CCL was trained only on binary and multiple-choice classification tasks. For the generation task, we used a system prompt that instructs the model to “rephrase the sentence to express an opposite sentiment” — without explicitly instructing it to convert negative sentences into positive ones.\n\nTable 4. Percentage of sentences that changed from negative to positive sentiment\n| phi4-mini-it    | Amazon | Yelp   |\n| ----            | -----  | -----  |\n| ICL (5-shot)     | 83.28  | 80.85  |\n| CCL (5-shot)     | **89.59**  | **84.77**  |\n| | | |\n\n- As shown in Table 4, CCL achieves better generalization on this unseen generation task, indicating its broader applicability beyond classification settings.\n\n> Writing and Clarity.\n\n- Thank you for your valuable suggestion to improve the clarity of the paper. We fully agree that including an algorithm will enhance readability, and we will incorporate it in the revised version accordingly.\n\n[3] Liu et al., Learning Causal Semantic Representation for Out-of-Distribution Prediction, NeurIPS 2021.\\\n[4] Lu et al., Invariant Causal Representation Learning for Out-of-Distribution Generalization, ICLR 2021."}, {"review_id": "a5l02UslUX", "reviewer": "Reviewer_Vb3H", "summary": "This paper addresses the brittleness of ICL under distribution shift by learning causal representations of examples. The proposed CCL framework uses a two-branch VAE to disentangle latent causal factors from environment factors, then retrieves demonstrations whose inferred causal factors are closest to those of the test query. Under a linear-causal SCM, the authors prove that selecting by causal similarity yields parameter updates closer to the true causal predictor and strictly lower test error than input-based retrieval.", "questions": "1. How does retrieval and ICL accuracy change when $\\dim(c)$ is 32 or 128, or when ELBO reconstruction loss varies?  \n2. Could you provide a ballpark estimate of the computational resources (GPU hours and peak RAM requirements) for VAE training on, say, 100 K examples?  \n3. Have you applied CCL retrieval to open-ended or chain-of-thought prompts, and what modifications would be needed?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n- The paper introduces an approach of integrating causal representation learning into ICL, focusing on robustness to domain shifts and providing a clear motivation for why causal factors should be prioritized over surface-level similarity.  \n- Under a well-specified linear-causal SCM, the authors offer rigorous theoretical analysis: Theorem 3.3 demonstrates that nearest-neighbor on input can fail arbitrarily badly, while Theorem 3.4 guarantees that causal-based retrieval accelerates convergence to the true predictor and reduces test error.  \n- Empirically, CCL shows strong performance across diverse settings: it achieves 100% causal match recovery on synthetic OOD splits and improvements on four NLP OOD benchmarks. Moreover, the method is stable to embedding choice, with minimal variation between text-embedding-3-small and multilingual-e5-large.  \n- The authors provide extensive reproducibility details in the appendix, including full ELBO derivations, proof sketches, synthetic data generation protocols, MGSM split rationale, and exact prompt templates.\n\n---\n\n## Weaknesses\n- Reported gains lack any statistical significance testing or confidence intervals, raising concerns about whether smaller improvements are robust or within experimental noise.  \n- There is no ablation study on the latent representation, such as varying $\\dim(c)$ or analyzing the impact of reconstruction quality on retrieval and ICL performance, leaving hyperparameter sensitivity uncharacterized.  \n- All experiments focus on classification or short-answer tasks, so it remains unknown how CCL would extend to free-form generation or chain-of-thought prompting scenarios.", "was_revised": true, "final_justification": "The rebuttal addressed my earlier concerns by adding robustness checks, hyperparameter sensitivity analysis, and preliminary evidence for generalization beyond classification tasks.\n\n### Resolved issues\n- Added multi-seed experiments showing stable performance gains across OOD benchmarks, supporting statistical robustness.\n- Provided an ablation on latent dimension size, showing that CCL’s performance is not overly sensitive to this parameter.\n- Presented an additional experiment on a paraphrasing generation task, demonstrating applicability beyond classification.\n\n### Remaining considerations\n- The evaluation still focuses primarily on short-answer tasks; while the generation example is promising, broader testing on complex reasoning or long-form generation would further strengthen generality claims.\n\nOverall, the paper makes a well-supported contribution by integrating causal representation learning into ICL for improved robustness under distribution shift. The rebuttal strengthened both the empirical and theoretical case, and I consider this a strong and relevant contribution to the field.", "rebuttal": "We truly appreciate your suggestions, which have provided valuable directions for strengthening the validity and credibility of our method. We have carefully addressed your comments in the revised version by incorporating the following additional analyses and experiments.\n\n> Reported gains lack any statistical significance testing or confidence intervals, raising concerns about whether smaller improvements are robust or within experimental noise.\n\n- We completely agree that verifying the robustness and reproducibility of our method through repeated trials is critical. In response to your comment, we conducted additional experiments using five different random seeds (\\{40, 41, 42, 43, 44\\}) to assess the stability of our VAE-based CCL model. The results are presented in Table 1, which reports the mean and standard deviation of CCL’s performance across OOD benchmarks.\n- Note that the LLM was set to perform classification deterministically using greedy decoding to eliminate randomness from the output side. Therefore, other baseline results are not affected by seed variation.\n\nTable 1. Accuracy of OOD benchmarks for Llama 3.2 3B IT\n|        | QNLI           | PIQA            | WSC273         | YELP            |\n| ----   | ---            | ----            | ------         | ----            |\n|  ZS    | 43.4 ($\\pm$ 0.00) | 71.3  ($\\pm$ 0.00) | 55.3 ($\\pm$ 0.00) | 89.0  ($\\pm$ 0.00) |\n|  LLM_R | 29.9 ($\\pm$ 0.00) | 69.9  ($\\pm$ 0.00) | 61.2 ($\\pm$ 0.00) | 79.5  ($\\pm$ 0.00) |\n|  ICL   | 68.1 ($\\pm$ 0.00) | 69.0  ($\\pm$ 0.00) | 49.8 ($\\pm$ 0.00) | 75.8  ($\\pm$ 0.00) |\n| CCL    | **75.2** ($\\pm$ 0.45) | **72.4**  ($\\pm$ 1.12) |**58.98** ($\\pm$ 2.78) | **95.10** ($\\pm$ 0.25) |\n| | | | | |\n\n- These results demonstrate that CCL achieves consistent improvements across different runs, indicating that its performance is stable and reproducible.\n\n> There is no ablation study on the latent representation, such as varying $dim(c)$ or analyzing the impact of reconstruction quality on retrieval and ICL performance, leaving hyperparameter sensitivity uncharacterized.\n\n- Thank you for highlighting this important point. We conducted an ablation study to investigate how sensitive CCL’s performance is to the dimension of the latent causal variable $c$. As shown in Table 2, even with a reduced latent dimension, CCL maintains strong performance across OOD benchmarks.\n- This suggests that CCL is not overly sensitive to the choice of $dim(c)$. Identifying an optimal dimension that ensures generalizable performance across domains — while minimizing memory footprint — is a promising direction for future research.\n\nTable 2. Accuracy on OOD benchmarks with varying $dim(c)$\n| dim( c )      | QNLI | PIQA  | WSC273 | YELP  | Avg. |\n| ----          | ---  | ----  | -----  | ----  | ---  |\n|  128          | 69.5 | **72.6**  | 56.1   | 93.5  | 72.9 |\n|  256          | **75.3** | 71.6  | 60.4   | 94.9  | 75.6 |\n|  1024 | 75.2 | 70.5  | **61.9**   | **95.4**  | **75.7** |\n| | | | | | |\n\n> All experiments focus on classification or short-answer tasks, so it remains unknown how CCL would extend to free-form generation or chain-of-thought prompting scenarios.\n\n- We appreciate this insightful comment. To explore the generalizability of CCL beyond classification, we conducted a new experiment on an unseen paraphrasing task, where the goal is to rewrite a negative sentence into one with an opposite (positive) sentiment.\n- While CCL was trained only on binary and multiple-choice classification tasks, it was applied here to a generation task. The system prompt instructed the model to “rephrase the sentence to express an opposite sentiment” — without explicitly directing it to convert negative sentences into positive ones.\n\nTable 3. Percentage of sentences that changed from negative to positive sentiment\n| phi4-mini-it    | Amazon | Yelp   |\n| ----            | -----  | -----  |\n| ICL (5-shot)     | 83.28  | 80.85  |\n| CCL (5-shot)     | **89.59**  | **84.77**  |\n| | | |\n\n- As shown in Table 3, CCL demonstrates superior performance compared to ICL, even on an unseen generation task. This suggests that CCL may hold promise for broader application beyond classification, potentially including text generation.\n\nWe sincerely thank you for your thoughtful and constructive feedback."}, {"review_id": "QTxE7HVifC", "reviewer": "Reviewer_2GL1", "summary": "This paper introduces a new method called Causal-aware In-context Learning (CCL) to improve the performance of large language models (LLMs) on out-of-distribution (OOD) datasets, where traditional in-context learning (ICL) methods fall short. The proposed CCL method is inspired by causal representation learning and aims to select better demonstration sets for LLMs by focusing on the underlying causal relationships in the data rather than superficial features. To achieve this, CCL uses a novel VAE-based causal representation learning technique to identify and capture these causal features. The paper provides both theoretical and empirical evidence that CCL improves the OOD generalization performance of LLMs by selecting demonstration sets that are causally related to the target query. By doing so, CCL helps LLMs to generalize better even in unseen environments.", "questions": "N/A", "limitations": "Yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n1. The methodology is designed to prevent the LLM from relying on misleading \"shortcuts\" or spurious correlations in the data. By forcing the model to consider the causal structure, CCL encourages a deeper understanding of the task, which in turn leads to better performance.\n\nWeakness:\n1. The process of learning causal representations with a VAE and then using this model to select demonstrations for every new query could be computationally expensive. The paper may not adequately address the practical scalability of this approach to very large models or datasets, which could be a significant barrier to real-world application.\n2. The method assumes that causal and non-causal features can be cleanly separated into distinct latent variables. In practice, especially in complex domains like natural language, these factors might be heavily entangled, making a clean separation difficult and potentially limiting the effectiveness of the approach.\n3. The theoretical model might oversimplify the nature of causality in language. It assumes a relatively straightforward mapping from latent variables to observed data, which may not capture the intricate, context-dependent, and multi-layered causal relationships inherent in language and reasoning.", "was_revised": true, "final_justification": "I appreciate the analyses, but my concern remains: a single‐layer VAE—even with dual reconstruction—still oversimplifies language’s deeply entangled, context-dependent causal factors. Without broader benchmarks or formal guarantees, it’s unclear whether CCL truly learns causal structure rather than surface correlations.\n\nMy score keeps unchanged.", "rebuttal": "We sincerely thank you for the constructive and forward-looking feedback, which has been very helpful in guiding improvements to our work. I hope the responses below sufficiently address your concerns and questions. If anything remains unclear or requires further clarification, we would be more than happy to provide additional explanation.\n\n> The process of learning causal representations with a VAE and then using this model to select demonstrations for every new query could be computationally expensive. The paper may not adequately address the practical scalability of this approach to very large models or datasets, which could be a significant barrier to real-world application.\n\n- We appreciate this insightful comment. Many retrieval-based ICL methods — such as LLM-R [1], TTF [2] — train a retriever to find optimal demonstrations for a given input. These approaches typically rely on open-source language models as retrievers, requiring access to the full or partial model weights during training and inference.\n\n- In contrast, CCL only requires a relatively lightweight VAE for training, and its latent variables ($c$, $s$) are significantly more compact than raw text embeddings. The parameter size of our VAE is 6.3M, which is approximately 0.2% of the smallest language model we evaluated (LLaMA 3.2 3B Instruct).\n\n- To quantify the computational overhead, we compared CCL with embedding-based ICL without any retriever training. The average time to extract embeddings using OpenAI’s text-embedding-3-small API is about 0.416s per sample, while the multilingual-e5-large-instruct open-source model takes about 0.032s.\n\n- For CCL:\n  - VAE training takes approximately 20.4 minutes across 5 random seeds.\n  - A single training iteration for one sample takes 0.06s.\n  - Inference time for generating $c$ and $s$ embeddings is 0.003s per sample.\n  - Construction of an 8-shot prompt via k-means clustering takes about 0.2s, same as ICL.\n  - Inference using LLaMA 3.2B IT on a single A6000 GPU shows average response times of:\n     - CCL: 0.10s\n     - ICL: 0.16s\n\n- Thus, the total time per sample is 0.78s for CCL vs. 0.63s for ICL — a marginal increase. In our view, when the VAE is pre-trained and $c$ embeddings are pre-computed for the example pool, the inference-time overhead could be sufficiently small to make CCL practically usable in real-world scenarios.\n\n> The method assumes that causal and non-causal features can be cleanly separated into distinct latent variables. In practice, especially in complex domains like natural language, these factors might be heavily entangled, making a clean separation difficult and potentially limiting the effectiveness of the approach.\n\n- We agree that in natural language, the disentanglement of causal and non-causal factors can be inherently challenging due to semantic and contextual entanglement. To explore this concern, we performed a qualitative analysis of the learned latent features to better understand how $c$ and $s$ are interpreted in practice.\n\n- As mentioned in the paper, our objective is to decompose each sentence into components relevant to task objectives (domain-invariant) and those that are not (domain-variant). For instance, in sentiment analysis:\n  - Domain-invariant features capture polarity (positive/negative sentiment).\n  - Domain-variant features depend on dataset characteristics — e.g., service or food for Yelp, product quality or usability for Amazon.\n\n- To visualize the meaning encoded in $c$ and $s$, we decoded sentence embeddings with one latent dimension zeroed out. Specifically:\n  1. Infer $c$ and $s$ from input embedding $x$.\n  2. Set $s=0$ to generate $x_{s=0}'$ (highlighting $c$).\n  3. Set $c=0$ to generate $x_{c=0}'$ (highlighting $s$).\n  4. Find nearest words to each decoded embedding to interpret feature semantics.\n\nQualitative example (Yelp): “the red velvet pancakes were horrible and brown, and potatos were over cooked and bland.. would not recommend”\n\nTable 1. Top 5 nearest words\n\n| $x$                | $x_{s=0}'$         | $x_{c=0}'$ |\n| --------           | --------           | --------   |\n| horribleappetizers | unappetizing       | review     |\n| pancakes           | flavorless         | reviewers  |\n| potatos            | horribleappetizers | critiques  |\n| hadhorrible        | inedible           | soggy      |\n| bad                | trashed            | reviews    |\n| | | |\n\n- The original embedding ($x$) captures both semantic and contextual tokens (e.g., “pancakes”, “potatos”). In contrast, $x_{s=0}'$ clusters strongly around negative sentiment expressions, while $x_{c=0}'$ clusters around meta-content like “review”, which aligns with dataset-specific artifacts.\n\nQualitative Example 2 (Amazon dataset):\nOriginal negative sentence: “Worked for about 4 months. DVD player will not eject or accept disks. Do not buy.”\n\nTable 2. Top 5 nearest words\n| $x$       | $x_{s=0}'$  | $x_{c=0}'$ |\n| --------  | --------    | -------- |\n| dvd       | unusable    | reverb   |\n| eject     | expired     | throw    |\n| disks     | cancelled   | film     |\n| unusable  | crappy      | review   |\n| purchased | trashed     | trip     |\n| |  |  |\n\n- Similarly, in Table 2, $x_{s=0}'$ captures negative sentiment towards product usability, while $x_{c=0}'$ retrieves terms more loosely associated with contextual metadata.\n\n- These observations suggest that CCL learns interpretable and partially disentangled representations. While a perfect separation may not always be attainable in natural language, the model demonstrates meaningful structuring of causal factors.\n\n> The theoretical model might oversimplify the nature of causality in language. It assumes a relatively straightforward mapping from latent variables to observed data, which may not capture the intricate, context-dependent, and multi-layered causal relationships inherent in language and reasoning.\n\nThank you for raising this important point. We agree that causal reasoning in natural language is inherently complex and often multi-layered. However, recent advancements in **text embedding models** — many of which are widely used in real-world applications like retrieval-augmented generation (RAG) — suggest that **rich semantic and structural information** can be captured within a **single embedding**. This strengthens the feasibility of learning informative latent variables like $c$ from natural language inputs.\n\n- In our model, the **latent causal variable $c$ is trained not only to help reconstruct the input $x$, but also the output $y$**. This dual reconstruction objective encourages $c$ to **encode information necessary for solving the downstream task**, thereby aligning the learned representation with the task-specific causal features.\n\n- As discussed earlier and demonstrated in Table 1 and Table 2, the VAE is capable of meaningfully disentangling $c$ (domain-invariant) from $s$ (domain-variant). This supports our claim that the simplified causal structure adopted by CCL can still yield interpretable and functionally distinct representations that are aligned with the our goals for ICL in OOD setting.\n\n- Furthermore, we observed a **correlation between the reconstruction quality of the output $y$ and the accuracy on NLP tasks** — suggesting that **improving the expressiveness of $c$ is a potential path to handling tasks that require reasoning about complex causal relationships**.\n\n- Nevertheless, we acknowledge that this **single-layer generative structure may encounter limitations in handling the full complexity of certain reasoning tasks**. In particular, tasks that involve hierarchical or multi-hop reasoning may require richer representations of causality than what our current VAE-based framework offers.\n\n- To address this, we are actively exploring how to **extend CCL beyond the current framework**. Specifically, we are investigating leveraging transformer-based language models, rather than VAEs, to infer latent variables in tasks that require deeper reasoning.\n\n- Transformer architectures are known to implicitly capture hierarchical reasoning structures, and prior work [3] shows that **deeper transformers can handle complex tasks even without explicit chain-of-thought prompting**.\n\n- Moreover, **circuit-based interpretability studies [4] suggest that it is possible to identify which hidden states within transformers contribute most significantly to reasoning processes**. Inspired by these findings, we are exploring approaches to identify and utilize **specific hidden states that encode task-relevant information**, with the goal of **inferring latent causal features such as $c$** more effectively.\n\nThis direction represents a potential avenue for enabling CCL to handle more layered causal structures in natural language.\n\n[1] Wang et al., Learning to Retrieve In-Context Examples for Large Language Models, EACL 2024\\\n[2] Liu et al., Unraveling the Mechanics of Learning-Based Demonstration Selection for In-Context Learning, ACL 2025\\\n[3] Li et al., Chain of Thought Empowers Transformers to Solve Inherently Serial Problems, ICLR 2024\\\n[4] Dunefsky et al., Transcoders Find Interpretable LLM Feature Circuits, NeurIPS 2024"}, {"review_id": "ynYTH2ej6G", "reviewer": "Reviewer_1SuJ", "summary": "This paper proposes Causal-aware In-Context Learning (CCL) to address the performance limitations of traditional in-context learning (ICL) under out-of-distribution (OOD) settings. The authors argue that conventional ICL methods, which rely on superficial similarity in input xxx, often fail when demonstration samples come from different semantic domains despite having structurally similar forms. To remedy this, the paper leverages causal representation learning, using a VAE-based framework to extract task-relevant latent variables c that remain invariant across environments, while separating out spurious factors s. CCL selects demonstrations based on similarity in ccc, rather than x, leading to improved OOD generalization. Both theoretical analysis and empirical results confirm that demonstrations can diverge from the surface form of x as long as the semantic (causal) alignment is preserved.", "questions": "1.\tCan provided an example of x, c and s? In particular, is the modeling of s strictly necessary for the success of CCL, or could a simplified formulation using only c suffice for identifying useful demonstrations? \n\n2.\tHow you select x for each test query for generating c and s? Is there a separate VAE trained per task/domain? How is the demonstration pool determined for a new query during test time—does it rely on fixed candidate sets from the source distribution?\n\n3.\tIn Table 4, CCL underperforms on some tasks. Have the authors conducted any error analysis to investigate these cases? Is there evidence that these failures stem from incorrect estimation of the causal variables c? \n\n4.\tMore broadly, could the authors discuss the identifiability of the causal representations under their VAE framework and whether the observed performance drop is due to limitations in causal inference?", "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": "Strength\n\n•1.  Novelty & Significance: The paper introduces a principled and theoretically grounded method for demonstration selection in ICL by emphasizing causal similarity. This represents a meaningful step forward for robust ICL under distribution shift.\n\n•2.  Theoretical Depth: The method is supported by a rigorous theoretical framework. The authors clearly define their assumptions and provide formal guarantees demonstrating the benefits of selecting demonstrations based on causal representations. This theoretical grounding makes the proposed approach compelling and credible.•\n\n3.  Comprehensive Experiments: Experiments on both synthetic and real-world datasets demonstrate the effectiveness of the method. The design clearly isolates and evaluates the role of causal representations, validating the method's robustness in OOD scenarios.\n\nWeakness:\nSee Questions", "was_revised": true, "final_justification": "The clarification resolved my questions. I keep my positive review for this paper.", "rebuttal": "Thank you very much for your detailed and insightful review. Your feedback has been instrumental in helping us refine and improve the clarity and robustness of our work. We have carefully considered each of your comments. We hope that our responses below effectively clarify the concerns, and we remain open to providing further clarification where needed.\n\n> Can provided an example of x, c and s?\n\nThank you for this important question. We illustrate how these latent variables behave in practice using a qualitative decoding analysis from our trained VAE under the CCL framework.\n\nTo assess the interpretability of the latent variables, we first infer $c$ and $s$ from a given input sentence embedding $x$, and then decode variants of the input by zeroing out either $c$ or $s$\n- $x_{s=0}'$ retains only the information encoded in $c$ (task-relevant).\n- $x_{c=0}'$ retains only the information encoded in $s$ (domain-specific).\n\nBecause $x_{s=0}'$ and $x_{c=0}'$ lie in the same embedding space as $x$, we retrieved nearest neighbors to reveal the types of semantic content each latent variable preserves.\n\nQualitative Example 1 (Yelp dataset):\nOriginal negative sentence: “the red velvet pancakes were horrible and brown, and potatos were over cooked and bland.. would not recommend”\n\nTable 1. Top 5 nearest words\n\n| $x$                | $x_{s=0}'$         | $x_{c=0}'$ |\n| --------           | --------           | --------   |\n| horribleappetizers | unappetizing       | review     |\n| pancakes           | flavorless         | reviewers  |\n| potatos            | horribleappetizers | critiques  |\n| hadhorrible        | inedible           | soggy      |\n| bad                | trashed            | reviews    |\n\nIn Table 1,\n- $x$ retrieves both semantic and domain-specific tokens (e.g., “pancakes”, “potatos”).\n- $x_{s=0}'$ clusters around sentiment-related tokens, e.g., “unappetizing”, “inedible”.\n- $x_{c=0}'$ surfaces meta-review terms such as “review”, “critique”, which likely reflect domain-specific features.\n\nQualitative Example 2 (Amazon dataset):\nOriginal negative sentence: “Worked for about 4 months. DVD player will not eject or accept disks. Do not buy.”\n\nTable 2. Top 5 nearest words\n| $x$       | $x_{s=0}'$  | $x_{c=0}'$ |\n| -  | -   | - |\n| dvd       | unusable    | reverb   |\n| eject     | expired     | throw    |\n| disks     | cancelled   | film     |\n| unusable  | crappy      | review   |\n| purchased | trashed     | trip     |\n\nSimilarly, in Table 2,\n- $x_{s=0}’$ captures failure-related sentiment (e.g., “unusable”, “trashed”).\n- $x_{c=0}’$ retrieves words like “film” and “review”, showing latent connection to domain/style.\n\nThese examples highlight that the CCL is able to effectively separate task-relevant and task-irrelevant factors across domains.\n\n> In particular, is the modeling of s strictly necessary for the success of CCL, or could a simplified formulation using only c suffice for identifying useful demonstrations?\n\n- The reason we explicitly model two latent variables, $c$ and $s$, is to isolate domain-invariant features ($c$) for the purpose of robust retrieval particularly in OOD settings.\n- If we were to use a single latent variable ($z$) without decomposing it into $c$ and $s$, then $z$ would have to represent both the invariant and variant factors embedded in $x$. This conflation leads to two key problems:\n  - (1) Entangled representation: without disentanglement, the latent variable may capture spurious correlations tied to specific domains or datasets. Consequently, retrieval based on $z$ may select examples that are similar on superficial terms (e.g., domain, writing style), rather than the underlying invariant feature.\n  - (2) Inconsistent behavior across domains: Since s is highly sensitive to environmental factors (e.g., dataset source), a $z$ that mixes $c$ and $s$ will behave inconsistently under distribution shifts, undermining the core motivation of CCL — OOD generalization.\n- This point is reinforced both theoretically (Section 3.4) and empirically. Theorems 3.3 and 3.4 show that selecting demonstrations based on $x$ (or entangled $z$) can lead to large prediction errors due to confounding from $s$, whereas focusing on $c$ leads to improved convergence and lower test error.\n- In synthetic experiments (Table 1 in the paper), retrieval based on $c$ clearly outperforms that based on $x$ or $s$ when the goal is to retrieve task-relevant demonstrations, especially in OOD scenarios. This confirms that disentangling is not only theoretically desirable but necessary in practice to ensure that $c$ learns the invariant features.\n\n> How you select x for each test query for generating c and s? Is there a separate VAE trained per task/domain?\n\nWe appreciate the opportunity to clarify this point. The VAE used in CCL is trained on a collection of benchmark datasets spanning a wide range of tasks and domains. In other words, a single shared VAE is trained using data from diverse tasks and domains, enabling it to generalize and infer latent representations ($c$, $s$) across heterogeneous inputs.\n\nAt test time, for each test query $x_t$, we compute its embedding and pass it through the pretrained VAE to infer its corresponding latent causal variable $c_t$ and spurious variable $s_t$. No task- or domain-specific retraining of the VAE is required.\n\n> How is the demonstration pool determined for a new query during test time—does it rely on fixed candidate sets from the source distribution?\n\nCCL is specifically designed to be robust in out-of-distribution (OOD) settings. During inference, we assume access to a fixed and task-agnostic demonstration pool, which does not include examples from the same distribution as the test query.\n\nGiven a new test query $x_t$, we first extract its sentence embedding and use the pretrained VAE to infer $c_t$. Then, using similarity in the $c$-space, we retrieve the nearest examples from the fixed pool to construct a demonstration prompt. This retrieval process is entirely independent of task-specific labels or distributions.\n\n> In Table 4, CCL underperforms on some tasks. Have the authors conducted any error analysis to investigate these cases?\n\nThank you for raising this important point. We conducted an error analysis to explore the correlation between model performance and the quality of $c$ inferred by the VAE.\n\nSpecifically, we measured the average reconstruction loss of both input $x$ and output $y$ on in-distribution (ID) and out-of-distribution (OOD) datasets across four task types:\n- Coreference resolution: ID = WSC, OOD = WSC273\n- Natural language inference: ID = RTE, OOD = QNLI\n- Sentiment analysis: ID = Sentiment140, OOD = Yelp\n- Commonsense reasoning: ID = COPA, OOD = PIQA\n\nFor each dataset, we calculated the reconstruction loss on the test set using the pretrained VAE.\n\nInterestingly, we found that the commonsense reasoning tasks (COPA/PIQA) exhibited the highest reconstruction loss, particularly on the output $y$. This suggests that the VAE had greater difficulty modeling causal representations in this domain, potentially due to the more abstract or diverse semantics involved.\n\nThese findings imply that weaker performance on certain tasks (e.g., PIQA) may indeed stem from less accurate inference of $c$, highlighting a valuable direction for improving the robustness of causal variable estimation.\n\nImportantly, we appreciate the reviewer’s insight, which has led us to a new line of investigation. In our current setup, we balanced the training dataset sizes across domains. However, we now realize that equal data volume may not translate to equal modeling difficulty across tasks.\n\n> More broadly, could the authors discuss the identifiability of the causal representations under their VAE framework and whether the observed performance drop is due to limitations in causal inference?\n\nThank you for raising this important question. We would like to address your concern from both a conceptual and empirical perspective.\n\n- Our approach is inspired by works such as CEVAE [1], where a VAE is used to estimate the effect of unobserved confounders on treatment and outcome variables. Importantly, in such settings, the goal is not to explicitly recover the exact identity of the latent confounder, but rather to learn a latent variable that captures its functional influence on treatment and outcome.\n- Similarly, our framework does not aim to provide full identifiability of the latent causal variable $c$, but instead learns a functional approximation that is sufficient for improved generalization under distribution shift.\n- In line with this, our VAE-based model learns two latent variables — $c$ (domain-invariant, task-relevant) and $s$ (domain-variant) — whose separation is functionally meaningful. This is supported empirically:\n  - In synthetic experiments (Table 1 in the manuscript), we show that $c$ and $s$ align with ground-truth task and environment variables.\n  - In real data (Tables 1 and 2), decoding from $c$ and $s$ yields clearly different semantics, showing the model effectively separates domain-invariant and domain-variant information.\n- However, as mentioned earlier, we believe that the performance drop in certain tasks is less about the disentanglement between $c$ and $s$, and more about the expressiveness or optimization of $c$ itself in predicting $y$.\n- A promising direction is to enhance the VAE’s representation capacity. The current VAE in CCL has only 6.3M parameters (~0.2% the size of LLaMA 3.2 3B). A larger or more expressive VAE could lead to better modeling of $c$.\n- Your review motivated us to explore the correlation between the reconstruction loss of $y$ (from $c$) and NLP task performance. This suggests a new direction for improving the robustness of CCL by better optimizing the latent causal representation $c$. We will include this analysis and discussion in the revision.\n\n[1] Louizos et al., Causal Effect Inference with Deep Latent-Variable Models, NeurIPS 2017"}]

Author rebuttal:

Meta-review: This paper introduces Causal-aware In-context Learning (CCL), which is a principled framework that integrates causal representation learning into in-context learning to improve out-of-distribution generalization. The work is theoretically well-grounded (Theorems 3.3 and 3.4 show the benefits of causal retrieval over surface similarity) and supported by experiments across synthetic, multilingual math reasoning, and diverse NLP benchmarks. The reviewers initially raised concerns regarding interpretability of causal factors, baseline coverage, and robustness. During rebuttal, the authors provided useful information including qualitative analyses and tests on more challenging benchmarks such as MMLU and generation tasks. I feel that these clarifications and new results resolved most issues. I would recommend acceptance.
Decision: Accept (poster)