ICLR: In-Context Learning of Representations (2501.00070v2)

Abstract: Recent work has demonstrated that semantics specified by pretraining data influence how representations of different concepts are organized in a 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.

• The paper demonstrates that in-context examples can trigger a reorganization of internal representations, aligning them with arbitrary, context-defined graph structures.
• The authors employ an in-context graph tracing task and analyze model activations using PCA and the Dirichlet energy metric to quantify this emergent behavior.
• The findings reveal that scaling context not only improves task accuracy but also uncovers implicit energy minimization processes coexisting with pretrained semantic priors.

This paper, "ICLR: In-Context Learning of Representations" (2501.00070), investigates whether LLMs can alter their internal representations of concepts based on information provided in the input context, even if this contradicts their pretraining knowledge. The core idea is to see if models can adopt novel, context-specified semantics for concepts.

To paper this, the authors design a toy task called "in-context graph tracing". This task involves defining a graph where the nodes are represented by ordinary concepts (like "apple", "bird") that the LLM would have encountered extensively during pretraining. The connectivity of the graph, however, is arbitrary and doesn't align with the pretraining semantics of these concepts (e.g., "apple" might be connected to "car" on a grid). The LLM is then given a sequence of tokens representing a random walk or neighbor pairs on this predefined graph as in-context examples. The model's task is to predict the next valid node in the sequence, based only on the graph structure implied by the context.

The key implementation detail for analysis is examining the LLM's intermediate representations. The authors collect token activations from different layers of models like Llama3.1-8B, compute the mean activation vector for each concept (node) within a recent context window, and then visualize these mean representations using Principal Component Analysis (PCA).

The empirical results show a striking phenomenon: as the amount of in-context data increases, the internal representations of the concepts undergo a sudden reorganization. While initial layers or limited context might show representation structures reflecting pretraining semantics, deeper layers with sufficient context show representations organized according to the in-context graph structure. For example, concepts arranged on a grid in the context would have their representations align in a grid-like pattern in the model's latent space (Figure 1). Similarly, a ring structure defined in context would emerge in the representations (Figure 2).

The paper further explores the case where the in-context structure conflicts with a strong semantic prior. Using concepts like "Monday", "Tuesday", etc., which naturally form a circular structure in pretraining data, the authors define a different, non-semantic ring structure in the context. They find that the semantic prior is still present in the leading principal components of the representations, but the context-specified structure emerges in subsequent principal components (Figure 3). This indicates that the context doesn't necessarily override the prior but can induce alternative, task-specific representations coexisting with the pretrained ones. Despite the representation conflict, the model's task accuracy increases with context, showing it can learn to follow the new rule (Figure A.5).

To quantify this organizational change and understand its relationship with task performance, the authors introduce the Dirichlet energy metric. For a set of node representations ${\boldsymbol{X} \in \mathbb R^{n \times d}$ and a graph $\mathcal{G}$ with adjacency matrix $\boldsymbol{A}$, the Dirichlet energy is defined as: $E_\mathcal{G}({\boldsymbol{X}) = \sum_{i,j} \boldsymbol{A}_{i,j} \|{\boldsymbol{x}_i - {\boldsymbol{x}_j\|^2$. Intuitively, this measures how "close" the representations of connected nodes are in the latent space. A lower Dirichlet energy indicates better alignment with the graph structure. The authors show that as context size increases, the Dirichlet energy of the concept representations decreases, reaching a minimum around the same point where the model's accuracy on the graph tracing task rapidly improves (Figure 4).

This coincident emergent behavior in both representation structure (low energy, clear PCA visualization) and task performance suggests an implicit learning process. The authors hypothesize that the model implicitly performs an energy minimization to infer the context-specified structure (Hypothesis 5.1). They provide a theoretical link between this hypothesis and their PCA visualizations via spectral graph theory. They prove that if the model's representations minimize Dirichlet energy in a non-degenerate way, the leading principal components will correspond to the spectral embeddings of the graph (Theorem 5.2 Informal Version). Spectral embeddings are known to visually represent graph structure well (Figure 7, Figure 8).

The sharp, phase transition-like increase in accuracy and decrease in energy as context scales is further analyzed by drawing an analogy to bond percolation in graphs. The point at which the model "gets" the structure seems to correspond to the emergence of a large connected component in the graph defined by the contextual relationships. The transition point for accuracy scaling with graph size shows power-law-like behavior (Figure 6), reminiscent of phase transitions in physical systems.

Practical Implications and Implementation Considerations:

1. Context Scaling Unlocks Capabilities: The work provides strong evidence that simply scaling the amount of in-context examples can drastically change how an LLM represents information internally, potentially enabling new capabilities aligned with the context rather than solely relying on pretraining. This suggests that for certain tasks requiring structured understanding from demonstrations, providing a sufficient quantity of high-quality examples is critical and might lead to non-linear performance gains.
2. Representation Analysis: Techniques like PCA on layer activations can be a valuable tool for debugging and understanding what an LLM is learning in-context. If a task involves structure (spatial, relational, temporal), visualizing representations might reveal if the model is encoding that structure.
3. Handling Semantic Conflict: When using ICL for tasks where the context redefines common concepts, be aware that pretrained semantics might persist, especially in certain dimensions of the representation space. Models might still perform the task correctly by utilizing other dimensions, but internal conflict exists.
4. Energy Minimization as a Lens: The Dirichlet energy metric offers a quantitative way to evaluate how well internal representations align with external structural information provided in-context. This could potentially be used as a training objective or a monitoring tool during ICL training or deployment, though the trivial solution needs careful handling (e.g., using standardized representations as in Appendix A.3).
5. Computational Requirements: Analyzing intermediate activations and performing PCA, especially across many layers and for many tokens, requires significant computational resources. Access to tools like NDIF/NNsight (mentioned in the paper) facilitates this. The analysis is often post-hoc and not directly part of the inference process for the task itself.
6. Limitations: PCA visualizations are inherently low-dimensional projections of high-dimensional data and can be misleading; while the paper provides theoretical backing and empirical correlation (Table A.2), fully understanding the causal link between these components and output behavior requires more extensive causal interventions. The percolation analogy is suggestive but requires further verification on larger graphs.

In summary, the paper provides a practical experimental framework to demonstrate that LLMs can learn and represent novel, context-defined semantics for concepts. This in-context representation learning is an emergent phenomenon tied to the amount of context provided, potentially driven by an implicit energy minimization process that identifies the underlying structure of the contextual data. This highlights the power of context scaling for unlocking new model behaviors and provides tools (PCA, Dirichlet energy) for analyzing the internal state changes that enable these behaviors.