Geometry of Decision Making in Language Models

AI summary

Overview

Research area: Mechanistic interpretability and representation geometry of large language models, specifically the intrinsic dimension of hidden states during multiple-choice decision making.

Technical level: Advanced. The paper assumes familiarity with transformer internals (residual stream, MLP blocks, unembedding), manifold hypothesis concepts, and nearest-neighbor-based dimension estimators.

Scope: A large-scale empirical study of how the intrinsic dimension of hidden representations evolves across layers in 28 open-weight transformer models, and how that geometry aligns with the moment a model commits to an answer.

What This Paper Is About

LLMs generalize well across tasks, but it is unclear how they internally move from reading a prompt to committing to a specific answer. Existing interpretability work (circuits, probing) shows where information lives and what it represents, but not how the geometry of the representation space reshapes itself to support a decision. This paper measures the intrinsic dimension (ID) of hidden states layer by layer in a multiple-choice question answering (MCQA) setting, and asks whether geometric compression signals the onset of decisive predictions.

Key Contributions

1. A layer-wise geometric characterization of decision making. The authors estimate ID at two points in every transformer block (MLP output and residual post) for the final (decision) token, using three estimators: Maximum Likelihood Estimation (MLE), Two Nearest Neighbors (TwoNN), and the Generalized Ratios Intrinsic Dimension Estimator (GRIDE).

2. A large-scale empirical study spanning 28 open-weight transformer models across multiple architecture families (Llama2, GPT-2, Mistral, Phi, Gemma, LLaMA, Pythia), covering both real-world benchmarks and template-based synthetic reasoning tasks.

3. Evidence linking ID peaks to decisiveness. By projecting intermediate representations into vocabulary space via the unembedding matrix (in the spirit of Logit Lens), the authors localize the layer where accuracy sharply rises and show it consistently follows an ID peak.

4. An analysis of training dynamics and in-context learning, including Pythia models tracked across training checkpoints and few-shot prompting scaled up to 50 shots, plus a public codebase and results release.

Main Findings

• A consistent "hunchback" ID profile. Across models and tasks, the MLP output ID increases in early layers, peaks at mid-network depth, and declines in later layers. The hump only appears in settings where accuracy rises substantially above baseline, which the authors treat as a marker of non-trivial abstraction and task-specific decision making.

• Residual post-activations behave differently from MLP outputs. Residual post ID changes are smoother and more monotonic, consistent with gradual accumulation of state across layers. MLP outputs show sharper, more localized ID transitions, consistent with them acting as targeted, high-leverage corrections toward task-relevant directions.

• ID peaks coincide with, or slightly precede, decisive accuracy jumps. The sharpest accuracy increases always follow the ID peak. The authors explicitly state this relationship should be read as correlational, not causal.

• Reasoning tasks compress more sharply than retrieval tasks. On COPA (causal/contextual reasoning), the ID peak is sharp and the post-peak drop aligns with decisive accuracy gains. On MMLU-STEM (fact retrieval), the ID trend is flatter and accuracy increases monotonically with no prominent compression phase.

• Few-shot prompting lowers final-layer ID. Increasing the number of in-context examples reduces the final-layer ID, especially in MLP outputs, indicating more efficient compression of the input space. This is most salient in reasoning-heavy tasks.

• Well-performing models peak earlier and drop faster. Models with stronger performance exhibit earlier ID peaks and steeper ID declines, and larger models tend to reach ID peaks earlier in the layer stack while maintaining lower terminal ID. Model families such as LLaMA and Pythia show distinct ID trends, pointing to architectural influence on representational geometry.

• Normalized ID trajectories are consistent across model sizes. When aligned by relative depth (0 to 1), models of varying sizes follow similar normalized ID profiles, with high inter-model correlation, suggesting a shared geometric inductive bias.

• Training dynamics show a reverse pattern in generative reasoning. In the Pythia suite on the arithmetic dataset, accuracy begins to rise immediately after ID reaches a low (described as a reverse peak), and the transitions are more gradual than in MCQA tasks, suggesting continuous integration of symbolic structure rather than an abrupt commitment.

• The ID range is narrow regardless of model size. ID values typically fall in the range (5, 37), while extrinsic hidden dimensions range over (768, 4096), across both synthetic and real-world datasets.

• Final-layer ID is only a weak proxy for accuracy. Final-layer ID negatively correlates with accuracy across several datasets in the LLaMA family, but the correlation magnitude varies widely by dataset. When only models performing above baseline accuracy are considered, the correlation disappears. This contrasts with prior work that reported a stronger link between last-layer ID and accuracy/error.

• In-context learning saturates unevenly. With up to 50-shot prompting, accuracy decreases after some number of examples for smaller models (possibly because they extract surface patterns rather than the underlying arithmetic operation), while larger models improve. The 2.8B model is noted as an exception to the trend of larger models benefiting. Only 410M to 6.9B Pythia models were tested in this setting.

Methodology in Plain English

The researchers treat the model's internal state as a cloud of points: for each layer, they collect the hidden vector at the position of the final token (the token that must produce the answer) across every example in a dataset. That cloud of vectors lies on some lower-dimensional surface, and the intrinsic dimension is the minimum number of dimensions needed to describe that surface without much information loss.

They estimate this with three distance-based methods. MLE looks at the neighborhood around each point and infers a local dimension, then aggregates across points. TwoNN uses only the ratio between the distance to the second-nearest and first-nearest neighbor, exploiting a known statistical distribution. GRIDE generalizes TwoNN by using higher-order neighbor ratios, which makes it more robust to noise and curvature.

They do this at two places inside each transformer block: the MLP output (before it is written back to the residual stream) and the residual post (after the update). Separately, they take the same hidden vectors, multiply them by the unembedding matrix to project them back into vocabulary space, and measure how often that projection predicts the correct answer. Plotting accuracy alongside ID across layers lets them see where decisions solidify.

For real-world evaluation they use MCQA-formatted prompts over linguistic ability (CoLA), topic knowledge (AG News), field-specific knowledge (MMLU, split by STEM, humanities, social sciences, and other), sentiment analysis (Rotten Tomatoes, SST2), and reasoning (COPA, COLD). Correct and wrong options are shuffled to avoid positional bias. For controlled study they use synthetic template tasks: the Greater Than task, where the model must assign higher probability to continuations 44 through 99 than to 00 through 42 given a year prefix, and an arithmetic task with operands x1 in [0,99] and x2 in [1,50] and operations "plus" or "times". They use the Pythia suite, 16 autoregressive transformers from 14M to 6.9B parameters, with 154 publicly released training checkpoints per model, to watch geometry evolve during training. Representations were saved using Transformer-Lens.

Why This Matters

Impact on research. The paper reframes interpretability around representational geometry, offering a complement to circuit-level and probing-based analyses. It also provides a cautionary result: last-layer intrinsic dimension is not a reliable universal proxy for model accuracy, which tempers earlier claims in the literature and suggests that where you measure ID in the layer stack matters more than the final value.

Real-world applications:

• Model diagnostics and debugging. Layer-wise ID profiles could reveal whether a model is genuinely abstracting a task or latching onto surface patterns, which the authors observed in smaller Pythia models on arithmetic.
• Evaluation without labels. ID trends during in-context learning are proposed as a potentially unsupervised signal for ICL efficiency and saturation, useful when labeled data is scarce.
• Prompt engineering. The finding that additional few-shot examples compact the representation space gives a geometric justification for shot selection, and the observed saturation point suggests when adding more examples stops helping.
• Architecture and scaling decisions. Differences between model families and the tendency of larger models to peak earlier and end lower hint at how architectural choices shape the geometry of decision making.

Industry relevance. The work uses open-weight models and a public codebase, so practitioners can run the same analysis on in-house models. Identifying the layer where a model becomes decisive is directly relevant to intervention methods such as steering, activation editing, and early-exit strategies. The observation that a token-level projection at a specific layer already supports accurate prediction is relevant to efficient inference.

Future Directions

• Better ID estimators. The authors note that current estimators provide only a noisy estimate of the underlying manifold, and propose revalidating results with more advanced or newly developed intrinsic dimension estimators.

• Extending beyond single-token prediction. The current setup analyzes only the decision token, not open-ended autoregressive generation. Extending the analysis to natural language generation tasks is flagged as difficult because of autoregressive decoding, but is an explicit open problem.

• Looking inside transformer blocks. The study only uses features transformed by transformer blocks (MLP output and residual post). The authors suggest the same experimental setup could be applied to hidden representations within the blocks.

• Multi-manifold and broader data settings. The paper assumes representations come from a single manifold and notes the Hidalgo estimator as a possible extension for datasets with multiple underlying manifolds, along with testing on more language data sources.

Target Audience

Interpretability and mechanistic interpretability researchers; scientists studying representation geometry and the manifold hypothesis in deep networks; practitioners working on LLM evaluation, prompt engineering, or activation-level interventions; and graduate students with a background in transformer internals and dimensionality estimation who want a concrete, large-scale empirical reference for how representation geometry relates to decision making.

Authors’ abstract

Large Language Models (LLMs) show strong generalization across diverse tasks, yet the internal decision-making processes behind their predictions remain opaque. In this work, we study the geometry of hidden representations in LLMs through the lens of \textit{intrinsic dimension} (ID), focusing specifically on decision-making dynamics in a multiple-choice question answering (MCQA) setting. We perform a large-scale study, with 28 open-weight transformer models and estimate ID across layers using multiple estimators, while also quantifying per-layer performance on MCQA tasks. Our findings reveal a consistent ID pattern across models: early layers operate on low-dimensional manifolds, middle layers expand this space, and later layers compress it again, converging to decision-relevant representations. Together, these results suggest LLMs implicitly learn to project linguistic inputs onto structured, low-dimensional manifolds aligned with task-specific decisions, providing new geometric insights into how generalization and reasoning emerge in language models.