Context Structure Reshapes the Representational Geometry of Language Models

TL;DR

The paper investigates how the structure of in-context cues reshapes the internal representations of large language models during in-context learning. By applying four geometric metrics to Gemma-2 across natural language, grid-world, few-shot, and riddle tasks, it reveals a dichotomy: continual prediction contexts drive representational straightening and dimensionality reduction that correlates with improved prediction, while few-shot and other structured tasks show dissociated, phase-specific geometry indicating alternative mechanisms. The authors propose a Swiss Army knife view of LLM computation, where task structure selects among multiple strategies rather than a single universal mechanism. These findings urge task-aware interpretability and broader testing across architectures to understand the diverse computational tools LLMs deploy during ICL, with future work including causal interventions and model-family diversity.

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.

Figures (11)

• Figure 1: Models select computational strategies based on task structure, which influence the representational structure and geometrical measures over them
• Figure 2: Straightening in natural language context, and experimental design(A) Each transformer layer modifies the sequence representation, leading to changes in geometry. (B) Within each layer (ex. $L^p$) we first computed token curvature between all triplets of adjacent tokens in a sequence and took their average to get a sequence curvature. To measure how representation of a sequence changes as it gets transformed by each layer, we computed sequence straightening as the difference between the sequence curvature in first and subsequent transformer layers. (C) Changes in straightening are tracked across layers, focusing on the middle layers. (D) We compare natural versus random text to measure straightening induced by context. (E) Models show substantial straightening in middle layers (L15--25) for natural text compared to random. (F--H) We investigate geometry in diverse ICL settings: (F,G) Latent structure induction in Grid Worlds, (H) Few-shot learning, and (I) Semantic reasoning in Riddles
• Figure 3: Contextual adaptation in continual prediction task is reflected in straightening. (A) We probed the model's knowledge of the grid using three test conditions: (i) Short context, (ii) Long context, and (iii) Long context with repeat. (B) Logit values for neighbor versus non-neighbor tokens. While these values are similar in the short context, neighbor logits are significantly higher in the long context. (Inset) Neighbor logits are similar between the standard long context and the 'repeat' condition. (C) Straightening of target sequences across layers; long context exhibits substantially greater straightening in the middle layers. (D) Average straightening (L15-25) for short, long, and repeat conditions. (E) Experimental setup with progressively longer context windows to interrogate learning dynamics. Same color scheme is used in panels F-H. (F) Divergence of neighbor and non-neighbor logits as context length increases. (G) Corresponding increase in straightening in middle layers with longer context. (H) Relationship between the behavioral logit difference (neighbor minus non-neighbor) and straightening. (Inset) relationship between raw neighbor logits and straightening.
• Figure 4: Straightening generalizes to other representational measures.(A) Changes in target sequence representation measured via Menger straightening; long context induces higher straightening. (B) Effective dimensionality of target sequences; long context sequences exhibit lower effective dimensionality. (C) Elongation of target sequences; long context trajectories are more anisotropic, biased toward the first eigenvector direction.
• Figure 5: Straightening generalizes to learning over latent graphs (A) Latent graph structure: random walks occur on latent nodes, but the model observes only emitted tokens (e.g., "rain"). (B) Neighbor versus non-neighbor logits; long context increases neighbor logits compared to non-neighbors (C) Straightening in middle layers for short and long context. (D) Zero-shot probe: specific transitions between children of connected latent nodes were withheld during context (dashed red line) and introduced only at test time. (E) Neighbor vs. non-neighbor logits for Short, Long, and Long (0-shot) conditions. The 0-shot logits separate effectively, mirroring the fully observed Long context. (F) Straightening for Short vs. Long (0-shot) transitions; straightening is larger for inferred transitions later in context.