Disentangling Latent Shifts of In-Context Learning Through Self-Training

Thank you for your thoughtful comments. We're glad you found our method straightforward, significant, and clearly presented. We appreciate your recognition of self-training with pseudo-labels as a way to address inference-time bottlenecks and improve performance.

There is no section on limitations. I would be interested in seeing some discussion on this with respect to both the chosen method and results. ...

Thank you for raising this point. Below are key limitations of STICL, which will be included in the revised manuscript.

1. Computational cost: STICL introduces extra overhead due to adapter fine-tuning during self-training. While lighter than full model fine-tuning, it is still costlier than standard ICL, which avoids weight updates. However, STICL offsets this by removing demonstrations from the input during inference. For instance, processing 16 demonstrations with Llama 3 (8B) takes ~120x longer than a 0-shot setup due to the self-attention bottleneck. Based on our measurements, self-training with 100 unlabeled instances and 16 demonstrations has a computational cost equivalent to ~2100 inferences in a 16-shot setup, after which the fine-tuning cost is balanced out.

2. Applicability: STICL relies on unlabeled data, making it less suitable for extremely resource-limited settings. In our experiments, 100 unlabeled instances with up to 32 demonstrations sufficed for strong performance. However, acquiring even modest amounts of unlabeled data could be challenging in some cases.

3. Large demonstration sets: Although STICL efficiently encodes demonstrations into adapters to overcome context length limitations, the method has not been extensively tested with very large demonstration sets. Our findings suggest that using multiple adapters with manageable demonstration sizes is more effective as the number of demonstrations increases. While this approach could theoretically scale further, its effectiveness with much larger sets remains untested, and additional adapters increase computational costs, creating a tradeoff between scalability and efficiency.

Some of the discussion on weak-to-strong generalisation in section 2 is unclear. It does not explain why the student is expected to outperform the teacher or about what local-to-global consistency is in the context of in-context learning.

Thank you for pointing this out. We will clarify these concepts in the revised manuscript.

The concept of weak-to-strong generalization in Section 2 assumes the student model leverages data structure and robustness to local perturbations to improve beyond the teacher's pseudo-labels. While the teacher relies on ICL with fixed demonstrations, the student benefits from fine-tuning, refining noisy pseudo-labels to align with true labels, and learning consistent representations. Local-to-global consistency enables the student to maintain stable outputs under small input changes (local consistency) and generalize these patterns to new or complex queries (global consistency). Prior work (e.g., Lang et al., 2024) establishes theoretical bounds on the student's performance, showing that its error is directly proportional to pseudo-label error and inversely proportional to robustness, which we approximate using the Lipschitz constant.

… is there a way you could experimentally determine whether the demonstrations have been disentangled?

Thank you for this question. This prompted us to validate further STICL's ability to disentangle shifts, where we designed an experiment to assess how faithfully encoded demonstrations are captured in the adapter's weights. Using 1000 examples per GLUE dataset with Llama 3 (8B), we encoded a single demonstration and assessed the student model's ability to reproduce it by calculating the BERTScore (F1) similarity between the student's output and the original demonstration. The averaged results are shown below:

The consistently high BERTScore values demonstrate that the student model faithfully retrieves the encoded demonstration from the adapter.

Questions

Q1

Yes, the queries used for training the student are drawn from the training set, with their labels removed to simulate an unlabeled dataset. For evaluation, queries are exclusively drawn from the test set.

Q2

To the best of our knowledge, stability in ICL is not limited to transformer-based models. State-space models may face numerical stability challenges and prompt drift, diffusion models risk compounded errors during iterative denoising, and RNNs often exhibit sensitivity to sequence length and order dependence.

Q3

Please refer to our earlier response on the limitations.

Thank you for your valuable comments and suggestions, especially for proposing ideas for expanding our method to other setups. We are glad that you find our paper clear and interesting. Thank you for recognizing the potential of our method to mitigate the demonstration ordering issues, reduce inference cost, and improve ID and OOD generalization.

It's more natural for me to take 16 samples from k 16 examples to construct 16-shot demonstrations for STICL-R, rather than construct k separate 16-shot sets for k adapters in STICL-S.

Thank you for raising this point. The motivation behind STICL-R is to provide a comparison with STICL-S, which we propose as the superior and more practical approach. STICL-R requires more labeled data, as it resamples demonstrations in each iteration, yet it is still outperformed by STICL-S. Our experiments aimed to validate the intuition that resampling can degrade performance, likely due to overloading the adapter with too many demonstrations over the course of multiple epochs. We actually tested the proposed approach using k different 16-shot sets with a single adapter. However, this approach yielded limited success. We found that dedicating a separate adapter to each set resulted in improved performance. The multiple adapters can be merged into a single adapter for inference, ensuring efficiency while leveraging the benefits of multiple adapters during training.

Questions

Q1

Pseudo-label correction is the process where the student model, in a teacher-student framework, refines noisy pseudo-labels generated by the teacher, aligning them closer to true labels. This correction links the student’s true error to its error on pseudo-labeled data, leveraging the structure of the data. It assumes two types of points:

• "Good" points with correct pseudo-labels.
• "Bad" points with incorrect pseudo-labels.

Bad points are assumed to be surrounded by good points in the feature space ("neighborhoods"). If the model is robust to small perturbations, it can use good points to correct bad ones.

A formal bound can be derived (Lang et al., 2024) to show that the true error of the student model depends on:

1. The rate of pseudo-label errors.
2. The robustness of the classifier in local neighborhoods.
3. The proportion of bad points that are close to good points.

If these conditions are met, pseudo-label correction improves performance beyond learning from noisy labels. We will revise the manuscript to clarify this phenomenon.

Q2

We approximate the Lipschitz constant by computing the Frobenius norm of the input-output Jacobian matrix of the network. In Appendix B, we lay out how this approximation connects the Jacobian matrix to the Lipschitz constant. Specifically, for a given method, we first complete the training phase. Afterward, we compute the Jacobian matrices for selected instances, using the embeddings as input variables and the penultimate layer activations as output variables. We will revise our manuscript to clarify this process.

Q3

Thank you for this interesting question! We additionally evaluated STICL-S with Llama 3 (8B) in a 16-shot setup with 16 additional demonstrations encoded in the adapter (32 labeled instances in total). We compared this configuration to the standard 32-shot setup and two 0-shot STICL-S variants using 16 and 32 labeled instances, respectively. In the table below, we denote the STICL variant using the format n/d, where n represents the number of shots (n-shot) and d indicates the number of demonstrations encoded in each adapter. The results (averaged over 10 runs) are as follows:

STICL-S in the 16-shot setup with 16 demonstrations encoded outperforms standard 32-shot ICL and STICL-S (0/16) across all datasets, leveraging additional context during inference. However, it performs slightly worse than 0-shot STICL-S variants with 32 labeled instances (0/32), likely due to the self-training process exclusive to 0-shot setups. Nonetheless, the model performs well in n-shot setups where n > 0.

Q4

Thank you for this valuable suggestion. We reflected on this idea and concluded that exploring joint training of multiple tasks and analyzing its impact on OOD generalization is indeed a promising direction, which we plan to explore in future work.

Thank you for your thoughtful comments and suggestions. We appreciate that you found the self-training analysis interesting and useful, and that you recognize the practical significance of our work.

One main weakness (and source of confusion for me) is the terminology of "in-context learning" when the models are being fine-tuned via self-training. … The method is much more expensive than n-shot, and in general, the computational aspects need to be considered for fair comparisons.

Thank you for raising this point. While STICL incorporates fine-tuning through self-training, its primary goal remains to enhance ICL. Crucially, ICL is an integral part of STICL, as the teacher model relies on ICL to generate pseudo-labels. In this framework, fine-tuning is used specifically to align the student model's outputs with the teacher's pseudo-labels. For this reason, we decided to keep "ICL" in the method's name, as the approach principally relies on ICL. Although one of the key advantages of ICL over fine-tuning is that it eliminates the need for weight updates, ICL has significant limitations. As we noted in the Introduction, ICL struggles with effectively leveraging information from long contexts and is constrained by the context window size, limiting its ability to incorporate multiple demonstrations. STICL addresses these drawbacks by combining ICL with lightweight fine-tuning of adapter modules, overcoming the context size limitation while preserving the foundational principles of ICL.

Regarding computational costs, although fine-tuning is generally more expensive than n-shot ICL, the fine-tuning of adapters in STICL is computationally lightweight. It can also offset some of its cost in practice by eliminating the need for repeated demonstrations during inference, effectively enabling a 0-shot setup. Prompted by your remark, we conducted additional experiments to compare the computational costs of n-shot ICL and STICL. For example, with Llama 3 (8B) processing 16 demonstrations from GLUE datasets, inference takes approximately 120 times longer than a 0-shot setup (processing only the query). Additionally, we calculated that in this setup, the computational cost of self-training with a single adapter is roughly equivalent to performing 2100 inferences in a 16-shot setup. This means that after about 2100 inferences, the time invested in fine-tuning is effectively offset by the efficiency gains during inference. We will include this analysis in the revised manuscript.

... I don't think this disentanglement introduces a fundamentally new perspective - in-context learning can be approximated by (the more expensive) finetuning on related task data. …

We understand your suggestion that ICL can be approximated as a form of fine-tuning, a perspective we agree with. However, our contribution is not focused on introducing a new framing of ICL. Instead, we view our work as a practical exploration of disentangling latent shifts to enhance and simplify the ICL process, drawing inspiration from similar approaches in the literature (Hendel et al., 2023; Liu et al., 2023). One concrete benefit of this disentanglement is the elimination of demonstrations from the input context, which effectively shortens the prompt and addresses one of ICL's key limitations: the context window size.

Disentangling latent shifts not only improves performance but also provides insights into the internal mechanisms of ICL, potentially aiding model interpretability. While this is not our primary focus, similar ideas have been explored, such as Todd et al. (2024)'s use of function vectors to represent and extract task functions from LLMs. Our work complements these studies by offering a practical approach to enhancing ICL and contributing to model interpretability.

Questions

Q1

Yes, we fine-tune using the same number of labeled instances as in the other methods to ensure a fair comparison across approaches.

Q2

For STICL, demonstrations are labeled instances drawn from the training set and are used by the teacher model to generate pseudo-labels for queries, which are sampled from the unlabeled set. During test time, the information from these demonstrations is effectively encoded in the adapter, so they are not used as part of the input context. Instead, queries from the test set are directly used to evaluate the models, with the adapter providing the necessary task-specific context.

Q3

Exactly, the "unlabeled samples" are drawn from the training set with their labels removed. Using ground truth labels instead of pseudo-labels to fine-tune the student model improves performance, though it would require labeled data in practice. This supports the hypothesis that pseudo-label correction is key to weak-to-strong generalization, as STICL’s student model refines corrupted pseudo-labels to better align with ground truth labels.

Q4

Yes, we believe this captures the main idea of STICL.

We'd like to kindly remind you about our responses to your review. If there's anything further you'd like us to clarify, we'd be happy to provide additional details.

Thank you for your additional feedback and support for the acceptance of our paper. We appreciate the time and effort you’ve taken to engage with our work and provide constructive suggestions. We would like to take this opportunity to further clarify our views on the points you raised:

• Clarification of "ST-ICL": We appreciate your feedback regarding the terminology and agree that it is important to maintain clarity around the definition of ICL as update-less learning. While we intentionally framed our method to emphasize its connection to ICL through the teacher model, we recognize that the current terminology could lead to potential confusion. To address this, we will revise the expanded name of the method to Self-Training on ICL, which we believe more accurately reflects its relationship to the teacher model and ensures consistency with established definitions. Additionally, we acknowledge the angle toward fine-tuning that you highlighted. Our method, at its core, is indeed a lightweight fine-tuning approach that integrates the benefits of ICL through the teacher. We will address this in the manuscript to provide a better understanding of the method's positioning.

• Perspective on disentanglement: While we understand your suggestion to de-emphasize the perspective on disentanglement, we respectfully maintain our position as outlined in our initial response. We believe that disentangling latent shifts is central to both the practical benefits and the conceptual insights offered by our approach.

Thank you for your valuable feedback. We are glad that you find the paper well-explained and easy to follow while offering strong empirical evidence supporting the effectiveness of STICL.

While Table 3 discusses OOD generalizability, it would be beneficial to include a cross-tabulation across datasets to show the performance in-domain.

We initially chose not to repeat the numbers from Table 1 in Table 3 to maintain conciseness and avoid redundancy. However, we understand that presenting OOD results alongside ID results may aid in comparing the performances directly. We will update the table in the revised manuscript.

Comparison with MetaICL: This study’s use of adapters for STICL could be interpreted as an adaptation of MetaICL, making it important to explicitly clarify the differences between the two approaches.

We acknowledge that MetaICL is related to our work. We aim to clarify the differences between MetaICL and STICL and explain why we initially chose not to emphasize MetaICL, though we will now include it in the revised manuscript. While both methods aim to enhance ICL through some form of fine-tuning beforehand, they differ significantly in their underlying principles, making MetaICL less central to the focus of our paper. Specifically, we view STICL as more closely related to the disentanglement approach proposed by Dai et al. (2023), which we discuss in the paper. The key distinctions between MetaICL and STICL are:

• MetaICL fine-tunes the entire model using supervised training across multiple tasks, suited for smaller models like GPT-2 Large (774M). It relies on labeled data for task generalization but does not address latent shifts between demonstrations and queries;

• STICL: uses a teacher-student self-training framework, generating pseudo-labels for adaptation without additional labeled data. It updates only adapter modules, making it efficient for larger models, and explicitly disentangles latent shifts.

We have now conducted additional experiments with MetaICL and made necessary adaptations to enable a fair comparison with STICL-S. Since STICL uses 100 unlabeled instances, we incorporated their corresponding labels into MetaICL's supervised fine-tuning process to align the setups.

The experiments were conducted on Llama 3 (8B), with configurations of 16 labeled + 100 unlabeled instances for STICL-S and 116 labeled instances for MetaICL. We use batches of 16 instances for MetaICL in individual prompts, requiring 8 iterations to complete the fine-tuning for 116 instances. The results (averaged over 10 runs) are as follows:

STICL-S outperforms MetaICL across all datasets despite using fewer labeled instances during training. This improvement is likely due to weak-to-strong generalization taking effect in STICL-S, where the additional labeled data seems to have a marginal effect on MetaICL. Overall, these results further underscore the potential of combining labeled and unlabeled data through self-training, as demonstrated by STICL. We will include the additional findings in the revised manuscript.

Questions

Q1

We randomly draw subsets of demonstrations from the training set.

Q2

We use rank r=8 for LoRA. Other configuration parameters are provided in Appendix D.

Q3

To our knowledge, no prior studies have proposed this exact approach of shuffling demonstrations, likely due to the limited number of works employing fine-tuning to improve ICL. Our motivation arises from well-known challenges in ICL related to demonstration ordering. Liu et al. (2024) observed how LLMs tend to overemphasize information at the beginning and end of a context due to primacy and recency biases. Shuffling demonstrations across epochs can mitigate this issue by varying their positions and stabilizing the ICL process.

Q4

The relationship between the Lipschitz constant and weak-to-strong generalization may not be immediately clear, so we appreciate the opportunity to clarify. As explained in Section 4.1, a key factor for weak-to-strong generalization is robustness to input variations — the model's ability to maintain stable outputs under small perturbations (Wei et al., 2021; Lang et al., 2024). To test this, we approximated the Lipschitz constant, which bounds the maximum output change for any input change (Khromov & Singh, 2024), by calculating the norms of input-output Jacobian matrices for STICL, PBFT, and ICL. Our results show that STICL has a lower approximated Lipschitz constant, indicating greater stability and providing empirical evidence of its local consistency.