Where does In-context Learning Happen in Large Language Models?

We would like to thank the reviewer for their time and comments on our work. Regarding the listed weaknesses of our paper:

"The practical utility of the findings is questionable."

We respectfully disagree that our findings hold little practical utility. LLMs are increasingly being adapted to, and used as, task-specific models; For instance, LLMs have shown great potential in replacing traditional supervised Machine Translation models. Our results hold incredibly high practical utility for scenarios where we require an LLM to adapt to a specific task using in-context examples and prompts, which will be used for every forward-pass of the model. As the length of the in-context examples increases (which we show does not significantly alter the position of critical layers), our results become even more relevant.

doubt whether the proposed improvements to inference efficiency in Section 5 can be applied in practice.

While the exact position of critical layers is task-specific, finding them requires a one-time cost using only forward-passes of the model (e.g. no gradients or backpropagation is used). This is quite straightforward to implement, is highly reproducible, and requires no additional memory beyond forward inference.

The experiments are entirely based on multi-head Attention, which limits the applicability of these findings. Most current models use other attention methods, such as the grouped-query attention used in Llama-2. The findings in this paper may not apply to new attention variants, which already consider partial context attention.

We do in fact use Llama-2 in our experiments although we did not denote this in the paper, as it was the default Llama family available through hugging face at the time. Specifically, the checkpoints we use are https://huggingface.co/meta-llama/Llama-2-7b-hf and https://huggingface.co/meta-llama/Llama-2-7b-chat-hf. We sincerely apologise for the confusion.

However, we also note that (to the best of our knowledge) Llama-2 uses standard multi-head attention. Grouped query attention (GQA) is used in the more recent Llama-3, which was released only a month ago and is therefore out of the scope of our submission. We analyze a broad set of highly utilized LLMs which suggests that our findings are broadly relevant to the LLM community. Our overall findings (that some layers are critical to in-context learning and that attention to the entire context is no longer necessary past a critical point) are also not specific to any particular form of attention.

Dear reviewer, you raised two weakness which we provided responses to... we hope you might reconsider your overall assessment about the contributions of the paper.

We would like to thank the reviewer for their time and comments on our work. Regarding the listed weaknesses and questions:

Layer-wise masking. The authors only release first j bottom layers and conclude that bottom layers are important.

Our conclusion is not that bottom layers are important and top layers are not, but rather that there are critical layers. We write that “Doing so allows us to study how critical each layer is, where critical layers is loosely defined as those that have a large negative impact when masked.. (line 231-232) ”.

We explicitly test the relevance of individual layers across the entire model. Showing here that dropping an individual layer (layer-wise masking) significantly harms performance does suggest that this layer is important because performance would not suffer if its functionality was re-implemented by a different layer.

If the authors mask layers with a reversed order, I suppose the conclusion would change to that top layers are important.

The reviewer may be referring to layer-from-context masking (Figure 2,3),

The masking order that we have used, retains the original order of processing that the transformer sees during training.

"All masks operate from the j-th layer ($\ell_j$ ) onwards, i.e. masking from $\ell_{20}$ means causally masking out attention to all context positions from $\ell_{20:n_\ell}$, where $n_\ell$ is the total number of layers. To construct Fig 2, we increment $\ell$ from 1 to $n_\ell$ and apply the set of masks $\{m(j, u)\}^{\ell:n_\ell}$ in each experiment and observe the performance of the model" (lines 122:125)

The reverse order would mean skipping several layers of attention before passing the representations to the later layers, which may result in token representations that fundamentally differ from what the later layers have seen during training. It would not be clear whether performance was impacted because the context was masked until that point, or because the masked token representations no longer match the expected representation distribution at that layer.

pretrained LLMs solve task with multiple modules. And different layers may share similar functions for solving that task.

We do not present any findings in this paper that conflict with this interpretation! In our masking experiments, we only mask over the context, not the entire test input provided to the LLM. The goal of the paper is to demonstrate where in-context learning (in-context task recognition) happens, not where the LLM "solves the task". (Line 130) If this is unclear, we hope to be able to clarify.

Token masking. The three token masking strategies are not comprehensive.

To recap, the three token masking strategies are

• Has no instructions, and masks examples ($\overline{`Ex`}^{Mask}$)
• Has instructions and masks examples ($`Instr`\overline{`Ex`}^{Mask}$)
• Both instructions and examples masked ($\overline{`Instr`}\overline{`Ex`}^{Mask}$) -- “dropping all of them”

We described the significance of $\overline{`Instr`}\overline{`Ex`}^{Mask}$ i.e., “dropping all of them” in Section 4.1, Layer-from Context masking.

"Under this causal masking treatment masking from layer ℓ, the model must rely on the representations of the target input sentence from layer ℓ + 1 only to complete the task; if the target sentence representations do not already encode the target task (translation into a specific language) then the model will fail to generate translations." (line 126:129)

To put in another way, “Dropping all of them” and showing that performance continues to be maintained past a certain layer, demonstrates that the major work of in-context learning has occurred before that layer (Figure 2, 3). This finding is extended to a varying number of In-context examples (Figure 6).

many existing works show that ICL examples affect the performance much.

We agree that ICL examples are known to affect performance. Our experimental results are averaged over 5 random seeds (random in-context samples are used), to present findings which generalise beyond the choice of specific examples.

Line 117: can j be equal to I, i.e., (j<=i)?

Yes, thank you for catching this typo!

Thank you for the deep appreciation of the work!

It's not clear why “task recognition” v/s actual task performance was treated as a metric here. It seems that the learnings would be very different between the two. Is task recognition in itself important enough to celebrate over computational wins? Some clarity would be good in the paper to ensure readers understand the importance of "task recognition". Also, is there a cutoff (like BLEU > 0 means task recognition)?

It’s a fair point. The process of task recognition, happens over several layers. We think task performance as a metric is easier for people to understand.. Because finally the practical takeaway is that you can reach the model's ceiling task performance without needing full-processing over the context - that’s the computational win.

Sec 4.3 could use better pointers. It's not clear which model is instruction tuned v/s not.

Understood, we will fix it.

Sec 6.2 presents an interesting experiment; but the figure needs to be improved. The explanations don’t match the legends of the figure (for instance, what is true/false). Also, do we mask the source sentence as well?

We will improve the captioning of the figure. True/False refers to whether there are instructions provided (True) vs not provided (False). We do not mask any of the tokens, but remove the entire attention layer in Section 6.2 (as opposed to Section 3).

Thank you for the insightful comments on the paper and for the positive view of our work. We provide our response to the weaknesses (the questions are closely related)

One notable weakness is the limited exploration of why different models exhibit varying behaviors in terms of their "task recognition point" and critical layers. GPT-Neo has more severe critical layers compared to other models but do not provide a thorough analysis of potential reasons for this difference.

Unfortunately, the differences are not due to easily observable hyperparameters like model size or architecture. To put in another way, why do large models exhibit different characteristics? Training data and training dynamics is a typical suspect. While very interesting, we feel that a thorough analysis is perhaps outside the scope of this paper.

The paper doesn't adequately address potential confounding factors that could influence the results. For example, the impact of different tokenization schemes across models or the potential effects of the specific prompt format used are not discussed in depth.

We would like to discuss with the reviewer on the suggested confounding factors.

• Tokenisation schemes, we’re not sure how an analysis on tokenisation scheme would be relevant to the problem being studied here. We used GPTNeo, Llama2, Bloom, Starcoder, all of which have different tokenisation schemes.

• Specific prompt format. There are two main prompt formats used in the paper, reflecting the task.

1a. “Translate from {L1} to {L2}: Q: {source_sentence} A:, 1b. “Translate from {L1} to {L2}: Q: {source_sentence} A:,

2. Write a function to find the longest chain which can be formed from the given set of pairs. Q: {program_description}, A:

We hope this gives the reviewer some confidence that there is sufficient coverage of the confounding factors highlighted.

Instead, the potential confounding factors that we have investigated more pertinent to the research question are

• With/without instructions
• Number of In-context Examples
• Different models Families
• Instruction Tuned VS Not Instruction Tuned Models
• 2 Different Tasks

The paper's focus on machine translation and code generation tasks, while valuable, raises questions about the broader applicability of the findings. The authors could strengthen their work by discussing potential limitations in generalizing these results to other types of tasks or by providing a theoretical framework that explains why these findings might (or might not) extend to other domains of language model application.

Thank you for the suggestion, we will consider whether it’s possible to frame this better.

The paper uses models of different sizes (2.7B to 7B parameters) but doesn't systematically analyze how model size affects the location of the "task recognition point" or the distribution of critical layers. A more structured comparison across model sizes could reveal important trends.

Thank you for the suggestion, we think this does not change the main point or contribution of the paper but acknowledge that it can be an interesting addition for extended experiments that we are happy to carry out

Limitations

Most sincere apologies we seem to have missed this in the current draft... We would be extremely explicit about the scope of the experiments, and also include those the reviewer has highlighted. (exploration of different models, tasks, sizes).

Currently in the paper, we have several "future work". i.e. limitations that we have not figured out.

• “It is not immediately clear why GPTNEO has such critical layers and suffers compared to the other models, although we note that this is unlikely to be due to size or model architecture as BLOOM is also around the same size as GPTNEO and performs more similarly to LLAMA. We suspect that it could be due to training data or some other factor related to the training dynamics but leave this for future work. “ (line 245:249)
• "Potential reasons for this difference might be due to cross-entropy loss associated with task tuning for MT vs non-specific training on large corpora. We leave this as an open question for future work." (line 299-301)