CCL: Causal-aware In-context Learning for Out-of-Distribution Generalization

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.

Can provided an example of x, c and s?

Thank 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.

To 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$

• $x_{s=0}'$ retains only the information encoded in $c$ (task-relevant).
• $x_{c=0}'$ retains only the information encoded in $s$ (domain-specific).

Because $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.

Qualitative Example 1 (Yelp dataset): Original negative sentence: “the red velvet pancakes were horrible and brown, and potatos were over cooked and bland.. would not recommend”

Table 1. Top 5 nearest words

In Table 1,

• $x$ retrieves both semantic and domain-specific tokens (e.g., “pancakes”, “potatos”).
• $x_{s=0}'$ clusters around sentiment-related tokens, e.g., “unappetizing”, “inedible”.
• $x_{c=0}'$ surfaces meta-review terms such as “review”, “critique”, which likely reflect domain-specific features.

Qualitative Example 2 (Amazon dataset): Original negative sentence: “Worked for about 4 months. DVD player will not eject or accept disks. Do not buy.”

Table 2. Top 5 nearest words

Similarly, in Table 2,

• $x_{s=0}’$ captures failure-related sentiment (e.g., “unusable”, “trashed”).
• $x_{c=0}’$ retrieves words like “film” and “review”, showing latent connection to domain/style.

These examples highlight that the CCL is able to effectively separate task-relevant and task-irrelevant factors across domains.

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?

• 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.
• 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:
  • (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.
  • (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.
• 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.
• 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.

How you select x for each test query for generating c and s? Is there a separate VAE trained per task/domain?

We 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.

At 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.

How is the demonstration pool determined for a new query during test time—does it rely on fixed candidate sets from the source distribution?

CCL 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.

Given 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.

In Table 4, CCL underperforms on some tasks. Have the authors conducted any error analysis to investigate these cases?

Thank 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.

Specifically, 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:

• Coreference resolution: ID = WSC, OOD = WSC273
• Natural language inference: ID = RTE, OOD = QNLI
• Sentiment analysis: ID = Sentiment140, OOD = Yelp
• Commonsense reasoning: ID = COPA, OOD = PIQA

For each dataset, we calculated the reconstruction loss on the test set using the pretrained VAE.

Interestingly, 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.

These 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.

Importantly, 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.

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?

Thank you for raising this important question. We would like to address your concern from both a conceptual and empirical perspective.

• 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.
• 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.
• 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:
  • In synthetic experiments (Table 1 in the manuscript), we show that $c$ and $s$ align with ground-truth task and environment variables.
  • 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.
• 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$.
• 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$.
• 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.

[1] Louizos et al., Causal Effect Inference with Deep Latent-Variable Models, NeurIPS 2017

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.

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.

• 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.

• 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).

• 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.

• For CCL:

  • VAE training takes approximately 20.4 minutes across 5 random seeds.
  • A single training iteration for one sample takes 0.06s.
  • Inference time for generating $c$ and $s$ embeddings is 0.003s per sample.
  • Construction of an 8-shot prompt via k-means clustering takes about 0.2s, same as ICL.
  • Inference using LLaMA 3.2B IT on a single A6000 GPU shows average response times of:
    • CCL: 0.10s
    • ICL: 0.16s

• 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.

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.

• 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.

• 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:

  • Domain-invariant features capture polarity (positive/negative sentiment).
  • Domain-variant features depend on dataset characteristics — e.g., service or food for Yelp, product quality or usability for Amazon.

• To visualize the meaning encoded in $c$ and $s$, we decoded sentence embeddings with one latent dimension zeroed out. Specifically:

  1. Infer $c$ and $s$ from input embedding $x$.
  2. Set $s=0$ to generate $x_{s=0}'$ (highlighting $c$).
  3. Set $c=0$ to generate $x_{c=0}'$ (highlighting $s$).
  4. Find nearest words to each decoded embedding to interpret feature semantics.

Qualitative example (Yelp): “the red velvet pancakes were horrible and brown, and potatos were over cooked and bland.. would not recommend”

Table 1. Top 5 nearest words

• 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.

Qualitative Example 2 (Amazon dataset): Original negative sentence: “Worked for about 4 months. DVD player will not eject or accept disks. Do not buy.”

Table 2. Top 5 nearest words

• 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.

• 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.

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.

Thank 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.

• 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.

• 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.

• 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.

• 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.

• 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.

• 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.

• 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.

This direction represents a potential avenue for enabling CCL to handle more layered causal structures in natural language.

[1] Wang et al., Learning to Retrieve In-Context Examples for Large Language Models, EACL 2024
[2] Liu et al., Unraveling the Mechanics of Learning-Based Demonstration Selection for In-Context Learning, ACL 2025
[3] Li et al., Chain of Thought Empowers Transformers to Solve Inherently Serial Problems, ICLR 2024
[4] Dunefsky et al., Transcoders Find Interpretable LLM Feature Circuits, NeurIPS 2024

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.

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.

Reported gains lack any statistical significance testing or confidence intervals, raising concerns about whether smaller improvements are robust or within experimental noise.

• 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.
• 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.

Table 1. Accuracy of OOD benchmarks for Llama 3.2 3B IT

• These results demonstrate that CCL achieves consistent improvements across different runs, indicating that its performance is stable and reproducible.

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.

• 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.
• 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.

Table 2. Accuracy on OOD benchmarks with varying $dim(c)$

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.

• 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.
• 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.

Table 3. Percentage of sentences that changed from negative to positive sentiment

• 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.

We sincerely thank you for your thoughtful and constructive feedback.

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.

Interpretability Challenge. While the paper adopts causal learning, the proposed latent causal variables lack interpretability.

• 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].

• 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.

• 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.

• 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).

• 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.

• 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:

  • Zeroing out $s$ to generate $x’_{s=0}$, and
  • Zeroing out $c$ to generate $x'_{c=0}$.

• 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$.

Qualitative Example 1 (Yelp dataset): Original negative sentence: “the red velvet pancakes were horrible and brown, and potatos were over cooked and bland.. would not recommend”

Table 1. Top 5 nearest words

• 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).

Qualitative Example 2 (Amazon dataset): Original negative sentence: “Worked for about 4 months. DVD player will not eject or accept disks. Do not buy.”

Table 2. Top 5 nearest words

• 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.

• Our qualitative analysis indicates that CCL may learn to separate causal factors in natural language, even with a relatively simple VAE-based architecture.

Insufficient Experiments. The paper lacks comparison with state-of-the-art learning-based and learning-free retrievers (e.g., [1–2]).

Thank 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.

• We position CCL differently along two key dimensions:\

• 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.

• 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.

• 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.

• 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.

Table 3. Accuarcy performance comparision

Need for Deeper Explanation. The method appears to leverage output labels (y) from the demonstration set in a way that resembles TTF [1].

• 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.

• 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.

• 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.

The proposed method appears to have limited generalizability to tasks beyond those evaluated

• 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.

• 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.

Table 4. Percentage of sentences that changed from negative to positive sentiment

• As shown in Table 4, CCL achieves better generalization on this unseen generation task, indicating its broader applicability beyond classification settings.

Writing and Clarity.

• 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.

[3] Liu et al., Learning Causal Semantic Representation for Out-of-Distribution Prediction, NeurIPS 2021.
[4] Lu et al., Invariant Causal Representation Learning for Out-of-Distribution Generalization, ICLR 2021.

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.