In-Context Learning in Large Language Models: A Neuroscience-inspired Analysis of Representations

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

Significance of our contribution: Regarding the final point about RSA having been applied in various NLP contexts previously, we do acknowledge that it's not a wholly novel technique in NLP. But as we discuss in the related works section, this prior work has focused on small-scale models. To our knowledge, our work is the first to demonstrate the viability of leveraging RSA to analyze emergent behaviors in large state-of-the-art language models. Given the rapidly increasing scale and complexity of LLMs, developing interpretability tools that scale efficiently is an important open challenge. By scaling up RSA to analyze a 70B parameter model, we aim to demonstrate RSA's promise as a tool for characterizing complex behaviors in large modern LLMs. And of course, no one else has specifically investigated ICL with RSA.

Please also refer to these new figures we created to visually describe our main contribution:

• Overview of methods: https://figshare.com/s/1b8983f60c9f81075c32
• Task and experiment design: https://figshare.com/s/e96a108f58e1ed381619
• Attention ratio analysis (ARA): https://figshare.com/s/9541214a2a8fbd37ba83

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.

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.

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.

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.

Continued in part 2.

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.

Applications of Attention Ration Analysis: The approach of analyzing attention ratios may have applicability beyond question and answer. It can be applied to identify potential hallucinations for any continuation assuming the model’s response could be decomposed into a claim and support structure. If a strong correlation exists, as the one found in our experiment, this threshold could serve as a measure to assist with prompt engineering. This would be done by finding phrasings of a question which maintained the threshold within a certain range based on observed values and their relationship to a desired property of the output.

Thank you once again for your thoughtful review and detailed comments. We hope that our response addresses your concerns.

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.

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.

Presentation issues:

• 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.

• 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.

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.

Thank you once again for your thoughtful review and detailed comments. We hope that our response addresses your concerns.

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.

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.

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.

Thank you once again for your thoughtful review and detailed comments. We hope that our response addresses your concerns.

We thank the reviewer for their comments about our design decisions. We have added new experiments to the Appendix Sections C2 (Figures 17 and 18) and B2 (Figure 9) to address some of the concerns and we clarify others here.

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.

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.

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.

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.

Continued in part 2.

We thank the reviewer for appreciating the premise of our work, that latent layers of LLMs should show changes with ICL that are reflected in behavior and allowing us to elaborate on our specific contributions. Our contributions are threefold:

1. We carefully design tasks that allow us to investigate what the structure of latent representations should be (hypothesis driven design). This allows us to ask questions like what the reviewer suggests “What do we expect to have happened to those representations and attention if the behavior has improved?”. As such, the question here is not whether it is “possible for the representations and attention to not reflect that improvement?” as the reviewer asks, but how to measure what these representational and attentional changes should look like, and what methods enable us to measure the changes.

2. We propose using a tried and tested method from neuroscience, RSA and correlation with a hypothesis matrix, to test the extent to which a representational structure (measured by RSA) becomes more similar to a hypothesized structure (correlation with a hypothesis matrix) following ICL. We can then measure the extent to which gradual changes in this correspondence lead to changes in behavioral error. As with neuroscience studies, it is not enough to show changes in representation but quantifiably how ICL changes representation, and how these changes correspond with changes in behavior. Another significance of this approach is that it allows us to decode answers to certain prompts directly from the representations, even if they were not reflected in behavior. This is an impactful finding, allowing future uses of open LLMs, in which some information can be directly captured from the layers.

3. We propose ARA (Attention Ratio Analysis), a novel method to measure how attention changes following ICL. As above, we also investigate how changes in attention following ICL correspond with changes in behavior

Layer choice: 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.

Details of the models: We chose two open-source decoder models from HuggingFace. Both models leveraged instruction fine tuning.

1. Vicuna v1.3 is fine-tuned from LLaMA with supervised instruction fine-tuning. The training data is around 125K conversations collected from ShareGPT.com.

2. Llama 2 was pretrained on 2 trillion tokens of data from publicly available sources. The fine-tuning data includes publicly available instruction datasets, as well as over one million new human-annotated examples.

Typos and Latex errors: We apologize for our errors and thank the reviewer for bringing them to our attention. We have fixed all typos Latex errors.

Definition of “correct” answer in the persona injection task: For evaluating the persona injection task, all results measure the model's ability to generate the truthful response. When asked to deceive the expectation is that the model should follow the earlier instructions to tell the truth and ignore instructions to the contrary.

The role of RSA: RSA allows us to directly compare the representational geometry to a hypothesized model space. Specifically, RSA involves constructing representational similarity matrices that capture the pairwise similarities of instances in a representation, and then comparing these RSMs to an a priori similarity matrix that encodes some hypothesized relational structure. So, one of the notable benefits of RSA is to provide a simple and intuitive way to test correspondence to relational hypotheses. The reviewer's final point about probing sidesteps the various issues with parametric approaches mentioned in the paper. Although we use probing as well in the paper, we would like to emphasize RSA's ability to provide a holistic characterization of the representational geometry changes, rather than just improving decodability of one variable.

Thank you once again for your thoughtful review and detailed comments. We hope that our response addresses your concerns.

Please also refer to these new figures we created to visually describe our main contribution:

• Overview of methods: https://figshare.com/s/1b8983f60c9f81075c32
• Task and experiment design: https://figshare.com/s/e96a108f58e1ed381619
• Attention ratio analysis (ARA): https://figshare.com/s/9541214a2a8fbd37ba83