How does representation impact in-context learning: An exploration on a synthetic task

Thank you for your valuable comments and suggestions and the encouraging words. We appreciate the time and effort you took to review our paper. Please find below our responses to your questions and the changes we have made to address them.

Q1,Q2. 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.

Q3, 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.

Q4, Yes. We give a experiment on a new NLP dataset to demonstrate that conclusion obtained in synthetic tasks can be extended to NLP task in Appendix C.2.

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

Q6, We regret that our writing makes you misunderstand here. Our intention in using PCA was to elucidate why we named them the "in-weights component" and "in-context component". However, we did not mean to imply that we are specifically using PCA. To avoid this ambiguity, we have revised the paper.

Q7, 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.

Q8, 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.

Q9, We will revise these typos.

[1] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.11895, 2022.

Thank you for your valuable comments and suggestions. We appreciate your feedback and have made several changes to address your concerns. Please find below our responses to each point you raised.

Weakness 1

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.

Weakness 2

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.

Question

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

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.

Minor edits

We fix the issues. We will continue to check potential grammar issues.

Point 1: please note that there are several typos and grammar errors in this text: We will fix the grammar errors.

Point 2: I feel that this proof, in its current state, is still difficult for me to understand. Please respect our time for this rebuttal. There exist two elements that influence the comprehensibility of a proof: the complexity of the issue at hand and the quality of the proof's composition. If the former is the case, it cannot serve as a valid reason to dismiss our paper. However, if it's the latter, please provide us with detailed feedback rather than a mere "However, I feel that this proof, in its current state, is still difficult for me to understand". This behavior is irresponsible.

Point 3: I see that most reviewers share my concern that the finding — that a good in-weight component is required for in-context learning — is more or less expected. Whether it holds for larger models that are strong in both aspects still seems to be a matter of debate. The most reviewers concern about the experiments about D_fix->D_rnd and "in-context component is related to in-context learning" is within expectation, instead of the conclusion you mentioned that "that a good in-weight component is required for in-context learning". You also don't raise the concern in your first review. What's more, "that a good in-weight component is required for in-context learning — is more or less expected" and "Whether it holds for larger models that are strong in both aspects still seems to be a matter of debate" are somewhat contradicted with each other. If it is within expectation, then why you questions about "Whether it holds for larger models that are strong in both aspects still seems to be a matter of debate"?

We appreciate the time and effort you took to review our paper. Please find below our responses to your concerns and the changes we have made to address them.

Weakness 1

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

weakness 2

We regret that our writing makes you misunderstand here. The role of D_rnd and D_fix setting is not to seperate the impacts of in-context and in-weights component. Instead, we use the probe method to seperate it. In the revision, we rewrite the Section 2 to make it clear. In the revision, we provided a formal definition of the in-weights and in-context component and an analysis of the D_fix and D_rnd settings. We think the revision can solve your concern. (Section 2.2 and the paragraph start with "analysis" in section 2.3 are directly related to this).

Weakness 3

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.

Question

Q1: In this paper, we use linear probe, (The probe model contains only one linear layer). One epoch is enough for adequate training. Training for more epochs doesn't helpful in our problem. And we find that more complicate probe models don't have obvious gain.

Q2: We disable the attention layers by replacing the attention layers in the Transformer with identity maps, whose outputs are equal to their inputs. We add this content in the paper.

Thank you for your valuable comments and suggestions. Please find below our responses to your concerns and the changes we have made to address them.

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.

Weakness 2.

We agree with you that there is a distribution shift. However, we disgree with your logic that ``Distribution shift->drop of performance". The reason is that

1. There are no drops of the performance when we swith from D_rnd to D_fix, although there are also "distribution shift" exsits.
2. Using D_{rnd and fix} can improve the results on D_rnd compared with only trained on D_rnd. D_{rnd and fix} also introduce the distribution shift. Our logic is that `` Distribution shift -> unaffected in-weights component and affected in-context component -> drop of performance". Our experiment reveal that training on D_fix and then evaluate on D_rnd, only the learning of the in-context component cannot be transfered. And due to this factor, the performance drops.

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.

Question

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

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

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

Q4, we rewrite the paper to fix the problem.

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

Q6, It it a typo. Actually it is D_rnd^fix => D_rnd. Thank you for your pointing out. we have fixed it.

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

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.

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.

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