ICLR: In-Context Learning of Representations

Authors: Core Francisco Park, Andrew Lee, Ekdeep Singh Lubana, Yongyi Yang, Maya Okawa, Kento Nishi, Martin Wattenberg, Hidenori Tanaka
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: In-Context Learning, Representational Geometry, World Models, Emergence, Percolation

Abstract:
Recent work demonstrates that structured patterns in pretraining data influence how representations of different concepts are organized in a large language model’s (LLM) internals, with such representations then driving downstream abilities. Given the open-ended nature of LLMs, e.g., their ability to in-context learn novel tasks, we ask whether models can flexibly alter their semantically grounded organization of concepts. Specifically, if we provide in-context exemplars wherein a concept plays a different role than what the pretraining data suggests, can models infer these novel semantics and reorganize representations in accordance with them? To answer this question, we define a toy “graph tracing” task wherein the nodes of the graph are referenced via concepts seen during training (e.g., apple, bird, etc.), and the connectivity of the graph is defined via some predefined structure (e.g., a square grid). Given exemplars that indicate traces of random walks on the graph, we analyze intermediate representations of the model and find that as the amount of context is scaled, there is a sudden re-organization of representations according to the graph’s structure. Further, we find that when reference concepts have correlations in their semantics (e.g., Monday, Tuesday, etc.), the context-specified graph structure is still present in the representations, but is unable to dominate the pretrained structure. To explain these results, we analogize our task to energy minimization for a predefined graph topology, which shows getting non-trivial performance on the task requires for the model to infer a connected component. Overall, our findings indicate context-size may be an underappreciated scaling axis that can flexibly re-organize model representations, unlocking novel capabilities.

Decision: accept (poster)
Reviewer scores: 8, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730627297424

thread_type:
official_review

<thread_object>
creator: reviewer_hdrb

<content>
summary:
This paper explores how in-context learning in large language models allows a reorganization of representations to fit novel task-specific semantics. The authors propose to override the pre-trained meaning of the words, thus the internal organization of the LLM representations with a semantic extracted from the neighbor relationships of names placed at the nodes of a graph. 

The authors show that the LLM can internalize the graph structure learned in-context for sufficiently long input sequence sampled from the graph. They also connect their findings to energy minimization in graph and percolation theories, offering insight into how models might adapt internal representations based on in-context structure.

strengths:
1. The paper presents a new task framework to study in-context reorganization in LLMs. It connects empirical observations with a theoretical foundation, connecting context-scaling with energy minimization. 

2. The graph tracing task using well-defined structures (grid, ring) is an effective method for observing changes in LLM representations in a controlled way.

3. The paper is well-written and relatively easy to follow; the main goal is clearly stated, and the set of experiments are spot-on

weaknesses:
The findings are based only on the analysis of llama3 representations, so extending the experiment to at least another model class would be easy. 

The constructed task is somewhat artificial, and it is not immediately clear how we can extend these findings in the domain of natural language.

questions:
Have you tried to see if your finding holds also on models other than Llama?

l371, 372 Have you tried to check empirically if the coordinates z2 and z3 are aligned with the principal axis you find with PCA?

What does DGP stand for in line 249?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730677101828

thread_type:
official_review

<thread_object>
creator: reviewer_ebvh

<content>
summary:
The paper studies how the different representations are organized in a large language model’s internals. It reveals the sudden reorganization of representations as the context length increases.

strengths:
1. The problem is interesting and important. 
2. The experiment results are interesting.

weaknesses:
1. Lack of quantitative metrics of “the representations mirror the grid/ring structure” in Section 2: It is unclear how the PCA plot relates to the conclusions that the representations reflect the graph structure. The results in Section 4 may be helpful but require more discussion.
2. The use of PCA: The PCA is useful for visualization but is not convincing enough to draw rigorous conclusions. Without the manually added vertices in Figures 1(c), 2(c), and 3(c), the results are not interpretable. 
3. Equ. (4) is the probability of seeing one node with each token randomly sampled with replacement. It does not match the data generating procedure with the first token randomly sampled, and the remaining tokens are generated through random walk.
4. I do not find the percolation theory fits here. Although Figure 8(b) indicates a strong result, it is unclear how y and x connect to the percolation theory.

questions:
1. What is the layer of the figure 1(c)
2. In Section 3.1, it is claimed that a window of $N_w=200$ preceding tokens is included in the computation of representations. This number is greater than the minimal context length of Figure 2(b), which is 100.
3. How does Figure 3(c) prove that the LLM learns the in-context non-semantic structure?
4. The percolation relates to the thresholds of having an infinite connected path. Can authors explain more explicitly about its relation to the problem considered in the paper?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730677683242

thread_type:
official_review

<thread_object>
creator: reviewer_a5tz

<content>
summary:
This work designs a synthetic task to determine to what extent LLM  representations reflect in-context semantics rather than those observed during pre-training. 

To do this they generate a graph structure, where nodes consist of tokens that are highly likely to be observed during pre-training, but with novel dependencies defined by the structure. They then generate sequences following a random walk across the graph, which are fed into the LLM 'in-context' (i.e. without weight updates). They then show that the LLM can learn to organise its representations such that they mimic the underlying graph structure.

The task is to some extent adversarial, because the context neighbourhoods defined by the graph differ from the standard semantics of the tokens.

Evidence presented is both qualitative (using PCA) and quantitative (various measures for how well representations conform to the graph structure).

strengths:
The paper is excellently written, and the figures clearly get the central findings across the reader, making it a pleasure to read. The experimental design is original, elegant and well crafted in order to probe the hypothesis. Finally, the findings provide an interesting insight in the abilities of LLMs to zeroshot generalise to novel connective structures.

weaknesses:
The work relies quite heavily on PCA, which the authors admit can be somewhat misleading, though they provide theoretical justifications for why this reliance is valid.

On a higher level I am not quite sure whether the finding is particularly surprising. Natural language is supposed to have similar causal relationships modelled by dependency parses, and LLMs are capable of modelling it successfully. Moreover, they are supposedly able to learn to translate hitherto unseen languages entirely in context ([1]), which constitutes a much more complex graph than the ones assessed in this paper. As a result I am not certain that the findings expose a new capability. However, as stated in the strengths section the paper provides a very nice framework for understanding this capability which should bias towards acceptance.

[1] https://arxiv.org/abs/2403.05530

questions:
Have you or do you plan to extend your investigation to more complicated structures? The ones presented in the paper are fairly simple and it would be interesting if there are certain forms that are harder to represent.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730985086379

thread_type:
official_review

<thread_object>
creator: reviewer_kjbf

<content>
summary:
The paper investigates how LLMs form representations for novel concepts during in-context learning. Authors design a synthetic graph navigation task where the model has to learn from random walks between graph nodes are referenced as words from the training set (e.g. “apple”), meaning that the model has to learn a new representation for these words to accurately predict the next node in the graph traversal. They refer to this problem as “in-context graph tracing”.

Authors then study the hidden representations formed by LLMs as they learn this task in-context, visualizing their first few principial components. They find that, after a certain number of examples, the model abruptly re-organizes the representation space in a way that facilitates the graph traversal task, as learned from the context. They also find cases where the in-context learned representations do not dominate  the first principal components, but they are still present in subsequent components (e.g. 3rd and 4th in Figure 3). Finally, authors demonstrate a connection between the process of in-context learning of node representations and a form of graph energy minimization.

strengths:
- The paper investigates an interesting phenomena that may further our understanding of how LLMs learn in-context;
- The paper includes a simple and easy-to-reproduce experiment that can facilitate future analysis of LLM representations;
- The paper is generally well-written, if a bit unconventionally structured. It follows the natural flow of an investigation into how representations change during in-context learning;
- Authors go to reasonable length to ablate alternative hypotheses in their analysis.

weaknesses:
- The paper focuses primarily on just a single model: Llama 3 8B. It is possible that some of the authors findings reflect not the general ability of LLMs, but an artifact of Llama 3 models, or a property of models of this particular size. The paper would arguably be improved if authors verify their findings on other models: both of different size (e.g. 1B / 70B / 405B), different family (gemma, qwen, mistral, opt, etc) and of different type (e.g. instruct vs non-instruct). Failing that, authors could at least clarify this as potential limitation.

- The part where authors argue that the in-context learning is linked to Dirichlet energy minimization is interesting, but (arguably) a bit contrived. As authors explain, energy minimization would be solved simply by assigning all nodes to the same embedding. Perhaps it would be best to further explore the analogy with graph embedding learning as a possible alternative? (e.g. distance-preserving embeddings or similar)

- While I don't count this as a significant weakness, the paper could arguably be improved by further exploring *what* in the model causes the abrupt shift in representations. (e.g. is there any connection with induction heads, transformer circuits, or any other mechanism that could be responsible for the change).

questions:
### Questions:

Note that these questions are mostly asked out of curiosity and do not affect my overall recommendation.

1. Authors found that there is a shift in concept representations after the model has seen enough context. **Could this mean that the model does not 'understand' early examples in real few-shot learning tasks?** For instance, could it mean that the model would not be able to reason about information if it requires understanding the sequences it 'read' before the change in representations? If yes, would it mean that the model would benefit from repeating the early in-context examples again, once the model 'understands' them?

2. What happens to the representation shift if you increase the number of concepts? Is there a point by which the model is no longer able to shift the structure properly to match the in-context 'meanings'?

### Minor:

L186 task-specific? (more popular with a dash)

L269 “Given that we empirically do not observe this to be the case, we can safely assume this trivial solution does not arise in our experiments.” - slightly odd phrasing - you make an observation about your experiments that implies an assumption about the same experiments? Perhaps it would be best to paraphrase / clarify the footnote.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732245425274

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Rebuttal

comment:
Dear Reviewers,

We would like to thank you for your insightful feedback! We are pleased that the reviewers unanimously found our novel synthetic data generation process and the observation of in-context emergence of structured representations to be: “interesting phenomena” (R. KJBF), “experimental design is original, elegant and well crafted” (R. a5tz), “the experiment results are interesting” (R. EbvH), and “the set of experiments are spot-on” (R. hdrb). We also appreciate the feedback and constructive suggestions for validating the robustness of our observations, which motivated us to run an extensive set of new experiments across different combinations of large language models (LLMs), graph structures, and energy functions of interest. We are glad to report that our results hold robustly across these conditions!

To fully reflect your feedback, we have implemented the following updates in our revised manuscript:

**New Experimental Results:**


- **E1: Validation of claims on additional models (Figs. 11, 12, 13).** We thank R. KJBF and hdrb for highlighting the importance of verifying our claims on additional models to enhance the robustness of our conclusions. In response, we conducted experiments on four additional models (Llama3.2-1B, Gemma2-2B, Gemma2-9B, and Llama3.1-8B-Instruct) across three distinct graph structures (square grid, ring, and hexagonal grid).  The 42 new plots resulting from these experiments are presented in Figures 11, 12, and 13 in App. C.2 and C.3. Our qualitative findings from the PCA visualizations remain consistent (Fig. 11). 


- **E2: Representations are not degenerate**: We confirm that the representations are not degenerate by demonstrating that the Dirichlet energy results hold after both normalization (Fig. 12) and standardization (Fig. 13). This resolves R. KJBF’s concern that all nodes might be assigned to the same embedding, thereby reinforcing our original claims.


- **E3: The PCA dimensions are causally relevant (Table 1).**  We conducted experiments demonstrating that the top principal components encode information causally relevant to how our model solves the task. Specifically, in App. C.4, we performed interventions on the representation projected onto top PCs of the last token presented to the model, showing that we can manipulate the model's belief about its current node in the graph, regardless of the actual last token. This demonstrates that the principal components not only represent the graphical structures but also capture causally relevant information used by the model to solve our in-context learning tasks.


- **E4: Quantifying similarity of PCA components and spectral embeddings (Table 2).** In response to the comments from R. a5tz, EbvH, and hdrb, we have quantified the cosine similarity between the graph's spectral embeddings and the principal components identified in our study. This analysis, presented in Table 2 of App. C.5, further reinforces the connection between our empirical results and spectral graph theory.


- **E5: Energy quantification for the semantic prior experiment (Fig. 14).** In response to R. EbvH's comment, we have strengthened the argument made by Fig. 3 by adding new analyses in Fig. 15. In addition to clarifying the interpretations of the results in Fig. 3c, we have quantified the energy in the first and second principal components, which are believed to embed the semantic structure. And we have also quantified the energy in the third and fourth principal components, which appear to support the in-context structure. As expected, we found that the energy of the dimensions supporting the in-context structure decreases as more context is added.

**Manuscript Updates:**


- **M1: Strengthening motivation for our synthetic task (Sec. 2).** We have enhanced the motivation for our synthetic in-context learning task by discussing recent studies on LLM representations, the use of LLMs for in-context reinforcement learning tasks, and existing experiments on humans. In Section 2, we have better emphasized how our work relates to existing literature to provide stronger grounding and motivation for our study.


- **M2: Clarifying the use of PCA and Dirichlet energy (Sec. 3,4).**  R. KJBF and EbvH highlighted the subtlety of using PCA to measure structure and draw conclusions. Initially, we used PCA to visually present results and motivate further validation through the computation of Dirichlet energy. Furthermore, we have theoretically demonstrated how spectral graph theory motivates PCA to faithfully embed spectral embeddings of the graph, as verified in E4 above. The reviewers' feedback prompted us to better emphasize this logical connection. To this end, we have updated Section 3.1 to assure readers that proper quantification of structure will follow in Section 4.

With these substantial updates, we believe our revised manuscript now addresses all the concerns. Please let us know if you have any further questions!
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper investigates how LLMs can reorganize their internal representations during in-context learning to reflect novel task-specific structures, rather than being constrained by semantic relationships learned during pre-training. The authors demonstrate this through a "graph tracing" task where models must learn to navigate graphs whose nodes are labeled with common words (like "apple", "bird") but connected in arbitrary ways. The authors show that with sufficient context examples, models can suddenly reorganize their internal representations to mirror the graph structure rather than semantic relationships, as evidenced through PCA visualizations, Dirichlet energy measurements, and causal intervention experiments. The reviewers appreciated the experimental design of the work, thorough experimental validation on multiple models and graph types, as well as theoretical grounding in graph theory and energy minimization. The weaknesses were limited experimentation on Llama, reliance on PCA, and the synthetic nature of the task. Some of these were addressed during the rebuttal/discussion phase. Overall, the reviewers were in broad agreement to accept the work.