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

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 17:43:19

Abstract: Self-supervised large language models have demonstrated the ability to perform various tasks via in-context learning, but little is known about where the model locates the task with respect to prompt instructions and demonstration examples.
In this work, we attempt to characterize the region where large language models transition from recognizing the task to performing the task.
Through a series of layer-wise context-masking experiments on \textsc{GPTNeo2.7B}, \textsc{Bloom3B}, \textsc{Llama2-7b}, \textsc{Llama2-7b-chat}, \textsc{Starcoder2-3B} and \textsc{Starcoder2-7B} on Machine Translation and Code generation, we demonstrate evidence of a "task recognition" point where the task is encoded into the input representations and attention to context is no longer necessary. We further observe correspondence between the low performance when masking out entire layers, and the task recognition layers. Taking advantage of this redundancy results in 45\% computational savings when prompting with 5 examples, and task recognition achieved at layer 14 / 32 using an example with Machine Translation.

Main paper: 
  
 Appendix:  


Decision: This paper provides a detailed study on where in the LLM architecture is a task-related information is learnt. Once the task is learned, the authors propose to mask the attention to in-context examples for that task for that layer onwards. This results in compute savings and increased interpretability. The paper is well written and the experiments are thorough. Based on revisions from previous submissions to other conferences the reviewers have improved the paper.

Reviews: review: paper_summary: paper_summary: This paper investigates where in-context learning occurs within large language models, focusing specifically on machine translation and code generation tasks. The authors introduce a "layer-from context-masking" technique to identify at which layer an LLM transitions from task recognition to task execution during ICL. They apply this method to several models including GPT-Neo2.7B, BLOOM3B, LLaMA7B (base and chat), and StarCoder2 (3B and 7B).

The key finding is that models do not need to maintain attention over the entire context throughout all layers to perform the task. Instead, there appears to be a "task recognition point" after which attention to the context is no longer necessary. This point varies between LLMs but generally occurs in the middle layers. The authors characterize this as a three-phase process: initial processing where masking has little effect, a critical phase where masking greatly impacts performance, and a final phase where additional context processing yields minimal gains.

The paper also explores the roles of instructions versus examples, the differences between instruction-tuned and non-instruction-tuned models, and whether these phenomena generalize across different tasks. 
The authors highlight the potential practical implications of their findings, particularly for inference efficiency. They estimate that up to 45% computational savings could be achieved by eliminating unnecessary context processing in later layers.
strength_weakness: strength_weakness: - One notable weakness is the limited exploration of why different models exhibit varying behaviors in terms of their "task recognition point" and critical layers. For instance, the authors observe that GPT-Neo has more severe critical layers compared to other models but do not provide a thorough analysis of potential reasons for this difference. A more in-depth investigation into the architectural differences, training data, or other factors that might contribute to these variations would significantly enhance the paper's insights and generalizability.
- 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.
- 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.
- 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.
questions: questions: 1. You observe varying behaviors across different models, particularly in terms of the "task recognition point" and critical layers. Could you provide more insight into why these differences occur? 
2. Your study focuses on machine translation and code generation. How confident are you that these findings would generalize to other types of NLP tasks? 
3. Your study includes models of different sizes. Do you observe any consistent trends in how model size relates to the task recognition point or the distribution of critical layers?
limitations: limitations: In the NeurIPS Paper Checklist (after appendix), they mentioned that they have discussed the limitations in the conclusion section (Line#592). However, I didn't find any such thing.
paper_summary: paper_summary: In-context learning has emerged as an important paradigm in LLMs. In this paper, the authors attempt to characterize where models learn to “recognize” an in-context task. 

To do this, in-context portions (examples and/or instructions) are masked out after certain layers and are not included to generate model predictions.  An advantage of this style of masking is improved cost. Their experiments give clear results, indicating that for each of the selected tasks (machine translation, code generation), there exists a cutoff layer beyond which masking is acceptable without hindering ability to recognize the task. The most clear utility of this masking is computational savings. The authors also find evidence of 3-phase in-context learning, with the last phase accounting for little to no performance improvements.
strength_weakness: strength_weakness: I have some concerns about the definitions of “task recognition” [more in the Questions], as well as explanations on some sections (like Sec 4.3 and 6.2). It would be important to iron these out.
questions: questions: 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)?

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

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?
limitations: limitations: NA
paper_summary: paper_summary: This paper tries to locate where LLMs handle in-context learning (ICL) and interprets how LLMs perform ICL internally. The authors propose a layer-wise attention masking strategy and conclude that LLMs perform ICL in bottom layers.
strength_weakness: strength_weakness: However, I have several concerns about the work.
1. Layer-wise masking. The authors only release first j bottom layers and conclude that bottom layers are important. This is problematic. If the authors mask layers with a reversed order, I suppose the conclusion would change to that top layers are important. The reason is that pretrained LLMs solve task with multiple modules. And different layers may share similar functions for solving that task. Simply masking several layers or layers with specific order cannot prove whether these layers are important/useless. It’s similar to the glove game: you have 3 gloves with 2 left gloves and 1 right glove. Dropping 1 left glove randomly cannot prove this dropped one is useless, or conclude that the right one is the only necessary one. 
2. Token masking. The three token masking strategies are not comprehensive. There is no discussion about the individual examples. Either dropping all of them or keeping all of them cannot explain how ICL works, since many existing works show that ICL examples affect the performance much.

Overall, I think this paper tries to explore an important research question, and I would like to encourage the authors to address my two concerns.
questions: questions: Line 117: can j be equal to I, i.e., (j<=i)?
limitations: limitations: See weakness.
paper_summary: paper_summary: This study primarily investigates where in-context learning occurs within GPT-style models. Specifically, it explores the stage at which a model transitions from functioning as an in-context learner to a task-specific model. By applying layer masking to the instruction and in-context examples in machine translation and coding tasks, the authors observe performance changes to understand the internal mechanisms of in-context learning. They find that certain critical layers within the model are crucial for in-context learning.
strength_weakness: strength_weakness: 1. The practical utility of the findings is questionable. I doubt whether the proposed improvements to inference efficiency in Section 5 can be applied in practice. For example, critical layers differ across tasks, and even within the same task, they can vary between subtasks (e.g., en→fr vs. fr→en). Identifying critical layers typically requires significant cost.

2. 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.
questions: questions: Please refer to the weaknesses section.
limitations: limitations: --

Comments: title: Reply
Review clarifications
Review clarifications
comment: Thanks for the reply. However, my concerns are not well addressed in the reply: the layer-masking strategy is problematic, and thus the following obtained conclusions are not reliable to me. Therefore, I'll keep my score.
Thanks! My score remains the same. I do think that the motivations for task recognition should be explained better in the paper.
Thank you for clarifying and sharing your expected scenario. We would like to address this scenario, and your comments around it:

> this method need one forward-pass for each task, which makes it practical only when there are many requests for the same task. But that's not the case when we deploy our model in reality. In practice, we receive diverse requests from different users.

The scenario we describe of adapting a general purpose LLM to a specific task, after which the LLM processes many requests for the same task, is an increasingly popular technique. For instance, in both research [1] and industry [2], there has been a shift towards using prompted LLMs to handle machine translation, under which a fixed prompt and model is used to handle many requests for a single task (i.e. translation from one specified language to another). This is precisely a setting that we study, and we therefore argue our analysis has very real practical implications for these techniques.

Thank you again for your continued engagement with us!

[1] A Paradigm Shift: The Future of Machine Translation Lies with Large Language Models,  Lyu et al., 2024
[2] https://www.welocalize.com/insights/google-teams-up-with-welocalize-to-test-its-adaptive-translation-llm-solution/
Thank you for your reply. Let me clarify that in my previous response, I mistakenly wrote llama-3 as llama-2. As you mentioned in your reply, this method need one forward-pass for each task, which makes it practical only when there are many requests for the same task. But that's not the case when we deploy our model in reality. In practice, we receive diverse requests from different users.
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.
Dear reviewer, if our response has helped to clarify the misunderstandings, we greatly appreciate if you could reconsider your assessment!