Context Structure Reshapes the Representational Geometry of Language Models

Technical Summary: Context Structure Reshapes the Representational Geometry of Language Models

Problem Statement
Large Language Models (LLMs) demonstrate the ability to adapt to new tasks and learn structure within the context window (In-Context Learning, or ICL). While prior research has established that LLMs organize input sequences into "straighter" neural trajectories in deep layers to facilitate next-token prediction via linear extrapolation, it remains unclear whether this representational straightening is a universal mechanism for all forms of ICL. Specifically, the relationship between the structure of the context and how LLMs reshape their internal representations to adapt behavior is not fully understood. This paper investigates whether representational straightening occurs consistently across diverse in-context tasks or if it is contingent upon specific task structures.

Methodology
The authors analyzed the Gemma-2-27B model (base/pre-trained version) across three distinct classes of tasks designed to probe different cognitive demands and data structures:

1. Natural Language Long-Range Dependencies: Using the LAMBADA dataset, the study examined how models integrate broader context to predict the final word of a passage, contrasting this with a "random text" control where semantic structure was destroyed.
2. Grid World Tasks: Synthetic tasks requiring the model to infer latent graph structures from linear token sequences. These included:
   • Direct Mapping: Random walks on a 6x6 lattice of unique words.
   • Hierarchical Mapping: Inferring a hidden 4x4 latent graph from noisy observations of "child" tokens associated with latent nodes.
   • 0-shot Transitions: Testing the model's ability to infer transitions between latent nodes that were never explicitly observed in the context.
3. Structured Prediction Tasks:
   • Few-Shot Learning: Standard benchmarks (e.g., Country-Capital, algorithmic reasoning) formatted as alternating "Q: [Input]" and "A: [Output]" pairs.
   • Riddles: A semantic disambiguation task from the BIG-bench dataset, contrasting 0-shot and 8-shot performance.

Geometric and Behavioral Metrics
To quantify representational geometry, the authors employed four complementary metrics across transformer layers:

• Straightening: Measured as the reduction in local curvature (angle between consecutive transition vectors) from the first layer to subsequent layers.
• Menger Straightening: A complementary curvature measure based on the radius of the circumcircle of three consecutive points.
• Effective Dimensionality: Estimated via the Participation Ratio of PCA eigenvalues to measure the volume of the representation space.
• Elongation: Measured via the anisotropy of the trajectory (variance along the first vs. second principal component).

Behavioral performance was evaluated via logit differences (for grid worlds) and exact match accuracy (for few-shot tasks).

Key Results
The study reveals a fundamental dichotomy in how LLM representations change during ICL, dependent on the task structure:

1. Continual Prediction Tasks (Natural Language & Grid Worlds):

   • In tasks requiring the integration of long-range dependencies or the inference of latent graph structures, increasing context length leads to a consistent increase in representational straightening, particularly in middle layers (Layers 15–25).
   • This straightening correlates strongly with behavioral improvement (e.g., higher logit separation between valid neighbors and non-neighbors).
   • In hierarchical grid worlds, straightening persists even in 0-shot transition conditions, suggesting the model learns the underlying latent graph rather than merely memorizing token pairs.
   • Geometrically, these tasks show a contraction of the representation onto a lower-dimensional, more anisotropic (elongated) manifold.

2. Structured Prediction Tasks (Few-Shot & Riddles):

   • In tasks requiring the mapping of discrete examples to outputs (e.g., Country-Capital, Riddles), representational straightening is inconsistent and dissociates from the task's computational core.
   • Phase Dissociation: When decomposing few-shot examples into Question, Transition (formatting), and Answer phases:
     • The Transition phase (e.g., "\n A:") exhibits progressive straightening as more examples are added, behaving like a predictable continuation.
     • The Question phase often shows a drop in straightening after the first example.
     • Crucially, the Answer phase (where the actual solution is generated) shows no progressive straightening and, in some cases (like riddles), a significant decrease in straightening with increased context.
   • This indicates that the mechanism solving the core logic of the task (mapping input to output) does not rely on the linearization of neural trajectories observed in continual prediction.

Significance and Claims
The paper argues that In-Context Learning is not a monolithic process driven by a single geometric mechanism. Instead, the authors propose that LLMs function like a "Swiss Army knife," dynamically selecting computational strategies based on the specific structure of the task.

• For continual prediction tasks, the model utilizes a strategy of representational straightening to linearize the manifold and facilitate extrapolation.
• For structured mapping tasks, the model employs distinct mechanisms that do not necessarily manifest as trajectory straightening.

The authors conclude that while representational straightening is a robust signature of learning in continual prediction settings, it is not a universal indicator of ICL capability. This finding offers a cautionary perspective for mechanistic interpretability: relying on a single geometric signature (like straightening) may fail to capture the full spectrum of how models adapt to diverse contexts. The study emphasizes the need to test interpretability hypotheses across a diverse set of task structures rather than assuming a generalized representational structure underlies all ICL behaviors.