Context Structure Reshapes the Representational Geometry of Language Models

TL;DR

This paper investigates how large language models' internal representations change during in-context learning, revealing a task-dependent dichotomy in representational straightening linked to prediction improvement.

Contribution

It demonstrates that representational straightening occurs variably within tasks, depending on their structure, and introduces a dynamic view of how LLMs adapt their internal strategies.

Findings

In continual prediction tasks, increasing context leads to more straightened neural trajectories.

In structured prediction tasks, straightening occurs only during explicit structural phases.

Representational straightening correlates with prediction performance in some contexts.

Abstract

Large Language Models (LLMs) have been shown to organize the representations of input sequences into straighter neural trajectories in their deep layers, which has been hypothesized to facilitate next-token prediction via linear extrapolation. Language models can also adapt to diverse tasks and learn new structure in context, and recent work has shown that this in-context learning (ICL) can be reflected in representational changes. Here we bring these two lines of research together to explore whether representation straightening occurs \emph{within} a context during ICL. We measure representational straightening in Gemma 2 models across a diverse set of in-context tasks, and uncover a dichotomy in how LLMs' representations change in context. In continual prediction settings (e.g., natural language, grid world traversal tasks) we observe that increasing context increases the straightness of neural sequence trajectories, which is correlated with improvement in model prediction. Conversely, in structured prediction settings (e.g., few-shot tasks), straightening is inconsistent -- it is only present in phases of the task with explicit structure (e.g., repeating a template), but vanishes elsewhere. These results suggest that ICL is not a monolithic process. Instead, we propose that LLMs function like a Swiss Army knife: depending on task structure, the LLM dynamically selects between strategies, only some of which yield representational straightening.