Paper: How does representation impact in-context learning: An exploration on a synthetic task (ICLR-2024)

Reviewer comment (Writing):
How to understand "a good in-weight component" in Section 1.1. Could the paper provide more details? This is hard to understand from the current descriptions.

Author response:
Thank you for your pointing out. It is an important question. We rewrite the section 2.2 to solve these problems. In the revision of Section 2.2 , we give a formal definition about in-weights and in-context components, formal definition about what are good in-weights and in-context components and a decomposition of the problem.

Reviewer comment (Theory):
There is confusion about the approximation that uses a very simple network to estimate the large language models (LLMs). This is a bit incredible. If the approximation can be so good with respect to performance, why do we still increase the scale of LLMs? Would the approximation bring some disadvantages?

Author response:
We regret that our writing makes you misunderstand here. We are not asserting that a small network can approximate the Large Language Model (LLM). Our main points are: 1) The in-weights component plays a crucial role in in-context learning. This is evident as, under the perfect in-weights component assumption, even a simple network can exhibit proficient in-context learning. 2) We demonstrate that the neural network we've constructed can approximate the performance of the model trained on experimental data, thereby validating that our construction is meaningful.  
We revise the paper to reduce the possible of the misunderstanding.

Reviewer comment (Experiments):
In this paper, $i=40$ and $j=0$. The determination of such values is unclear.

Author response:
The selection is arbitrary. As stated in [1], the choice of i, j within a reasonable range will not affect the results.

Reviewer comment (Writing):
The paper could supplement some intuitive explanations about Definition 2 for better understanding.

Author response:
This definition is to ensure that the representation has a good in-weights component. (the definition of ``good in-weights component" is given in section 2) and the in-weights component can be easily readout using a linear layer.

Reviewer comment (Evaluation):
The limitation of this paper (which is also stated in Section 4) is worrying. Perhaps, due to the issue, the obtained conclusion is not general.

Author response:
We add the Section C in the Appendix. We give the following effort to try to solve your concern: 1) we analyze the previous discovery in NLP paper related to in-weights and in-context components. 2) we add a new experiment on a NLP dataset and our discoveries are held in this dataset.

Reviewer comment (Writing):
This paper looks like it was submitted in a hurry. There are many typos in the current form. For example,  
(1) It should use upper and lower quotation marks but not all lower quotation marks.  
(2) On the citation of the Adam optimizer.  
(3) In-weights comp. but not in-weights comp.  
(4) There should be a blank space between descent and (von Oswald et al., 2022). Similarly, between Transformer and (Vaswani et al., 2017), between this sequence and According, and between semantic” and i.e..  
These should be checked carefully for better representations.

Author response:
We will revise these typos.

Reviewer comment (Experiments):
Given how the different experiments and probes are set-up, the results look a bit obvious to me. For example, `D_fix` → `D_rand` results was quite trivial as during training time the model never saw such tasks. It would be interesting to understand within the same framework under what conditions does in-context learning not emerge. Are there cases where in-weight learning interferes with in-context learning and vice-versa? For example, what if you overtrain on `D_fix`/`D_rand` or under/over parameterise the model? What would your prediction be for this case?

Author response:
D_fix can be considered a specific instance of D_rand. Even in the "D_rand->D_rand" scenario, the test set includes tasks that the model has not previously encountered. This is because the number of potential mappings between factor values and labels is $n_e(n_e-1)(n_e-2)...(1)$, where n_e represents the total possible factor values of factor e. If n_e equals 10, then the number of potential mappings, $n_e(n_e-1)(n_e-2)...(1)$, is approximately 3.6*10^7. This figure greatly exceeds the number of sequences (10^5) present in the training set.

Reviewer comment (Experiments):
One of the strengths of the task is that it requires the model to do both in-context and in-weight learning to do well. Have the authors investigated the phenomenon in a purely language setting where this is also a requirement?

Author response:
Thanks for your suggestion. We add this in Section C (Appendix). In this section, we summary some prior practice works that evident both in-weights and in-context components are used in practice. And we give an experiment on a new NLP dataset to demonstrate that conclusion obtained in synthetic tasks can be extended to NLP task.

Reviewer comment (Reproducibility):
What is the loss function? Is it `cross-entropy loss`?

Author response:
Q1: yes, it is crossentropy. Thank you for your pointing out. We add the description in the paper.

Reviewer comment (Theory):
The proof by construction was a bit convoluted to understand, and quite a few notations were missing or unclear. At the moment, I am unsure if I fully understand it. Perhaps this was because of the lack of space or because it was written a bit at the end, but I would be happy to go through it again if expanded with intuitions.

Author response:
Q2： We have entirely rewritten the proof to enhance its readability (Section E in Appendix). To accomplish this, we've taken the following steps: 1) We've added a new paragraph to explain the notation, and 2) we've incorporated more explanations and intuitive reasoning throughout the proof.

Reviewer comment (Writing):
There are areas in the manuscript that require attention, such as citation error in Section 2.2.

Author response:
Thank for your pointing out. We have fixed it in the paper.  We will continue to check these issues.

Reviewer comment (Experiments):
The probe methodology requires further clarification. There is a lack of clarity regarding from which layer features are drawn and how the choice of layer affects the results. It is possible that some layers dominantly contain in-context components, while others are richer in in-weight components. Additionally, the rationale behind training each classifier for only one epoch, as opposed to ensuring thorough and adequate training, is not explained.

Author response:
**Point: There is a lack of clarity regarding from which layer features are drawn and how the choice of layer affects the results. It is possible that some layers dominantly contain in-context components, while others are richer in in-weight components.** We add the experiment in Section F (Appendix). In the experiment, we find that the in-context component concerntrate on the higher layer whether and in-weights component prefer lower layers. We also give the choosen layers.

**Point:  the rationale behind training each classifier for only one epoch, as opposed to ensuring thorough and adequate training, is not explained.** In this paper, we use linear probe, (The probe model contains only one linear layer). One epoch is enough for adequate training. And we find that more complicate probe model doesn't have obvious gain.

Reviewer comment (Theory):
A fundamental (somewhat unstated) assumption is that the in-context comp and in-weights comp are decomposed via this probing analysis and thus we can see the impact of both of them. However, the in-context component is clearly also impacted by concepts learned in the weights. It also already seems obvious that one would need to be able to achieve a high in-context comp score in order to achieve high in-context accuracy. For example, I fail to see what we learn regarding the second major claim ‘in-weights comp plays a crucial role in learning the in-context comp’. I am confused about (1) why this is surprising (2) if it is surprising, which of the experiments justifies this. I could see 6A potentially supporting this, but there are other factors at play. In 5, one trains from scratch so clearly there is some ‘warm up’ that is needed to get in-context learning ability.

Author response:
### Weakness 1.  
**Point: A fundamental (somewhat unstated) assumption is that the in-context comp and in-weights comp are decomposed via this probing analysis. However, the in-context component is clearly also impacted by concepts learned in the weights.**
In the revised version, we provide a formal definition of the problem along with the assumptions in Section 2.2.

**“Point:It also already seems obvious that one would need to be able to achieve a high in-context comp score in order to achieve high in-context accuracy.**
I understand the reviewer's logic in associating "in-context comp" and "in-context accuracy" with in-context learning due to their terminology and their relation to in-context examples. However, two points might have been overlooked. 1) The existence of in-context comp isn't natural; it's a result of our effort to decompose it from the representation. 2) Furthermore, our findings suggest a portion of the representation is directly tied to performance. This connection hasn't been explored in previous work, and this discovery suggests that enhancing representation learning may be a fundamental approach to improve in-context learning.


**Point: I fail to see what we learn regarding the second major claim ‘in-weights comp plays a crucial role in learning the in-context comp’**
1) Regarding the statement, "I could see 6A potentially supporting this, but there are other factors at play. In 5, one trains from scratch so clearly there is some ‘warm up’ that is needed to get in-context learning ability." We've supplemented this with a new experiment (Figure 9A in Appendix F), where we train D_fix for a brief period and still observe a significant improvement.
2) The use of both D_{fix} and D_{rnd} leads to considerably better results compared to utilizing D_{rnd} alone.

Reviewer comment (Evaluation):
Given this, perhaps the most surprising result is that the in-weights scores on 5A with D_rnd are poor and 6A does not improve to what is achievable training full on D_rnd, but only the in-context score improves. I’m not sure what to make of this though. For instance, this could be explained by saying that it just may not be important to form a representation that is predictive of the factor values without the context for this particular task since the task can be inferred from the context. Perhaps the authors can shed some light on that.

Author response:
### Weakness 3. 
**Point: Given this, perhaps the most surprising result is that the in-weights score on 5A with D_rnd are poor and 6A does not improve to what is achievable training full on D_rnd, but only the in-context score improves. I’m not sure what to make of this though.** 

Here, we give possible reasons for this obervation:

1) We don't have direct supervision signal for improving in-weights and in-context component. 

2) Under D_rnd setting, the model can achieve smaller loss by just improving the in-context component. Therefore, the model don't have motivation to improve the in-weights component and it even sacrifice in-weights component to obtain better in-context component.

Reviewer comment (Writing):
It would be helpful to more formally define the in-context and in-weights components earlier in the paper. Throughout the intro and early parts, it is clear what th paper wants to do, but it’s difficult to discern what quantifiably one wants to measure since ‘in-weights components’ and ‘in-context components’ are sort of nebulous terms until one reaches section 2.3. Honestly, I don’t know the best way of presenting it, but I just want to raise this as a potential issue for future readers in case the authors have alternatives.

Author response:
Q1 We rewrite the paper and give the formal definition in Section 2.2

Reviewer comment (Reproducibility):
When you compute the in-context comp score, are you using the very last element of the sequence at L?

Author response:
Q2 Yes, we use the last element of the sequence at L

Reviewer comment (Experiments):
When producing the probes, are you doing a train and test split and are we looking at the test accuracy. How is this split being done? Are we looking at the test score when we look at the plots?

Author response:
Q3 Yes, we look at the test score. We add a description for the dataset split of probe model in Appendix A.4.

Reviewer comment (Writing):
FYI I don’t think E is defined before being used in Section 2.3.

Author response:
Q4, we rewrite the paper to fix the problem.

Reviewer comment (Presentation):
Fig 6A should be D_fix->rnd not D_rnd->fix I think.

Author response:
Q5, Yes, it is a typo. It is D_rnd->D_fix

Reviewer comment (Experiments):
In, D_rnd^fix, do you use twice the amount of data as D_rnd? That is, does each epoch contain twice the data or are they both halved so the total is the same?

Author response:
Q7，We use same amout of data. We add a description for the dataset split of probe model in Appendix A.4.

Reviewer comment (Writing):
I don’t really understand the last paragraph that starts with “Comparison with previous work in analyzing the in-context learning mechanism.”

Author response:
Q8, this sentence is related to the proof details of the proposition. I move it to the Appendix for the audiences who have interesting in the proof.

Reviewer comment (Evaluation):
It would be helpful to compare 6A directly with 5B since they are plotting the same kind of results. It seems that in both you get good performance starting at roughly the same epochs. This I think supports my question about how much of this is just attributed to distribution shift: once you start training on the same data as you test on, performance improves.

Author response:
Q9, We give a directly comparison about the some results on Figure 9A (Appendix). Thanks for your advise, we will rethink it and try to find a better way for representation.

Reviewer comment (Writing):
Some typos: “copy-past” “fead forward” “How transformer solve this problem?” “analysis is to analysis”

Author response:
Q10, Thank you for your point out. We have fixed it in the paper.