Where does In-context Learning Happen in Large Language Models?
Authors: Suzanna Sia
David Mueller
Kevin Duh
Venue: NeurIPS 2024 poster (2024)
Keywords: Explainability
self-attention layers
redundancy
LLM
in-context learning
machine translation
text to code
Primary area: interpretability_and_explainability
TL;DR: We report that LLMs have a 'task recognition' point, where self-attention over the context is no longer necessary.

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 GPTNeo2.7B, Bloom3B, Starcoder2-7B, Llama3.1-8B, Llama3.1-8B-Instruct, 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. 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. Our findings also have implications for resource and parameter efficient fine-tuning; we observe a correspondence between strong fine-tuning performance of individual LoRA layers and the task recognition layers.

Where does In-context Learning Happen in Large Language Models?
Suzanna Sia Johns Hopkins University ssial8 jh. edu
David Mueller Johns Hopkins University dam@ cs. jhu edu
Kevin Duh Johns Hopkins University kevinduhêcs: -iu..
Abstract
Sepf-super 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 exam- ples. In this work, we attempt to characterize the region where large language models transition from recognizing the task to performing the task. Through a GPTN BLOOM3B,
series of layer-wise context-masking experiments and STARCODER22-75 LLAMA3. --BB, LLAMA3. 1 1-B.....--BBBB1 INSTRCTT,,,, on Machine Translation and Code generation, we demonstrate evidence of "task recognition" point where the task encoded into the input representations and attention to context is longer necessary. Taking advantage of this redundancy results in 45% computational savings when prompting with 5 examples, and task recognition achieved at layer 14 32 using example with Machine Translation. Our findings also have implication for resource and parameter efficient fine-tuning; we observe correspondence between fine-tuning performance of individual LoRA layers and the task recognition layers.
Introduction
In-context learning (ICL) refers the phenomenon which large generative pretrained transformers (GPTs) perform tasks with no gradient updates when shown task examples or descriptions in their context [13, 12]. Recent work on in-context learning has focused on prompt-ongineering, treating GPT models as black boxes and focusing on which examples to provide in-context [53, 14, 21, 22, 57]. However, many of these works apply surface level interventions leaving the internal mechanism of task recognition in GPT models largely not understood.
In this work, we ask where does in-context Learning occur in GPT models? Our view of In-context Learning is that of "task recognition" not "task learning" [69, 47]. While in-context learning in GPT models appears to be generally applicable any natural language task, study task location, focus on two tasks. Machine Translation (MT) and Code generation, in evaluating whether the model has recognized the task. For MT. the model must generate tokens in a diffferent language. For Code generation, the model must produce working program in the correct programming language. These two tasks are unlikely be "learnt" from following patterns, and are more complex than lookup in associative memory for simple Question-Answer tasks.
there is little to no ambiguity
We focus multi-head attention layers unit of study, as the self-attention mechanism is designed to allow the model to attend to it's context during generation of the target sentence [59]. Using causal masking over different parts of the context demonstrate that there exists "lask-recognition point after which attention to the context is no longer necessary (Section 4). Concurrent and previous
"Corresponding Author, suzyahyah@@gmail.cmm and Code Repository https:/ httpp::/github.... com/ wawyahyyhhwhere_ddae www.dddsssinccneeeeeeeeee learning hhwpsiggghpppp_iii____e LLMs
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

work on Task and Function Vectors [32, 58] have also characterised a similar phenomena where the activations induced by in-context examples can be used to control tasks in the model.
We further characterise this phenomena, by studying the effect of on various ablations of masking self-attention over the instructions, examples, and even the query sentence itself. We report that not only is it unnecessary to compute self-attention across the instructions and examples, in later layers self-attention over the query itself may also be redundant.
This work informs the design of efficient inference and training for LLMs with the following contributions
1. We discover large computational savings when the context is several times longer than the test source sentence, a typical phenomena in prompt engineering (Section 5).
2. Parameter efficient fine-tuning corresponding to the phenomena of in-context learning. We observe that very lightweight tine-tunin of LoRA parameters [34] are most effective at earlier layers of the model compared to the later ones (Section 6). The effectiveness of the LoRA training corresponds directly to the layers that occur before the 'task recognition point.
We further investigate the extent of MT task redundancy using differentiable Lo regularisation to train discrete attention head gates (Section 5.1) and find that only around 10% of the attention heads can be fully masked. This indicates that the attention- all of the context that can be redundant. This fundamentally differs from the literature in supervised learning where more than half of the attention heads can be pruned, and Transformers are highly specialised for particular tasks [61, 44, 7].
themselves are not redundant, is attention over
2 Background
In-Context Learning was first demonstrated by [13] who showed that GPT-3 could be used to perform huge variety of tasks without any task-specific parameters or training, by conditioning the model's generation on a prompt which included a few labeled examples of the task of interest. Since then, interest in using GPT models for ICL has grown significantly [45, 3, 66], with several recent works introducing methods such as Instruction-tuning [55, 63] or chain-of-thought prompting [64] to improve downstream ICL accuracy. One key characteristic of In-context Learning is its reliance on prompt examples demonstrating the task that the model should carry out [52].
In-context Learning as Task Recognition. Ostensibly, ICL work for nearly any task that can be defined or described in natural language, and therefore has potential for incredibly broad impact. However, ICL can often still underperform supervised fine-tuning [9], prompting research in analyzing the mechanisms underlying ICL. One line of work studies in-context learning with linear functions, typically linear regression, characterizing the learnability of these functions with ICL [39, 28] and even the learning algorithm a transformer uses [2, 18, 62). A second body of work suggests that in-context learning locates existing latent concepts (tasks) which have been already learnt during pretraining [69, 65]. Notably, [58] describe function vectors which are robust to changes in context. [30] try to characterise the extent of task recognition from the pre-training data. Although there have been many studies on task recognition, our work presents a complementary perspective for task recognition, by demonstrating that there exists a point in the model's layers where the task has been located and causal selffattention onto the context is no longer needed for the model to perform the task.²
Transformer Layers and Self-attention as the Unit of Study. Many works study layers of the model as natural unit of analysis for interpretability [33, 20, 48, 24, 8, 54]. We highlight some of the work which is more closely related to task performance. [68] study the layer-wise adaptability by a hidden-state variability ratio while [60] study evolution of representations in MT-supervised transformer models. [49] studies when model layers can be skipped by feeding intermediate representations into the final output layer a pre-trained supervised model. Our work
In our experiments investigating where "task recognition" happens, we consider the actual "task perfor- mance" score as it could be possible to recognise the task as Machine Translation or Code generation, yet perform less well on it.

---

adds to this body of work by considering the perspective of when and where layers are responsible for task location in in-context learning models.
The self-attention mechanism specifically has been highlighted as a source of redundancy by many previous and concurrent works [10, 46, 31]. This is due it's causal structure over the input symbols under the specific context of the input sequence within it's context window [51]. In this paper, we study a major source of causal redundancy the input, the "prompt examples" that are provided as input-output demonstrations to the model for "in-context learning".
Transformer overparameteriiztion and redundancy has been active area of research [19] with multiple works suggesting to adapt transformer inference depth [36, 15, 25]. While we draw inspiration from these. our main objective is not to compress models for inference, but to highlight the redundancy in computing over long context token sequences.
3 Data and Settings
Models We use GPTNE02.77 7B [11]. BLOOM3B [56], LLAMA3.1-8B and LLAMA3.1 1-..A...1BBB..... --B--uutttt in all of our experiments with Machine Translation. For code generation, we used LLAMA3. 1-8B- Instruct[23] and STARCODER2-7B[44 has 30 layers and 32 heads, I.IAMA2-7B and LLAMA3. 1-8B has 32 layers and 32 heads and STARCODER2 has 30 layers and 24 heads. The checkpoints we use are from Meta AI (for LLAMA) and the transformers library [67]. STARCODER2 and LLAMA models utilises grouped-query attention [1]. while the rest of the models use "regular" multi-head self-attention.
GPTNF02. 7B has 32 layers and 20 heads, B1.00M3B
GPTNEO was trained on The PILE [27], an 825GB text dataset which consists of roughly 98% English data. Despite being mostly monolingual, The PILE contains Europarl which GPTNEO was trained on at document level (rather than a sentence level). Conversely, BLOOM was trained on the ROOTS corpus [38], a composite collection of 498 datasets that were explicitly selected to be multilingual, representing 46 natural languages and 13 programming languages. LLAMA training data consists primarily of common crawl, C4, wikipedia, stackexchange as major sources. STARCODER2 was trained on Github as well as Arxiv and Wikipedia. To our knowledge, there has not been any reports of sentence level parallel corpora in the training datasets of these models.
Data We test our models using two datasets, FL.ORES [29] for Translation and HUMANEVAL. for Code generation. For FLORES,, we experiment with en ++ fr (main paper) and en en-> pt (appendix). Prompt examples are drawn from the development set. We evaluate the generations using BLEU scores, following the implementation from [50]. For HUMANEVAL[I6), we evaluate on the execution accuracy of the generated code using the Pass @1 metric. As HUMANEVAL does not have explicit train set, the prompt set is drawn from the Mostly Basic Python Program (MBPP) dataset [4]. To account for example selection and ordering effects, all inference runs repeated with 5 randomly sampled prompt example sets.
Prompt Format Our prompts may consist of instructions, examples, both, none. Importantly, we adopt neutral delimiters, "Q:" and "A:" to separate the prompt and the start machine generated text. This ensures that the models do not have any information from the delimiters what the task is and must recognise the task from examples.*
For the translation task, when no natural language instructions are used the model input will be Q: [source_sentencee } A: Instructions are given in natural language and take the form: Transl late from [L1) to (I2): Q: pouurcesentencee
A:, where I.1 English and L,2 French if the source and target languages are English and French respec- tively. Examples are given after instructions, and similarly delimited by Q: and A:. See Appendix:
Table 1 for an example.
'In-context learning models have been found be sensitive these order effects [42].
'In an earlier exploration, we found that supplying the model with language indicators only, e.g., "English:", "French:" or "English:". "Python:", was sufficient for strong models (LLAMA3.1 1-BB and LLAMA3. 1-BB. INSTRUCT) to perform the task without seeing any instructions examples in the context.
3

---

Ex
Mask
Inst Ex Mask
Instr, Ex
Mask
0
Atterdoo Block NOM
Attertion Pacck POM
Examples
Test Sentence
Instr.ctions
Exempies
Sentence
Instrutions
Examples
Test Sentence
Name
Instr
EX
Ex Mask
Query
N
Y
Q: ..... A.... Q: <query> A:
Instr, Ex
Mask
Query
Translate French to English: Q: At... Q: <query>> A:
Instr, Ex
Mask
Query
Y
Y
Translate French to English: Q: .... A Q: <qquery>> A:
Instr, Ex, Query
Mask
Translate French to English: Q: ..... A: Q=<query>> A:
Instr, Ex, Query Mask
Translate French to English:
A
<query>> A:
Figure 1: (Top): Graphical explanation of Masking the Attention over Instructions and Examples.
The leftmost image has instructions and masks examples (Instr, Ex ExMask), while the right image has both instructions and examples masked (Instr, Ex Mask) (Bottom): We demonstrate which components of the input prompt masked for each setting that we experiment with. The overline of the setting name describes which portion of the input is highlighted (and thus masked). N/Y refer to absence presence of either Instruction (Instr) Examples (Ex). Although (i.e. instructions and examples), in
primarily
concerned with the effects of masking out task lesk-identifying context some experiments additionally consider masking out the input query well.
For the code generation task, when no natural language instructions are used, the model input will be Q: lenggra_descrrree
where the program rrogram_descripp descripts is Instructions are given in natural language and take the form: "Write program for the following task:".
4 Where Does In-Context MT happen?
4.1 Analysis Methodology: Layer-from Masking
In-context learning differs from task-specific supervised learning in that, during test time, the desired task must be identified, or learned, from the context first and then applied to the input. what stage in the feed-forward computation does causal Large Language Model transition from an in-context learner translation or code-generation model? To explore this question, we introduce layer-from context-masking which masks out all attention weights to the context (instructions, examples, or queries) from certain layer onwards (see Figure for graphical description).
For Causal Decoder-only Transformer Language Models, given each position the Attention weight Ogj context positions j.j i be computed by 0ii softamax( QK"
va, )ij. Each element in ) is the dot product between query vector and key vector q k. where y WyT,, ky W for trained weight matrices W; and W We apply the attention mask over the input so that the attention score is (g: k;) + míj, u). Here, u are the tokens that we wish to mask, and
(QKT)
if
layer. payer.j.attre aatattnnmsssil torch. finfo aaatt_mmssttttt dtype) .min
i, (ssaat)) :u (end)
m (j u) = -00 if if Eu
implemented in practice the smallest floating point value for that datatype. All masks operate from the j-th layer (Ez) onwards, i.e. masking from Eg means zeroing out attention to all positions
Readers should note that there is a We and W,, weight matrix for each layer and each attention head, but we omit the notation this for readability.

---

in from feo:ne. where ne is the total number of layers. To construct Fig 2, we increment é, from 1 to ne and apply the set of masks {mi[j, u) fis in each experiment and observe the performance of the model.
When masking input tokens from layer l, the model must rely on only the information in the hidden state representations of the remaining, unmasked tokens from layer L + since representations of the masked tokens can no longer be incorporated moving forwards: if the unmasked representations do not already encode enough information complete the task (e.g.. Machine translation) then the model will fail to generate the correct output. Our intuition is the following: if, at layer e, the model can perform the target-task without attending to the task context-task-ionet task-idenniifying tokens such as instructions and examples- tten information about the task has already been incorporated into the query representations and the model has identified, or "recognized", the target-task by layer l.
4.2 Experiments on Layer-from context-masking
In Figure (Table) we show the various masking treatments that we apply to the input in our experi- ments. We ablate over 3 different task-context masking settings to test the impact of various parts of the context: providing only task examples and masking them from given layer (Ex Mask QQuery;; including the instruction and but only masking out the examples (Instr, Ex" Mask Query); and including instructions but masking them with the examples (Instr, Mask Query). As a control in our experiments, we also experiment with masking the entire input (Instr, Ex, Query Mask) to study whether the model needs to attend to any input beyond a certain layer, and with masking only the query (Instr, Ex, Query Mask to study whether masking the context vs the query have similar effects. For each masking setting, we apply the mask from all layers in the model (j = 1, .n) and observe how task performance is affected at each layer. When examples are provided in-context, we use 5 examples per prompt and we re-sample these examples to control for variance in example selection.
4.3 Results
Models do not need to maintain attention over the task context past a certain layer to perform the task. In all models, we observe that when applying masking from {m(j, u jé over the task context, models obtain their maximum performance well before the final layer, i.e., when & < ng. The results our experiment for en -> and fr + en are shown in Figure 2, and additional experiments for GPTNEO and BLOOM en pt and pt en are shown in Section A.4. Different models reach this platcau point at different layers: in GPTNEO this point occurs around layer 25, in BLOOM this point occurs around layer 15-20, and in LLAMA models this occurs around layer 13-15. As English is the dominant language in most model's training, models can successfully perform translation into English upon earlier layers of masking, than translation out of English. Once this plateau is reached, the models benefits only marginally, if at all, from retaining attention to the context, suggesting most of the task "location" has already occurred.
We observed that with LLAMA-3.1 models, when masking only task context, there a jump in performance from nearly negligible to nearly optimal in the course of few layers. Conversely, when the query is masked we see both performance begin to rise much later in the model and the approach to optimal performance occur much more slowly, often plateauing only few layers before the end of the model.
In some models, there may be a point where forward computation independent of even the query. There is also a point in the model where it no longer needs access any input query tokens. We find this effect to be much less pronounced in GPTNE02. 7B, BLOOM3B and STARCODER on the code generation task, and thus maybe a characteristic of Llama models training.
There exists critical layers for task location. Prior to the task recognition point, around the middle layers of the models, moving the context mask up layer results in significant increase to performance. We consider these critical layers, as instead of a gradual increase performance, we observe very steep jumps, accounting for more than 80% of the model's ''s ceiling performance for that task. We conjecture that the model is locating the correct task during processing in these middle layers, after which the context is no longer necessary to perform the task.
5

---

Masking LLAMA-3.1-8B context-from (en-fr)
Masking context-from (en-fr)
M
BLEU
BLEU
20
20
Ex Query
ExHatQuery
Query
Eueryyuery Instr, Query
Instr, Query
Query
Ex. Quereeyruu
Instr. Ex. Query
Instr, Ex. Queryta
Layer
Layer
Pevorry XLCOMOO comsutiooo
Paxingg LAMAS corset//cccen
Mast LLANS- IEEESTRNC conedee that
Figure 2: Layer-from context-masking experiments for LLAMA3. 1-8B, LLAMA3. 1AAMMA3.--B-INSTREEE 1-SR-INTTRUCT en en-frr (main figure), and GPTNE02. 7B, BLOOM3B, en - fr, LLAMA3. 1-8B. LLAMA3. 8B-INSTRUCT on fr en. The graphs show translation performance when masking contexts from the 7th layer onwards. Different lines indicate different masking treatments, as described in Figure 1. The dashed black line is the performance when masking of the input occurs.
Masking STARCODER2-7B context-from (en-py)
Masking
context-from (en-py)
ExteeiQuery
ExMeikQuery
Instr, Query
Instr, Exerr
Instr. Query
Instr. ExtasQuery
0.30
Instr.Ex.Query
Instr. Ex. Ex.Query*Me
PASS@1
PASS@1
0.15
0.00
0.00
20
25
20
25
Layer
Layer
Figure 3: Layer-from context-masking experiments for Starcoder2-3B, Starcoder2-7B, Llama7b, Llama7b-chat on text code generation task. The graphs show translation performance when masking contexts from the jih layer onwards. Different lines indicate different treatments of the instruction, described Figure The dashed black line is the performance when shown both examples and instructions without masking.
Overall, our findings suggest a 3-phase process to in-context learning: in the first phase, moving the mask up makes little difference performance, which close This suggests that the context has not influenced task location at all. In the second phase, shifting the mask upwards makes large difference performance, suggesting that model has started to locate the task but can improve significantly with more processing of the context. Finally, in the third phase, shifting the mask upwards again has little-to-no effect on the performance, suggesting that model has fully recognized the task translation and no longer requires the context to interpret the task.

---

). This is to expected as Masking GPTN2.78 78 Context From Layer (on In)
Masking BLOOM3R Context From Layer (en ft)
25
13
20
22
25
50
15
30
Layer
Layer
Figure 4: Layer-from context-masking experiments for GPTNeo and BLOOM on en + fr investigat- ing number of examples in the Ex" Mask mask setting. The dashed black line refers to no instructions and examples.
4.4 Instruction-tuned vs Non-instruction Tuned Models
When comparing non-instruction tuned LLAMA3. 1-8B instruction-tuncd models LLAMA3. 1- IN-SRTTUCTION, we do not observe any noticeable difference in where performance plateaus, ie., where the model no longer requires attention over the context. This occurs around layers 16 for both LLAMA models in en + fr and around layer 13 for fr + en. The main difference is that instruction-tuned model is able to achieve better performance in the earlier layers for the setting where instructions are present and examples are masked (Instr, Ex Mask), these models are tuned towards following instructions.
Overall find that the observation of task recognition layers and task recognition point is present across both con-instruction tuned and instruction tuned models, and that this presents itself similarly in both types of models.
4.5
Do models have distinct task recognition region regardless of the type of task? (Experiments on Code Generation)
For tasks that the model does not perform fluently, we do observe sharp increase at any particular layer. For instance, for code generation (HUMANEVAL) where the LLAMA2 model performs poorly, we can observe only a very gradual effect of masking the self-attention layers, and not a distinet increase as compared to the L.LAMA2's performance on Translation.
However when consider STARCODER2 while masking instructions instructions, i.e., the Instr,Ex* Enstr, Mask and Mask we again see the same pattern demonstrating the task recognition phenomena on layer 19 of the 3B model, and layer 20-23 of the 7B model.
To understand Starcoder2's strong performance the (Instr, Ex Mask condition, investigations found that instructions and test prompt alone contain sufficient information for the model to recognise that the task is to generate Python program, even though the model is not instruction tuned. This happens as the model is very specialised towards code generation and has a strong prior to generate python code given its prevalence in it's training data.
4.6 The Role of Instructions vs Examples
In separate experiments, we found that when shown only instructions and no examples, GPTNEO and BLOOM models are unable translate, and their performance is nearly at O BLEU Score. For GPTNEO and BLOOM we see that the behavior of the model is similar when no instructions are present (Ex* Mask ) and when instructions are masked (Instr, Ex Mask), However, the model is given complete access to instructions (Inat rEx Mask, it can use the intermediate processing of examples to reach baseline performance earlier.
7

---

(Prompt) Instrnations and Exemppees
Howd
Redundant Self-
Allenton Computation
0.6
Task Processing
Query and Gemerattio
(Masked)
(Prompt) Instructions Examples
Query
Generation
Figure 5: (Left) Illustration of redundancy in self-attention computation based on our masking setup (Instr, Ex Mask Querry. (Right) Visualisation of attention head masks for GPTNEO and BLOOM, learned with La(A = 0.01) regularisation under O 0-prompt train scheme in en fr. A value of 0 (in black) indicates that the attention head is effectively masked by the trained attention gate. Around 10% of attention heads are masked out i.e., redundant, with a majority of them occuring at the later layers for GPTNeo and distributed across layers for BLOOM. fr - en is availble in Section A.7.1
4.7 Does the Number of Prompts Affect Task Recognition?
In Section 4 we study context-masking with fixed number of prompts. However, it is not clear if the number of prompts affects how fast, layer-wise, the model is able to recognize the task. We plot these results for number prompt examples has little effect on which layer the task is recognized at. While there is some variation in performance when the context masked around the middle layers of the model, the final performance plateau occurs at the same layer regardless of the number of prompts.
fr in Figure 4, for both GPTNEO and BLOOM. In general,
find that the
Inference Efficiency
Speeding up transformer inference is of great interest the community [26]. We highlight the potential of speeding up inference time as a direct consequence of identifying where task recognition occurs in the model and redundancy attention from the query and generated sequence tokens over itself (Task Processing) and Self-attention over the prompt Instructions and Examples (Masked). If ceiling performance is achieved, then self- attention over the previous context becomes redundant (Redundant Self-attention Computation).
self-attention processing. In Figure 5, we illustrate self-
Our results indicate that of context-tokens all-together after certain point in the model, with little to no impact on downstream performance. Let &, be the th layer where subsequent layers and match the "ceiling" performance. Let be the number of prompt examples, where each example consists of pair parallel sentences. Then, for model with n4 layers, the amount of processing in terms of speed and memory saved is approximately (n. (r)rrr/n x (k/k+ 1).
achieve significant speedups in inference by removing the processing
mask out the attention of the context across
Using the example of LLAMA3. 1-BB (32 layers) on en-ffr, very close to it's ceiling score after processing the examples at layer 14 (L 14). we no longer need to process examples after ( = 14, under a prompt size of 5 the savings approximately 50%.
from Figure 3 that the model is
For instruction-tunel models which are typically deployed in production, even if we assume that no examples are provided, savings can be non-trivial as very long-form instructions are typically provided to the model in attempt to control it's behavior (prompt engineering).

---

PIEAPO 1-BcStteeenenii
Masting
Figure 6: Performance of models ((LLAMA3. 1-BB and LLAMA3. ----SSRUTTTTT for en+ff) trained with single LoRA layer, where each point the x-axis reflects single trained LoRA layer. The LoRA layer was trained without instructions, and with causal LM cross-entropy loss over next token prediction of parallel translation sentences. 400 and 800 refer to the size of the training set. The layers which are most amenable to lightweight fine-tuning occur in the earlier layers before the "task recognition" point.
Although we had demonstrated the redundancy of self-attention over the input context, the significance of this computational savings extends to all components of the transformer during forward inference. Since all subsequent layers of forward inference no longer rely on computations on previous token positions, all processing related to those redundant token positions (from the task recognition layer onwards) can be effectively removed.
Overall, observing redundancy over the context is not surprising- To explain why models can have such redundancy, we refer to [17] who identify phenomena where attention heads attend almost exclusively to delimiter and separator tokens such as [SEP], periods and commas. This is thought to act as a "no-op" as the value such tokens in changing the current hidden representation is very small. Note that it is then possible to mask entire Transformer layers and still achieve a sensible output due to residual connections the Transformer architecture at every layer.
5.1 Are There Specialised Attention Heads?
A well established finding for supervised encoder-decoder MT models, is that up to 90% of the attention heads can be pruned while minimising fall in translation performance [61, 6, 44). We note that asking about the extent of pruning is a slightly ill-formed research question, it depends on the type of pruning technique used. However broad trends of highly prunable models have been observed in the supervised MT paradigm. For instance, [5] studied attention- head importance for broader set of ICL tasks, finding that the most important heads for ICL occur in the middle layers of the model. We train discrete attention head gates with Lo regularisation for GPTNEO and BLOOM on en + fr (see Section A.7.1). Overall, we report that there are no "few" specialised heads, which directly contrasts with the literature on compression in supervised MT models [61, 44]. Potential reasons for this difference might be due to cross-entrops loss associated with task tuning for MT vs non-specific training on large corpora. We leave this as an open question for future work.
The Adaptability of Task Layers
The layers prior to "task recognition" should contain information about locating the MT task. To test this, we further explore the adaptability of these layers by lightweight fine-tuning experiments on LLAMA3. 1-8B and LLAMA3. 1-SB-INSTRUCH on en ++ fr.
We trained single Low-rank Adaptation matrix (LoRA; (34J) for each layer of the output projection while keeping the rest of the network frozen. 6 This means there were ne individual layers trained for ny experiments in Figure 6, where ny is the total number layers the model.
The model was shown parallel sentences as input, and layers were trained with no explicit translation instructions. We split the dev set of FLORES into 400 and 800 training examples and 200 dev examples, we repeated the experiments with 2 random seeds initialisations. Note that this setup is designed to tune the layers for task location. highly unlikely that the model can learn translation
'We also experimented with the training separate Key, Query and Value LoRA Layers but found this to be less effective.

---

knowledge from this small amount of supervision. The LoRA layers were trained for up to 50 epochs with batch size= 32, learning rate= 4, early stopping patience= and threshold= 0.01, with = 32, r = 8 and dropout= 0.05. These values are default and there was no hyper-parameter optimisation over the training parameters. The cross-entropy loss was computed across the entire sequence, and we used the best checkpoint on the 200 held out dev examples for evaluation.
Without any fine-tuning, performance close to 0 because the model will generate sequences continuing from the source language instead of doing translation. While each layer can be trained to perform better than no fine-tuning at all, tuning different layers have vastly different impacts on performance (see Figure 6). In particular, we find that high performing layers occur at the earlier middle parts of the network, with the peak occuring strictly before the "task loskklocating" layers from Section 4.
6.1 Task Locating Layers are critical for resource efficient fine-tuning
With half the number of training examples (400 instead of 800), the range of trainable layers drop very greatly. For the more challenging direction of generating in French, en - fr, reducing the number of training examples result in none of the layers being successfully fine-tuned for translation task location. For fr -t-een, the range of trainable layers is much more concentrated around layers 10 to 15, which occurs just before the 'task recognition' layers as shown in Figure 2.
We demonstrate that in contrast to common fine-tuning wisdom, additional tuning on later layers in the transformer network has much smaller impact on final performance, and this strong correlated with where the 'task locating' layers are in the model. The major reason for this discrepancy from conventional fine-tuning wisdom, that we are performing extremely lightweight parameter-efficien fine-tuning for task location, and not full fine-tuning on large datasets. Our results should thus be interpreted under the lens of highly resource efficient fine-tuning of layers for task location, and fundamentally different from the wisdom of "true" task fine-tuning.
Conclusion
We demonstrate evidence that In-context Causal Decoder models locate their task at specific layers during forward inference. To study this, we introduced causal masking of self-attention over the context from layer & onwards (Section 4). The findings generalise across 4 models of different sizes and in both non instruction-tunee and instruction-tuned models. We further identify certain layers as task critical, and show that this corresponds the task recognition point of the model (Section A.9) and is not influenced by increasing number of examples (Section 4.7) shown to the models.
Our central finding that models do not need to maintain attention over all of the context across every layer has direct implications for inference efficiency of transformers, with estimated up to 45% cost-savings for llama model with 5 examples (Section 5).
Contrary to common fine-tune fine-tuning wisdom, we show that it sometimes beneficial to target middle layers for fine-tuning the model which could be associated with task recognition ( (Section 6). Finally, we trained attention head gates using differentiable Lo regularisation (Section 5) and found that around 10% of attention heads can be masked. These are mostly distributed across the later layers of the model, providing some support for the idea that later layers are redundant. Although we have characterised this phenomena using Machine Translation and Code Generation, we believe that the broad findings are likely to generalise to other tasks.
7.1 Limitations (and Future Work)
There is limited exploration of why different models exhibit varying behaviors in terms of their "task recognition point" and critical layers. Unfortunately, the differences not due easily observable hyperparameters like model size or architecture. To put in another way, why do large models exhibit different characteristics?
This paper reports on empirical analysis and observations, and currently lacks a more theoretical framework that could explain why this phenomena is being observed.
10

---

Acknowledgments
We would like to thank all the anonymous reviewers for their invaluable comments and suggestions, as well as Daniel Kashabi and Mare Marone for feedback on earlier drafts.
References
[1] Ainslie, Lee- Lee-Th M. de Jong, Y. Zemlyanskiy, F. Lebron, and S. Sanghai. GQA: Training generalized multi-quen transformer models from multi-head checkpoints. In H. Bouamor, J. Pino, and K. Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 4895-4901, Singapore, Dec. 2023. Association for Computational Linguistics.
[2] E. Akyürek, D. Schuurmans, J. Andreas, T. Ma, and D. Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.156611 2022.
[3] E. Akyürek, B. Wang. Y. Kim, and Andreas. In-context language learning: Arhitectures and algorithms. arXiv preprint arXiv:2401. 12973, 2024.
[4] J. Austin, A. Odena, M. Nye, M. Bosma, H. Michalewski, D. Dohan, E. Jiang, C. Cai, M. Tenry Q Le, et al. Program synthesis with large language models. arXiv preprint arXiv:2108.07732, 2021.
[5] H. Bansal, K. Gopalakrishnan, Dingliwal, S. Bodapati, K. Kirchhoff, and D. Roth. Rethinking the role of scale for in-context learning: An innerpeetability-basee case study at 66 billion scale. In A. Rogers, J. Boyd-Graber, and N. Okazaki, editors, Proceedings of the 6lst Annual Meeting of the Association for Computational Linguistics ( (Volume 1: Long Papers), pages 11833-11856, Toronto, Canada, July 2023. Association for Computational Linguistics.
[6] M. Behnke and K. Heafield. Losing heads in the lottery: Pruning transformer attention neural machine translation. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 666426774, 2020.
[7] M. Behnke and K. Heafield. Pruning neural machine translation for speed using group lasso. In Proceedings of the sixth conference on machine translation, pages 1074-1086, 2021.
[8] Ben-Shaul and S. Dekel. Nearest class-center simplification through intermediate layers. In Topological, Algebraic and Geometric Learning Workshops 2022, pages 37-47. PMLR, 2022.
[9] K. Bhatia, A. Narayan, C. D. Sa, and C. Ré. Tart: A plug-and-play transformer module for task-agnostic reasoning, 2023.
[10] Y. Bian, J. Huang, X. Cai, J. Yuan, and K. Church. On attention redundancy: A comprehensive study. Proceedings of the 2021 conference of the north american chapter of the association for compulational linguistics: human languuge technologies, pages 930-945, 2021.
[11] S. Black, G. Leo, P. Wang, C. Leahy,, and S. Biderman. GPT-Neo: Large scale autoregressive language modeling with Mesh- Tessorrenso,, Mar. 2021.
[12] R. Bommasani, D. A. Hudson, E. Adeli, R. Altman, S. Arora, S. von Arx, M. S. Bernstein, J. Bohg, A. Bosselut, E. Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint 1xXiv:2108.072555 2021.
[13] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[14] B. Chen, 7. Zhang, N. Langrené, and S. 7hu. Unleashing the potential of prompt engineering in large language models: a comprehensive review. arXiv preprint arXiv:2310.1477: 2023.
[15] D. Chen, Li, M. Qiu, Z. Wang, B. Li, B. Ding, H. Deng, J. Huang, W. Lin, and J. Zhou. Adabert: Task adaptivee bert compression with differentiable neural architecture search. In International Joint Conference Artificial Intelligence, 2020.
11

---

[16] M. Chen, J. Tworek, H. Jun, Q. Yuan, H. P. d. O. Pinto. J. Kaplan, H. Edwards. Y. Burda, N. Joseph, G. Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374, 2021.
[17] K. Clark, U. Khandelwal, O. Levy, and C. D. Manning. What does BERT look at? an analysis of BERT's attention. In Proceedings of the 2019 ACL Workshop BlackboxNIPP. Analyzing and Interpreting Neural Networks for NLP, pages 276-286, Florence, Italy, Aug. 2019. Association for Computational Linguistics.
[18] D. Dai, Y. Sun, L. Dong, Y. Hao, Ma, Z. Sui, and F. Wei. Why can gpt learn in-context? lan- guage models implicitly perform gradient descent as meta-optimizers. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023.
[19] Dalvi, H. Sajjad, N. Durrani, and Belinkov. Analyzing redundancy pretrained transformer models. In Conference on Empirical Methods in Natural Language Processing, 2020.
[20] N. De Cao, M. S. Schlichtkrull, W. Aziz, and I. Titov. How do decisions emerge across layers in neural models? interpretation with differentiable masking. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 3243-3255 2020.
[21] Denny, V. Kumar, and N. Giacaman. Conversing with copilot: Exploring prompt engineering for solving csl problems using natural language. In Proceedings of the 54th ACM Technical Symposium on Computer Science Education V. I, pages 1 1136-1142, 1122, 2023.
[22] Q. Dong, L. Li, D. Dai, C. Zheng, Z. Wu, B. Chang, X. Sun, J. Xu, and Z. Sui. A survey on in-context learning, arXiv preprint "iXi:233010002344 2022.
[23] A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al-Dahle, A. Letman, A. Mathur, A. Schelten,
A. Yang, A. Fan, et al. The llama 3 herd of models. arXiv preprini arXiv:240722178. 2024.
[24] N. Durrani, H. Sajjad, F. Dalvi, and F. Alam. On the transformation of latent space in fine-tuned nlp models. arXiv preprint arXiv:2210. 2022.
[25] A. Fan, E. Grave, and A. Joulin. Reducing transformer depth on demand with structured dropout. arXiv preprint arXiv: 1909. 1556, 2019.
[26] Q. Fournier, G. M. Caron, and D. Aloise. A practical survey on faster and lighter transformers. ACM Computing Surveys, 35(14s):1-40) 2023.
[27] L. Gao, S. Biderman, S. Black, L. Golding, T. Hoppe, C. Foster, J. Phang, H. He, A. Thite, N. Nabeshima, al. The Pile: An 800GB dataset of diverse text for language modeling. arXiv preprint arXiv:2101.0003
2020.
[28) S. Garg, D. Tsipras, P. S. Liang, and G. Valiant. What can transformers learn in-context? case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598, 2022.
[29] N. Goyal, C. Gao, V. Chaudhary, P.-J. Chen, G. Wenzek, D. Ju, S. Krishnan, M. Ranzato, F. Guzmán, and A. Fan. The flores-101 evaluation benchmark for low-resource and multilingual machine translation. 2021.
[30] X. Han, D. Simig, T. Mihaylov, Y. Tsvetkov, A. Celikyilmaz, and T. Wang. Understanding in-context learning via supportive pretraining data. arXiv preprint arXiv:2306. 2023.
[31] S. He, G. Sun, Z. Shen, and A. Li. What matters in transformers? not all attention is needed. arXiv preprint arXiv:2406. 15786, 2024.
[32] R. Hendel, M. Geva, and A. Globerson. In-context learning creates task vectors. arXiv preprint arXiv:23 15916, 2023.
[33] Hewitt and P. Liang. Designing and interpreting probes with control tasks, arXiv preprint arXiv:1909.03368 2019.
12

---

[34] E. J. Hu. Y. Shen, P. Wallis, Z. Allen-Zhu, Low-rank adaptation large language models. arXiv preprint (tth:2000.096833 2021.
Zhu, Y. Li, S. Wang, L. Wang. and W. Chen. Lora:
[35] E. Jang, S. Gu, and B. Poole. Categorical reparameterization with gumbel-softmax, arXiv preprint arXiv: 161 1.111111111011 144, 2016.
[36] Y. Kaya, S. Hong, and T. Dumitras. Shallow-deep networks: network overthinking. In International conference on machine learning, pages 3301-3310. PMLR, 2019.
Understanding and mitigating
[37] P. Koehn. Europarl: A parallel corpus for statistical machine translation. In Proceedings of machine translation summit x: papers, pages 79-86, 2005.
[38] H. Laurençon, Saulnier, T. Wang, C. Akiki, Villanova del Moral, Le Scao, Von Werra,
C. Mou, E. González Ponferrada, H. Nguyen, et al. The bigscience roots corpus: A 1.6 tb composite multilingual dataset. Advances in Neural Information Processing Systems, 35:3 35:31809- 3 31826, 2022.
[39] S. Li, Z. Song, Y. Xia, T. Yu, and T. Zhou. The closeness of in-context learning and weight shifting for softmax regression. arXiv preprinl wrXiv:2304.13276 2023.
[40] C. Louizos, M. Welling, and D. P. Kingma. Learning sparse neural networks through 1_0 regularization. arXiv preprint 1rXiv:1712.01312, 2017.
[41] A. Lozhkov, R. Li. L. B. Allal, Cassano, Lamy-Poirier, N. Tazi, Tang, D. Pykhtar, J. Liu, Y. Wei, et al. Starcoder 2 and the stack v2: The next generation. arXiv preprint arXiv:2402.1912 19173, 2024.
[42] Y. Lu, M. Bartolo, A. Moore, S. Riedel, and P. Stenetorp. Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 8086-8098, Dublin, Ireland, May 2022. Association for Computational Linguistics.
[43] C. J. Maddison, A. Mnih, and Y. W. Teh. The concrete distribution: continuous relaxation of discrete random variables. arXiv preprint arXiv:1611 1.00772,1..01..... 2016.
[44] Michel, O. Levy, and G. Neubig. Are sixteen heads really better than one? Advances in neural information processing systems, 32, 2019.
[45] C. Olsson, N. Elhage, N. Nanda, N. Joseph, N. DasSarma, T. Henighan, B. Mann, A. Askell, Bai, A. Chen, et In-context learning and induction heads. arXiv preprint arXiv:2209. 1895, 2022.
[46] B. Pan, R. Panda, Y. Jiang, Z. Wang, R. Feris, and A. Oliva. la-redEt: Interpretability-wwaar redundancy reduction for vision transformers. Advances in Neural Information Processing Systems, 09224888-249111 2021.
[47] Pan. What in-context learning "learns" in-context: Disentangling task recognition and task learning. PhD thesis, Princeton University, 2023.
[48] A. Pasad, J.-C. Chou, and K. Livescu. Layer-wise analysis of a self-supervised speech repre- sentation model. In 202/ IEEE Automatic Speech Recognition and Understanding Workshop (ASRU), pages 914-921. IEEE, 2021.
[49] Phang, H. Liu, and S. R. Bowman. Fine-tuned transformers show clusters of similar represen- tations across layers. arXiv preprint arXiv: 2rXiv:2109.0888 100.08006 2021.
[50] M. Post. call for clarity in reporting bleu scores. arXiv preprint arXiv://80.00877,, 2018.
[51] R. Rohekar, Y. Gurwiez, and S. Nisimov. Causal interpretation of self-attention pre-trained transformers, Advances in Neural Information Processing Systems, 36, 2024.
13

---

[52] O. Rubin, J. Herzig, and J. Berant. Learning to retrieve prompts for in-context learning. In M. Carpuat, M.-C. de Marneffe, and I. V. Meza Ruiz, editors. Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human I. Language Technologies, pages 2675-2671, Seattle, United States, July 2022. Association for Computational Linguistics.
[53] P. Sahoo, A. K. Singh, S. Saha, V. Jain, S. Mondal, and A. Chadha. A systematic survey of prompt engineering in large language models: Techniques and applications. arXiv preprint arXiv:2402.07927, 2024.
[54] H. Sajjad, F. Dalvi, N. Durrani, and P. Nakov. On the effect of dropping layers of pre-trained transformer models. Computer Speech & Language, 77:101429, 2023.
[55] V. Sanh, A. Webson, C. Raffel, S. Bach, L. Sutawika, Z. Alyafeai, A. Chaffin, A. Stiegler,
A. Raja, M. Dey, M. S. Bari, C. Xu, U. Thakker, S. S. Sharma, E. Szezechla, T. Kim, G. Chh- ablani. N. Nayak, D. Datta, J. Chang, M. T.-J. Jiang, H. Wang, M. Manica, S. Shen, Z. X. Yong. H. Pandey, R. Bawden, T. Wang. T. Neeraj. Rozen, A. Sharma, A. Santilli, T. Fevry, J. A. Fries, R. Teehan, T. L. Scao, S. Biderman, L. Gao. T. Wolf, and A. M. Rush. Multi- task prompted training enables zero-shot task generalization. In International Conference on Learning Representations, 2022.
[56] T. L. Scao, A. Fan, C. Akiki, E. Pavlick, S. Ilié, D. Hesslow, R. Castagné, A. S. Luccioni, Yvon, M. Gallé, et al. Bloom: A 176b-parameter open-access multilingual language model. arXiv preprint arXiv:2211.05100 2022.
[57] S. Sia and K. Duh. In-context learning as maintaining coherency: A study of on-the-fly machine translation using large language models. arXiv preprini ar arXiv:2305.03573, 20233.
[58] E. Todd, M. L. Li, A. S. Sharma, A. Mueller, B. C. Wallace, and D. Bau. Function vectors in large language models. arXiv preprint arXiv:2310. 15213, 2023.
[59] A. Vaswani, Shazeer, N. Parmar, Uszkoreit, L.. Jones, A. N. Gomez, Kaiser, and I. Polosukhin. Attention is all you need, Advances in neural information processing systems, 30, 2017.
[60] E. Voita, R. Sennrich, and I. Titov. The bottom-up evolution of representations in the trans- former: study with machine translation and language modeling objectives. arXiv preprint 1909.0000 2019.
[61] E. Voita, D. Talbot, F. Moiseev, R. Sennrich, and I. Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 5797-5808, Florence, Italy, July 2019. Association for Computational Linguistics.
[62] von Oswald, Niklasson, F. Randazzo, Sacramento, A. Mordvintsev, A. Zhmoginov, and M. Vladymyrov. Transformers learn in-context by gradient descent, 2023.
[63] Y. Wang, S. Mishra, P. Alipoormolabashi, Y. Kordi, A. Mirzaei, A. Naik, A. Ashok, A. S. Dhanasekaran, A. Arunkumar, D. Stap. E. Pathak, G. Karamanolakis, H. Lai, Purohit, I. Mon- dal, J. Anderson, K. Kuznia, K. Doshi, K. K. Pal, M. Patel, M. Moradshahi, M. Parmar, M. Purohit, N. Varshney, P. R. Kaza, P. Verma, R. S. Puri, R. Karia, S. Doshi, S. K. Sampat, S. Mishra, S. Reddy A, S. Patro, T. Dixit, and X. Shen. HuperrNaaurlllnstuuciioo tion via declarative instructions on 1600+ NLP tasks. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 5085-5109, Abu Dhabi, United Arab Emirates, Dec. 2022. Association for Computational Linguistics.
Generaliza-
[64] J. Wei, X. Wang, D. Schuurmans. M. Bosma, brian ichter, F. Xia, E. H. Chi, Q. V. Le, and D. Zhou. Chain of thought prompting elicits reasoning in large language models. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022.
[65] N. Wies, Y. Levine, and A. Shashua. The learnability of in-context learning. arXiv preprint a7X1v:2303.08895, 2023.
14

---

[66] N. Wies, Y. Levine, and A. Shashua. The learnability of in-context learning. Advances in Neural Information Processing Systems, 36, 2024.
[67] Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, Cistac, T. Rault, R. Louf, M. Funtowicz, et al. Huggingface's s transformers: State-of-the-art natural language processing. arXiv preprint arXiv::910.03777 2019.
[68] S. Xie, Qiu, A. Pasad, L. Du, Q. Qu, and H. Mei. Hidden state variability of pretrained language models guide computation reduction for transfer learning. arXiv preprint (0::22200.00411 2022.
[69] S. M. Xie, A. Raghunathan, P. Liang, and T. Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080 2021.
15

---

A Appendix
A.1 Reproducibility
The Machine Translation dataset that we use, FLORES [29] and the code generation datasets 11 are fully open-source and well-known in the community. All models are open- source and freely available on Huggingface [67]. We used models of "reasonable" size (3B and 7B parameters) that can be run with consumer or academic grade GPUs, making our work reproducible to most academic institutions.
A.2 Impact Statement (Ethics and Societal Consequences)
There are no known ethical concerns as these are exploratory studies on open-source LLMss.
A.3 Prompt Format
Translate English to French.
Q: A discomfort which lasts
A: Un malaise qui dure
Q: HTML is language for formatting
A:
HTML un langage de formatage
Q: After you become comfortable with formatting..
Table I: A single continuous input sequence presented to the model for decoding a single test source sentence "After you become comfortable with formatting.". ". Given the entire sequence as input, the model proceeds to generate the target sequence.
A.4 Additional Results English & Portugese
In addition to the language pairs en + fr and fr + en, we also run experiments on English and Portugese language pairs, for en - pt. Due space limitations, we plot the results of those experiments here. Overall, we see largely identical trends on both directions of English and Portugese what observe on English and French translation tasks, leading us conclude that our conclusions generalize across different translation tasks.
A.5 Autoregressive Decoder only Transformer
The transformer consists of stacked blocks of self-attention, which itself consists of smaller units of self-attention heads that are concatenated before being fed through a fully connected layer. In autoregressive decoder-only transformers, training and inference adopts causal mask, where current positions are only able to attend to previous timesteps, instead of being able to attend to the entire input sequence. Unlike eecoder-decoder NMT models where source and target sentence have separate processing transformer blocks, decoder-only means that the same model weights are both used to "encode" the source sentence and "decode" the target sentence single continuous sequence.
A.6 Training with Autoregressive Translation
The original language modeling objective in GPT training involves predicting the entire input token sequence which consists of both the source and target sentence (shifted by position). We found this to produce slightly worse results than only minimising the negative log likelihood of predicting the target sentence be translated, and not the entire sequence. We consider this autoregressive translation training.
A.7 Lo Attention Gate Training
Training Details For Section A.8, We train using Adam Optimizer (3) 0.9, B2 0.999) with a batch size of 32, and learning rate of 0.001, early stopping patience of 10 and threshold of 0.01. We initialise attention head gates to be instead of random or 0.5 as this leads to faster convergence. We experiment with two different training settings, the 0 0-prompts Train setting
16

---

Masking GPTN2.78 (en-pt)
Kasking BLOOMBB (en-pt)
Ex Exery
Exte-Quory
Instr ERtTEER
Ex.Dsery
Diny
Query
Bett Ounyyua
24
Layer
Layer
Masking LLAMA-3.1-68 context-fro (en-pt)
Masking 3. 8-8--NSRUTT context-from
MN
BLED
Extolgguyy Italy Extuiquery Expunquuryy
Ex Extrguurry 2oor Ex Eueury Exear
Instr, Ex, Quary
En, Auu
2is0r Ex, Excrrr
Figure Layer-wise Context-masking results on the English + Portugese translation task. Critically, see nearly identical trends to what see in Figure on the English to French translation task, suggesting our results generalize across language pairs and "tasks".
and the -prompts Train setting. As described in Section A.6, we train the model by predicting only the target sentence, conditioned on the context. In the O-prompt setting, the context consists of the instructions and the source sentence to be translated. In the 5-prompt setting, the context consists of the instructions, 5 prompt examples, and the source sentence to be translated.
In the 0-prompt setting, the conditional prefix consists of the instructions and the source sentence to be translated. the 5-prompt setting, the conditional prefix consists of the instruction, source target sentence pairs, and the source sentence to be translated.
Data We used the first 10,000 lines of en + fr from WMT06 Europarl [37] for training. 7 To test the generalisability of trained attention head gates, we use a different test domain, FLORES [29] to reflect the scarcity of in-domain data. We also test an additional language direction en + in FLORES see training can generalise across languages.
Training Details We train using Adam Optimizer (31 0.9, 82 0.999) with batch size of 32, and learning rate of 0.001. We use a large early stopping patience of 10 and threshold of 0.01, and train for up to 100 epochs. This is due the nature of Lo training; improve over many iterations and would like the attention gates to keep training as long as there is no large loss in performance. We initialise attention head gates be instead of random 0.5 as this leads to much faster convergence and better performance. For the regularisation weight X, we search over hyperparameter set of (0.1, 0.01, 0.001, 0.0001) and found 0.01 performs best on the validation set.
do not expect performance to
A.7.1 Lo head masking experiments.
Additional experiments on LO head masking in the en- + direction.
Data available from https: hetpps:/www..o www.tt.ttt.....mmmmmmmmmmmmmm... arg/europarl/
17

---

Nesking
EUOCM 3B La Mesking
100
14
Attention Head
8
Attention
33 20 23 24232221
10
U4
04
90
0C
Layer
Layer
Figure 8: Visualisation of attention head masks for GPTNeo and BLOOM, learned with Lo 0.01) regularisation under 0 0-promppt train scheme en + fr. value of (in black) indicates that the attention head is effectively masked out by the trained attention gate. Around 10% of attention heads are masked out i.., redundant, with a majority of them occuring at the later layers for GPTNco and distributed across layers for BLOOM. fr - en is availble in Section A.7.1
A.7.2 Training Attention Head Gates with Lo regularisation
For scalable approach to pruning, we opt to train self-atten head gates following [61] using the technique of differentiable Lo regularization [40]. Let the attention head gates g E g Xny be set of trainable parameters, where n& is the number of attention heads per layer, and is the number of layers. Let the original output of each attention head be Uj. gated outputs Ūy are obtained by elementwise multiplication of the gate value gj. L.e., úÚ Oj @ U.. For {(x: y) }) source sentence (z) and target sentence (v) training pairs, a model f and loss function L, Lp regularisation adds a weighted penalty associated with the complexity of the parameters. The Lo loss is non-differentiable it involves raw counts of parameters. As work around. g be approximated with random variables drawn from a Binary concrete distribution [43. 35). We refer the reader to [40] for the relevant technical exposition. Details of training provided in Section A.7.
A.7.3 Generalisability of Lo gate training
We experiment with 0-prompts s and 5-promptss in training and using 0 (no regularisation) and 0.01. Lo training for the 0 0-prompt shows some gains for the O-prompts test case, and with no loss the 5-prompts test (Table 2). Notably, this persists in en pt, different language direction from training.
The robustness of translation performance under multiple testing conditions (number of prompts, datasets, language directions) gives some confidence that the trained discrete attention head gates from Lo support general ability translate (Table 2). In contrast, the soft attention head gates without regularisation (A 0) appear to overfit as they perform well on some conditions but deteriorate in others.
observe that 0 0-prompt training for Lo(A 0.01) also outperforms 5 5-prompts which slightly suprising since 5 5-prompts has more information in the prefix to locate the translation task. One possibility is that the model overfit to the Europarl domain where the training prompts were drawn from.
A.8 Studying Redundancy via Compression
(Note: This Appendix section is based on 1-BM models).
LL regularisation has the effect of reducing the magnitude of all Li regularisation has the effect of reducing the magnitude of several attention heads to very small value (but not exactly 0), while Lo regularisation has the effect of driving values to exactly 9The class of Concrete distributions was invented to work around the problem of automatic differentiation of stochastic computation graphs.
18

---

Base
0-prompt.s
5-prompts
Base
D-prompta
5-prompt.
A=0
A=
X=0
A=.01
==
A=.01
A=0
A=.01
O-prompts
18.3
21.4
20.1
18.9
19.3
6.7
15.7
13.2
6.4
5-prompts
24.3
24.5
24.1
23.6
24.2
25.9
19.6
25.8
24.3
26.0
Train: en ffr. Test: en fr
Train: fr, Test: >tpt
Table 2: Performance when using trained attention head gates for La with regularisation À = .01. = refers to training without regularisation. and prompts were used in the context for training. We highlight values which are greater or worse than 0.5 BLEU points from baseline. Note that as these are compression experiments, we do not expect Lo to perform better than baseline.
Pant warvee lart's
Nadrr 1,00P03 intrrete lacht
LLANST abaian--
LLANGO3
Figure 9: Layer-wise masking of self-attention heads for GPTNEO2.7 7B, BLOOM3B, LLAMA and LLAMA-CHAT on en + fr. The orange and blue dotted lines refer to the baselines (without masking) of O and 5 prompts with instructions. True/False refers to whether there are instructions provided (True) vs not provided (False). We observe critical layers near the middle and redundant layers towards the end of the model.
To what extent there specialised attention heads for MT in the GPT-style models? If there were specialised heads, we would expect the model to be highly compressable/punnble to select few heads. We plot grid map of learned attention gate values for en + fr, where O indicates that the head is masked out (Figure 8). We find that most of the masked heads are distributed at the later layers for GPTNeo and are distributed across layers for BLOOM.
A.9 Characterising Redundancy in Layers
Recently, [54] found that many layers in pre-trained transformers can be dropped with little harm to downstream tasks; moreover, it is well known neural MT transformer models are known have several redundant heads which are not necessary during test time [61, 44, 7]. However, it not clear if the same trends hold for in-context MT models, and how that redundancy is related to task location versus task execution. We focus on the task of en -> fr in this set of experiments.
We study the contributions of individual attention-layers by performing a simple layer-wise masking of all self-attention heads for single layer. When mask layer j. are masking the attention mechanism of layer j. that is the MLP of layer j acts directly on the output of layer j 1, rather than the output of the attention-head of layer j. Doing so allows us study how critical each layer is, where critical layers is loosely defined as those that have a large negative impact when masked.
We plot results for each layer all models, using three combinations of I0 examples, no instructions), (5 examples, instructions ), (5 examples, no instructions) in Figure 9."
In Section 4, we observed that there are layers for task location. In this section, we observe evidence that there are critical layers which correspond to the task locating layers, providing support for our earlier observations.
For instance for LLAMA2-7B en + fr, even in the scenarios when examples are provided, we can see a dip in performance around layer 15 to 18. Refering back to Figure 2, we see that this is where most of the task location with large jumps in performance had occurred.
Thee combination of I0 examples, no instructions] | is not meaningful as the model only receives "Q: <source sentence> A:" as the input and not expected to do the translation task.
19

---

For GPTNeo, we obseve a large set of contiguous layers which significantly decrease performance around layer 10 15. This also corresponds where most of the task location (large jumps in performance) had occurred for this model in Figure 2.
We note that the critical layers in different models have varying degrees of severity. 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 BLOOM is also around the same size as GPTNEO and performs more similarly to LLAMA. We suspect that it could be due training data or some other factor related the training dynamics but leave this for future work.
With regard to redundancy, we find that layers can be more safely removed towards the end without a noticeable loss in performance. We observe that for the less stable models, the model achieves close to baseline performance by layer-wise masking from f15 for GPTNEO, f26 for BI.OOM and 620 for LLAMA. This suggests that these later layers contain redundancy for translation.
20

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We outline our contributions precisely in our abstract and introduction.
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 in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long 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?
Answer: [Yes]
Justification:
We discuss our limitations 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 separate "Limitations" section in 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 how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on 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 the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to 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 address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an 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 Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
21

---

Justification:
We do not present theoretical results in this work.
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 the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to 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 supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We describe hyperparameters, models, and methods clearly. Our models are downloaded from Huggingface and are all opensource.
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, No answer 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 is 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 NeurIPS 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) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a 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 are welcome deseribe 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 to faithfully reproduce the main experimental results, as described in supplemental material?
22

---

Answer: [Yes]
Justification:
We will provide access to all of the code used to run experiments.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https: htpps://nipp...// cc/
pub public/ guides Cee//CodSSSSuaa s ssionPo ConPollcyy olicy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" 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 (ht tps: / //nips... cc/ public iillic//ae gu/guuioo des CodeS ubmi ss i onPo MaPPolccy icy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If 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 the paper) is 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 use publicly available dataset splits and our code will detail precisely what samples are used during in in-context learning.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of 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, or 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: [Yes]
Justification: Our core results present error bars over different random seeds and sets of in-context exemplars when applicable.
Guidelines:
The answer NA means that the paper does not 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 capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall 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).
23

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
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 emror 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 in the text.
8. Experiments Compute Resources
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?
Question:
Answer: [Yes]
Justification: We provide sufficient information to reproduce all of our experiments (all of our experiments can be run consume-grade GPUs such as RTX 6000).
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 of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't '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 https: /nwuriuuuup..eree httppubl.c/Ethccciiieeellel/
Answer: [Yes]
Justification:
We adhere to the NeurIPS code of ethics.
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 deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a 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: [Yes]
Justification: Our work contains an Impact Statement the Appendix.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment 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 to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If 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 system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: NA
Guidelines:
The answer NA means that the paper poses no 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 to 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 we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: 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?
Answer: [Yes]
Justification: All of our data and models are publicly available, and we reference their sources in the main paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.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.
assets are released, the license, copyright information, and terms of use the package should be provided. For popular datasets, papersw's thoode com/ dat aset has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
25

---

If this information is not available online, the authors are encouraged to reach out to the asset's '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: NA
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as 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.
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: [NA]
Justification: NA
Guidelines:
The answer NA means that the paper does not involve crowdsoureing 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 Ethies, 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 or Equivalent for Research with Human Subjects
Question: 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?
Answer: [NA]
Justification: N/A
Guidelines:
The answer NA means that the paper does involve crowdsourcing 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.
26


Reviews:
[{"review_id": "HqENZBmqMF", "reviewer": "Reviewer_iAtv", "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).\n\nThe 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.\n\nThe 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. \nThe 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.", "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? \n2. Your study focuses on machine translation and code generation. How confident are you that these findings would generalize to other types of NLP tasks? \n3. 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": "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.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The authors introduce a \"layer-from context-masking\" technique to probe the internal workings of large language models during in-context learning tasks. This method offers a fresh perspective on how these models process and utilize contextual information, providing insights that were not previously available. \n- The identification of a \"task recognition point\" and the characterization of a three-phase process in task recognition and execution are important.\n- The finding that attention to context can be removed after certain layers without significant performance degradation has direct practical applications for improving inference efficiency in LLMs, potentially leading to substantial computational savings.", "weaknesses": "- 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.\n- 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.\n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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)\n\n> 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.\n\nUnfortunately, 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. \n\n> 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.\n\nWe would like to discuss with the reviewer on the suggested confounding factors.\n\n* 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.\n\n* Specific prompt format. There are two main prompt formats used in the paper, reflecting the task. \n\n1a. “Translate from {L1} to {L2}: Q: {source_sentence} A:, \n1b. “Translate from {L1} to {L2}: Q: {source_sentence} A:,\n\n\n2. Write a function to find the longest chain which can be formed from the given set of pairs. Q: {program_description}, A:\n\nWe hope this gives the reviewer some confidence that there is sufficient coverage of the confounding factors highlighted. \n\nInstead, the potential confounding factors that we have investigated more pertinent to the research question are \n* With/without instructions\n* Number of In-context Examples\n* Different models Families\n* Instruction Tuned VS Not Instruction Tuned Models\n* 2 Different Tasks\n\n> 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.\n\nThank you for the suggestion, we will consider whether it’s possible to frame this better.\n\n> 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.\n\nThank 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\n\n> Limitations\n\nMost 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). \n\nCurrently in the paper, we have several \"future work\". i.e. limitations that we have not figured out.  \n* _“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.\n“_ (line 245:249)\n* _\"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)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KQDvrwi6Wn", "reviewer": "Reviewer_NuLE", "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. \n\nTo 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.", "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)?\n\nSec 4.3 could use better pointers. It's not clear which model is instruction tuned v/s not.\n\nSec 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": "NA", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "* This is an important area of research. The authors carve out a strong motivation, which is enhanced by clearly defined and executed experiments. \n* The paper is very well written, which makes it an interesting read. \n* The experiments are extremely thorough, some of them being: ex v/s Inst masking, MT v/s code, attention to context v/s input, # prompts analysis, attention heads study etc. I can tell that these took a lot of effort and will be of great importance to the community.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the deep appreciation of the work!\n\n>  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)?\n\nIt’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. \n\n> Sec 4.3 could use better pointers. It's not clear which model is instruction tuned v/s not.\n\nUnderstood, we will fix it.\n\n> 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?\n\nWe 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).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "m8pDuevuWy", "reviewer": "Reviewer_m9xg", "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.", "questions": "Line 117: can j be equal to I, i.e., (j<=i)?", "limitations": "See weakness.", "rating": 3, "confidence": 4, "soundness": 1, "presentation": 3, "contribution": 3, "strengths": "The motivation of this paper is clear, and the research question sounds interesting.", "weaknesses": "However, I have several concerns about the work.\n1. 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. \n2. 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.\n\nOverall, I think this paper tries to explore an important research question, and I would like to encourage the authors to address my two concerns.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their time and comments on our work. Regarding the listed weaknesses and questions:\n\n> Layer-wise masking. The authors only release first j bottom layers and conclude that bottom layers are important.\n\n**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) ”. \n\nWe 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. \n\n> If the authors mask layers with a reversed order, I suppose the conclusion would change to that top layers are important.\n\nThe reviewer may be referring to _layer-from-context masking_ (Figure 2,3), \n\nThe masking order that we have used, retains the original order of processing that the transformer sees during training. \n\n\"All masks operate from the j-th layer ($\\ell_j$ ) onwards, i.e. masking from $\\ell_{20}$ means causally masking\nout attention to all context positions from $\\ell_{20:n_\\ell}$, where $n_\\ell$ is the total number of layers. To construct\nFig 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)\n\nThe 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.\n\n> pretrained LLMs solve task with multiple modules. And different layers may share similar functions for solving that task.\n\nWe 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.\n\n\n\n> Token masking. The three token masking strategies are not comprehensive.\n\nTo recap, the three token masking strategies are\n* Has no instructions, and masks examples ($\\overline{\\texttt{Ex}}^{Mask}$)\n* Has instructions and masks examples ($\\texttt{Instr}\\overline{\\texttt{Ex}}^{Mask}$)\n* Both instructions and examples masked ($\\overline{\\texttt{Instr}}\\overline{\\texttt{Ex}}^{Mask}$) -- “dropping all of them”\n\nWe described the significance of $\\overline{\\texttt{Instr}}\\overline{\\texttt{Ex}}^{Mask}$ i.e., “dropping all of them” in Section 4.1, Layer-from Context masking.  \n\n\"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)\n\nTo 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).\n\n\n> many existing works show that ICL examples affect the performance much.\n\n 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. \n\n> Line 117: can j be equal to I, i.e., (j<=i)?\n\nYes, thank you for catching this typo!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HC132ry9j4", "reviewer": "Reviewer_z6j9", "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.", "questions": "Please refer to the weaknesses section.", "limitations": "--", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "The experimental design is well-conceived and includes an analysis that elucidates the internal mechanisms of in-context learning.", "weaknesses": "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.\n\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their time and comments on our work. Regarding the listed weaknesses of our paper:\n\n> \"The practical utility of the findings is questionable.\"\n\n**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. \n\n> doubt whether the proposed improvements to inference efficiency in Section 5 can be applied in practice.\n\nWhile 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.\n\n\n> 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.\n\n**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. \n\nHowever, 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept (poster)