Table of Contents
Fetching ...

ICLR: In-Context Learning of Representations

Core Francisco Park, Andrew Lee, Ekdeep Singh Lubana, Yongyi Yang, Maya Okawa, Kento Nishi, Martin Wattenberg, Hidenori Tanaka

TL;DR

The paper investigates whether in-context exemplars can override pretrained semantic structure in large language models by presenting a graph-tracing task where tokens map to graph nodes and connectivity is defined independently of semantics. It demonstrates that, with enough in-context examples, representations reorganize to reflect the graph structure, as evidenced by PCA visualizations and a decreasing Dirichlet energy, and that this reorganization occurs despite semantic correlations. The authors propose an energy minimization framework to explain the emergence, linking the observed PCA structure to spectral embeddings of the inferred graph. The findings suggest that scaling context can unlock novel, context-specified representations and capabilities in LLMs, with potential implications for adaptive world representations and reasoning. However, the reorganization is not always dominant over pretrained priors, and the phenomenon appears to follow phase-transition-like dynamics tied to context length and graph size.

Abstract

Recent work has demonstrated that semantics specified by pretraining data influence how representations of different concepts are organized in a large language model (LLM). However, given the open-ended nature of LLMs, e.g., their ability to in-context learn, we can ask whether models alter these pretraining semantics to adopt alternative, context-specified ones. Specifically, if we provide in-context exemplars wherein a concept plays a different role than what the pretraining data suggests, do models reorganize their representations in accordance with these novel semantics? To answer this question, we take inspiration from the theory of conceptual role semantics and 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 from pretrained semantic representations to in-context representations aligned with the graph 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, providing evidence towards an implicit optimization process to infer context-specified semantics. Overall, our findings indicate scaling context-size can flexibly re-organize model representations, possibly unlocking novel capabilities.

ICLR: In-Context Learning of Representations

TL;DR

The paper investigates whether in-context exemplars can override pretrained semantic structure in large language models by presenting a graph-tracing task where tokens map to graph nodes and connectivity is defined independently of semantics. It demonstrates that, with enough in-context examples, representations reorganize to reflect the graph structure, as evidenced by PCA visualizations and a decreasing Dirichlet energy, and that this reorganization occurs despite semantic correlations. The authors propose an energy minimization framework to explain the emergence, linking the observed PCA structure to spectral embeddings of the inferred graph. The findings suggest that scaling context can unlock novel, context-specified representations and capabilities in LLMs, with potential implications for adaptive world representations and reasoning. However, the reorganization is not always dominant over pretrained priors, and the phenomenon appears to follow phase-transition-like dynamics tied to context length and graph size.

Abstract

Recent work has demonstrated that semantics specified by pretraining data influence how representations of different concepts are organized in a large language model (LLM). However, given the open-ended nature of LLMs, e.g., their ability to in-context learn, we can ask whether models alter these pretraining semantics to adopt alternative, context-specified ones. Specifically, if we provide in-context exemplars wherein a concept plays a different role than what the pretraining data suggests, do models reorganize their representations in accordance with these novel semantics? To answer this question, we take inspiration from the theory of conceptual role semantics and 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 from pretrained semantic representations to in-context representations aligned with the graph 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, providing evidence towards an implicit optimization process to infer context-specified semantics. Overall, our findings indicate scaling context-size can flexibly re-organize model representations, possibly unlocking novel capabilities.
Paper Structure (32 sections, 2 theorems, 13 equations, 20 figures, 2 tables)

This paper contains 32 sections, 2 theorems, 13 equations, 20 figures, 2 tables.

Key Result

Theorem 5.1

Let $\mathcal{G}$ be a graph and ${\boldsymbol{H}}\in \mathbb R^{n \times d}$ (where $n \geq d \geq 3$) be a matrix that minimizes Dirichlet energy on $\mathcal{G}$ with non-degenerated singular values, then the first two principal components of ${\boldsymbol{H}}$ will be ${\boldsymbol{z}}^{(2)}$ an

Figures (20)

  • Figure 1: Alteration of representations in accordance with context-specified semantics (grid structure). (a) We randomly arrange a set of concepts on a grid that does not reflect any correlational semantics between the tokens. (b) We then generate sequences of tokens following a random walk on the grid, inputting it as context to a Llama-3.1-8B model. (c) The model's mean token representations projected onto the top two principal components. As the number of in-context exemplars increases, there is a formation of representations mirroring the grid structure underlying the data-generating process. Representations are from the residual stream activation following layer 26.
  • Figure 2: Alteration of representations in accordance with context-specified semantics (ring structure). (a) We randomly place concepts on a ring structure unrelated to their semantics. (b) We then generate sequences of tokens by randomly sampling neighboring pairs from the ring which is used as the input context to a Llama-3.1-8B model. (c) The model's mean representation of tokens projected onto the top two principal components. As the number of in-context exemplars increases, there is a formation of representations mirroring the ring structure underlying the data-generating process. The representations are from the residual stream activations.
  • Figure 3: In-context representations form in higher principal components in the presence of semantic priors. (a) (Purple) Semantic links underlying days of the week. (Dashed blue) We define a non-semantic graph structure by linking non-neighboring days and generate tokens from this graph. (b) (Purple) The ring geometry formed by semantic links established during pre-training remains intact in the first two principal components. (c) (Dashed blue) The non-semantic structure provided in-context can be seen in the third and fourth principal components. Note that the star structure in the first two components (b), which match the ground truth graphical structure of our data generating process (a), becomes a ring in the next two principal components (c). The representations are from the residual stream activation following layer 21.
  • Figure 4: A model continuously develops task representation as it learns to traverse novel graphs in-context. We plot the accuracy of graph traversal and the Dirichlet energy of the graph, computed from the model's internal representations, as functions of context length. We note that the Dirichlet energy never reaches a perfect zero---ruling out that the representations are learning a degenerate structure, as was also seen in the PCA visualizations in Sec. \ref{['sec:resultsv1']}. (a) A 4x4 grid graph with 16 nodes. (b) A circular ring with 10 nodes. (c) A "honey-comb" hexagonal lattice, with 30 nodes.
  • Figure 5: A memorization solution cannot explain Llama's ICL graph tracing performance. We plot the rule-following accuracy from Llama-3.1-8B outputs and accuracies from a simple 1-shot and 2-shot memorization hypothesis. (a) A ring graph with 50 nodes. (b) A square grid graph with 25 nodes. In both cases, we find that the memorization solution cannot explain the accuracy ascent curve. Instead, we find a slow phase and a fast phase, which we fit with a piecewise linear fit.
  • ...and 15 more figures

Theorems & Definitions (3)

  • Theorem 5.1: Informal Version of Theorem \ref{['thm:structrual-pca-full']}
  • Theorem B.1
  • proof