Paper: In-Context Learning in Large Language Models: A Neuroscience-inspired Analysis of Representations (ICLR-2024)

Reviewer comment (Novelty):
Some rationale for using RSA was a bit shaky. In its introduction, RSA is touted as overcoming parametrized probing methods, but this is a bit of a ‘red herring’, as there are already (increasingly many) varieties of parameter-free probing methods (e.g., https://aclanthology.org/2020.acl-main.383.pdf). That’s not to say that the parameter-free nature of RSA isn’t a positive, but it doesn’t really highlight why we should use RSA (or any neuroscience-based approach) in particular.

Author response:
**Benefits of RSA:** We agree that the parameter-free nature of RSA is not its only advantage, and that other parameter-free techniques exist. But RSA provides unique benefits beyond being parameter-free that make it well-suited as a tool for our goals in the paper. RSA allows direct comparison of the model's latent representational geometry to hypothesized relational structures specified a priori. This can be thought of as representational alignment between the representations and the ideal structure of the task. While probing can assess whether certain information is encoded in the representations, RSA provides a more holistic view of how well the entire relational structure of embeddings aligns with what an appropriate task representation should contain. In a nutshell, we think some of the key advantages of RSA are: 

- Allowing direct comparison of model representational geometry to hypothesized structures
- Providing a holistic characterization of representational geometry 
- Avoiding complex modeling assumptions / risk of spurious correlations associated with parametric probing

Reviewer comment (Presentation):
The text in the figures are fairly unreadable, due to size, but perhaps this is unavoidable. Some of their important parameters are not discussed (e.g., Fig 3 is merely described as ‘behavioral error’ but the nature of the plots are not given).

Author response:
**Definition of behavioral error:** We have included the details in figure captions. For example, in Fig 3, the caption describes the y axes as “absolute error for regression task” which is defined in equation (1). We have further clarified this in all figure captions, so it is easier to notice. We titled our y-axes as “behavioral error” so emphasize that in each task, we study behavioral changes that correlate with representation changes due to ICL.

Reviewer comment (Writing):
Can you please run a spell-check over your paper (e.g., ‘anayzing’, 'appraoch’, etc) and also check the grammar (e.g., “RSA avoid this risk’, ‘the models behavior’, ‘to reduces the effect’) and LaTeX errors (like ?? instead of references)?

Author response:
**Incomplete references and typos:** We apologize for our errors and thank the reviewer for bringing them to our attention. We have fixed all typos and Latex errors, and updated our references to include all information.

Reviewer comment (Reproducibility):
Sec 2.4 — why was only the final layer of the LLM considered? How are the subsets a, b, and c determined? Why is ‘a’ not included in the computation of A(a,b,c)_x? Does this only apply to question-answering, as the second sentence of the second paragraph suggests?

Author response:
**The role of ‘a’ in $A(a,b,c)_x$:** The input parameter $a$ which represents a substring of $x$ is referenced by the index $i$ in $A_i,j$ and $A_i,k$. As described in Section 2.4, $i$ indicates token indexes of substring $a$ after tokenization of string $x$. This means that $a$ is used in summations for both the numerator and denominator when adding the attention values for all pairs of tokens between $a$ and $b$ as well as $a$ and $c$. 

**Determining the prompt substrings $a$, $b$, and $c$:** We designed the reading comprehension tasks so that these substrings are well defined. Each prompt has informative and distracting components that clearly define these substrings. The last paragraph of section 3 describes these components for the reading comprehension task. 

**Layer choices:** We agree with the reviewer that investigating multiple layers and changes in representations across them is highly interesting. We experimented with the first, middle, and last layer before submission, and have now added more experiments to Appendix Sections C2 (Figures 17 and 18) and B2 (Figure 9) to show results for other layers of Llama-2 and Vicuna. Our results show that the middle layer and the last two layers demonstrate the same representational changes aligned with our hypothesis in the linear regression and persona injection tasks. The first layer of the network, however, does not show the same pattern. This is expected because the earlier layers in language models are known to reflect syntactical information as opposed to high-level task information. This observation applies to both models. We are happy to include more such experiments if the reviewer sees fit.

Reviewer comment (Experiments):
Sec 5 Why did you precede the instructions to be deceitful with instructions to be truthful? Would a fourth condition jump straight to instructions to be deceitful?

Author response:
**Persona injection prompt design:** We acknowledge the reviewer’s point that a prompt which directly instructs the model to lie is the most straightforward and would be the ideal experiment design. We initially conducted this experiment; however, Llama2 would output an end-of-sequence token immediately which prevented analysis of the model’s response with relation to the input prompt. We therefore attempted a succinct prompt injection which mirrors a simple adversarial attack. The prompt provided was the most succinct and simplistic adversarial prompt we found which changed the model’s response for both Vicuna 1.3 and Llama2.

Reviewer comment (Evaluation):
Perhaps my primary concern with this work is that it seems to focus quite a lot on presenting measures of behavioral changes due to ICL instead of representational changes. I believe the behavioral benefits of ICL have already been established (if not, please indicate what is new about these findings). Therefore, the exciting part of this work is the link between representational metrics and changes in inputs (and perhaps links between representational and behavioral changes). Many of the main figures, however, focus only on behavioral changes (Figure 3, Figure 4 a and b). Some figures are great (e.g., Figure 4 c), but I wish more were like that.

Author response:
We thank the reviewer for allowing us to clarify our work. Our focus is indeed on a) **measuring changes in representation** following ICL (as well as changes in alignment of representations with what would be expected of the task structures), and b) **measuring and quantifying changes in attention** following ICL, and c) **quantifying** how these representation and attention changes correlate with behavioral changes. In other words, our goal is to ground behavioral changes in representational changes. The reviewer is correct that having behavioral tasks that show an effect with ICL is not a novel contribution in itself. However, our specific focus requires tasks that satisfy specific design principles: have components that enabled us to hypothesize about the representational similarity structures, display similar behavioral changes for both Llama 70B and vicuna13B. **Figures 1c, 1d, all of figures 2, 4, 5 and figures 6 c-h** focus on representational changes. For example, figures 4a and 4b show the accuracy of a probing classifier trained on the models’ representation to predict the slope of the line from the prompt embeddings but from your comment it sounds like you have interpreted that as a behavioral result. We hope this improves clarity and are happy to update the paper accordingly. 

Similarly to neuroscience, studying representations requires designing tasks with specific representational relations across different components. Thus, coming up with careful hypothesis-driven tasks was a crucial step: we had to carefully design tasks that have meaningful components, and demonstrate behavioral improvements due to ICL in these two specific models. With these tasks, we can study representational changes related to task components and how they correlate with behavioral changes.

Reviewer comment (Theory):
I am somewhat hesitant to use attention values in interpretability work (see "Attention is not Explanation" by Jain and Wallace). Can the authors justify using attention measures to understand how the model is behaving?

Author response:
**Attention is not explanation:** We use attention in addition to model embeddings to measure representational changes, based on the hypothesis that token-level representational information may be lost in our choice of aggregation of token embeddings. Examining attention allows us to observe if there is a change in attention between specific parts of the prompt, and we do not mean to imply this is an explanation. For example, we would have reached similar conclusions even if the attention ratio defined in section 2 was always lower for truthful behavior than deceptive behavior. As stated in the paper in section 5, the attention values corresponding to different behaviors do exhibit a statistically significant difference in distributions.

Reviewer comment (Presentation):
Many of the violin plots used end up plotting uninterpretable extrapolations. For example, in Figure 3, by definition, no value can be less than 0, but the plots do go into negative values. The raw data are fine, but the plots seem confusing at first.

Author response:
- Violin plots: Thank you for pointing this out. All violin plots are fixed to not estimate negative values when representing non-zero distributions. We apologize if this caused any confusion.

Reviewer comment (Presentation):
Figure 5 a-c would really benefit from some sorts of labels to say what the axes are. I believe I understand this from examples in the appendix (e.g., Figure 7), but it should be very clear in the main paper.

Author response:
- Appendix Figure 12 and 5 a-c: We apologize for any confusion that the lack of labels may have caused. Including the prompts or even just the prompt IDs in the prompt similarity matrices would not yield a legible figure given the space constraints, but we have added clarification to the figure caption. To summarize, in this figure, each subfigure uses a distinct ordering of prompts aligned with the hypothesis they represent, so it is not visually perceptible that subfigure c is the average of the name and activity hypothesis.

Reviewer comment (Presentation):
Figure 12 in the appendix seems wrong to me, especially because there are no labels on the axes. In particular, c) is supposed to represent the average of a and b, but a and b look the same.

Author response:
- Appendix Figure 12 and 5 a-c: We apologize for any confusion that the lack of labels may have caused. Including the prompts or even just the prompt IDs in the prompt similarity matrices would not yield a legible figure given the space constraints, but we have added clarification to the figure caption. To summarize, in this figure, each subfigure uses a distinct ordering of prompts aligned with the hypothesis they represent, so it is not visually perceptible that subfigure c is the average of the name and activity hypothesis.

Reviewer comment (Reproducibility):
I only have a simple clarifying question: could the authors elaborate on how they "standardize the embeddings before computing pairwise cosine similarities" (section 2.1)?

Author response:
**Standardizing embeddings:** Similar to Timkey et. al (cited in section 2.1) we noticed that most prompt embeddings were dominated by one extremely large dimension pushing all similarities to high values. To mitigate the effect of such *rogue dimensions*, Timkey et al propose several methods including simply standardizing the embeddings before calculating cosine similarities.

Reviewer comment (Experiments):
In section 2.1, the way of composing \( n \) embeddings into one single embedding is to use max-pooling. Although max-pooling is a classical and popular method for representation composition, it is unclear how this method was selected and to what extent this selection impacts the final conclusion.

Author response:
**Embedding aggregation by max-pooling:** To address your question, we have included our analysis with mean pooling instead of max-pooling in the Appendix sections B2 (Figure 9), and C2 (Figures 17 and 18). 

In the regression task, mean pooling does not change the results for the middle and last layers, but the first layer’s representational correlation with our hypothesis seems to slightly increase with ICL when we aggregate the token embedding using mean-pooling instead of max-pooling. This observation applies to both models in both. 

In the persona injection task, mean-pooling, had a minimal impact on the correlation with the hypothesis matrices for all layers. All max pooled first layer token embeddings for 0-ICL produced the same encoding which led to a correlation of zero for all hypotheses. This is because a small subset of shared tokens, such as the trailing question mark and ‘what does’, collectively determined the elementwise max. We found that the middle layer and last layers had similar and high correlation for mean-pooled token embeddings of the prompt with both the activity and name hypothesis matrices. The correlation for both the middle and final layer increased when providing an ICL example for the name hypothesis for both models. Excluding the first layer, the inclusion of an ICL example improved correlation for a majority of model, layer, aggregation, and prompt type combinations. 

We are open to the reviewer’s suggestions in addition to this experiment.

Reviewer comment (Experiments):
In section 2.1, it says “the most relevant subset of tokens representing important components of the prompt are considered.” By manually (assumed) selecting the “most relevant” subset, it is uncertain whether the results can be used to represent LLM’s behavior.

Author response:
**Token choice for embeddings:** The tokens for embedding analyses were selected to represent the task portion of the prompt, excluding the ICL example tokens. In this way, the same tokens’ embeddings are considered from the prompt whether we are looking at a prompt with and or a prompt without ICL. It is not as manual a process since the question tokens are discoverable using the tokenizer and it is reasonable to assume that in any application scenario, the question/task tokens are distinguishable from the ICL example tokens. we have improved the clarity of our writing to more directly describe the token selection.

Reviewer comment (Experiments):
In section 2.2, it says “We designed a set of tasks with common components, such that representing those components is crucial in solving tasks.” If these components are crucial, more explanation would be valuable, such as how these components are designed and how we know that the combination of these components represents a way of solving tasks.

Author response:
**Task descriptions:** Each task is described in detail in its dedicated section (section 3, 4, and 5). For example, the reading comprehension task in section 3 is described to have informative and distracting components (sentences), and representing the informative sentence is crucial to solving the task because it contains the answer to the reading comprehension question. 

**Task design principles:** Similar to neuroscience, studying representations requires designing tasks with specific representational relations across different components. To identify such tasks, it is not sufficient to find tasks that show improvements with ICL, as what is needed are tasks with components that lead to representational similarities across different variations, about which we can make representational hypotheses. For instance, the reading comprehension task (Sections 3 and 5) allows us to vary *names* and *activities* separately and have hypotheses about the similarity of representations across multiple prompts with the same name, or the same activity. The same holds for multiple points that fall on the same line (Section 4), leading to the same slope. After successful ICL, one would expect improvements in *representational similarity or alignment with the hypothesis*. 

Thus, coming up with careful hypothesis-driven tasks was a crucial step: we had to carefully design tasks with meaningful components. This allowed us to have hypotheses about representations related to task components, continuous measures of improvement (e.g. similarity) after ICL, and how changes in representations and *changes in representational alignment with the hypothesis* correlate with behavioral changes following ICL.

Reviewer comment (Theory):
In section 2.4, the reason for focusing on the attention of the final layer is not explained.

Author response:
**Layer choice:**  In (https://aclanthology.org/P19-1356.pdf), Jawahar et al. used probing to demonstrate that earlier layers tended to capture syntactic information while later layers encode semantic information. Since we were interested in whether the model embeddings demonstrated a change in semantic encoding associated with the behavioral change, we focused on probing the last layer.

We agree with the reviewer that investigating multiple layers and changes in representations across them is highly interesting. **We experimented with the first, middle, and last layer before submission, and have now added more experiments to Appendix Sections C2 (Figures 17 and 18) and B2 (Figure 9)** to show results for other layers of Llama-2 and Vicuna. Our results show that the middle layer and the last two layers demonstrate the same representational changes aligned with our hypothesis in the linear regression and persona injection tasks. The first layer of the network, however, does not show the same pattern. This is expected because the earlier layers in language models are known to reflect syntactical information as opposed to high-level task information. This observation applies to both models. We are happy to include more such experiments if the reviewer sees fit.

Reviewer comment (Experiments):
There are some issues about experiment design, including the lack of clarity on why these three tasks were chosen, particularly the last one, instead of tasks where in-context learning has been widely used.

Author response:
Another constraint was identifying tasks that demonstrate an ICL-induced behavioral change for both models. We also wanted a set of tasks which covered both numeric and verbal reasoning. Given the vast discrepancy in model capacity (13B vs 70B parameters) and the desire to reduce noise in our experiment, we chose three tasks which were self-contained, relatively succinct, and elicited change in behavior with ICL. 

We are open to the reviewer’s suggestion for running other/more prevalent ICL tasks that can satisfy the above design criteria if they could kindly share their thoughts with us.

Reviewer comment (Writing):
The role of RSA in section 3 is not clear, and it is possible that something important has been missed there.

Author response:
**The role of RSA is section 3:** In section 3 we use probing classifiers and attention ratio analysis. RSA is not used in section 3, but the task in section 5 is based on the task in section 3 and there we use RSA with name and activity hypothesis matrices.

Reviewer comment (Evaluation):
In section 5, it is unclear whether the goal of this task is to study the robustness of the model or to understand how in-context learning is used for model inference. Although it is an interesting task, there is skepticism about the application scenario. The most important question is how this experiment can help demonstrate the impact of in-context learning.

Author response:
**ICL in section 5:** The intent of section 5 is to a. exhibit that providing a single ICL example improved both models’ abilities to respond truthfully (Figure 6 a-b) and b. to measure representational changes underlying the behavioral changes (Figure 6 c-h). We hope this clarifies how this section fits within the rest of the paper in that it does study the role of ICL in both behavioral and representational changes, and we are open to further suggestions by the reviewer.