🤖 machine learning

Beyond Fixed Directions: Adaptive Representation Analysis of Reasoning and Memorization in LLMs

This paper demonstrates that while reasoning and memorization tasks in LLMs are decodable via a single representation direction, this direction is not fixed but undergoes substantial geometric reorganization during reinforcement learning, even as the underlying information remains accessible.

Original authors: Shaheen Nabi

Published 2026-08-25

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Shaheen Nabi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the rapidly evolving world of artificial intelligence, large language models have become capable of solving complex puzzles and recalling vast amounts of information. Scientists have long wondered how these digital brains actually work. One prevailing idea suggests that specific behaviors, like logical reasoning or simple fact retrieval, are tied to distinct, fixed pathways within the model's internal memory. Imagine the model's mind as a vast, multi-dimensional space where every thought leaves a trace. If this idea holds true, then the difference between a model thinking through a math problem and one simply recalling a fact should be visible as a single, straight line pointing in a specific direction. This concept is so appealing that some researchers have tried to use these fixed lines as guides to train models to think better, assuming the path they found would remain steady and reliable.

However, a new study challenges the assumption that these paths stay put. Researchers took a closer look at a small language model, specifically a version known as Qwen3-0.6B, to see if this fixed-direction idea holds up when the model is put through rigorous training. They gathered a carefully balanced set of 400 examples, half of which required logical reasoning and half of which required factual recall. By stripping away distractions like sentence length or specific vocabulary, they ensured that any difference they found was truly about the type of thinking involved. Their first discovery was reassuring: they confirmed that, at any single moment in time, the difference between reasoning and remembering is indeed so clear that a single line can separate the two groups perfectly. In fact, this simple line worked just as well as a complex, high-dimensional analysis that looked at every possible angle of the data.

The story takes a turn when the researchers trained the model using a powerful technique called group-relative policy optimization, a method designed to boost reasoning skills. They expected that if the fixed-direction idea were true, the line separating reasoning from memory would stay in the same place, even as the model learned. Instead, they found that while the model still knew the difference between the two tasks perfectly well, the actual direction of that knowledge had shifted dramatically. After training, the line that once pointed clearly toward "reasoning" had rotated so much that it was no longer aligned with its original position. On average, the new direction was only about 45 percent similar to the old one, and in the deepest layers of the model, the shift was even more pronounced, with the internal representation drifting by more than half.

This finding reveals a crucial distinction between what a model knows and how it stores that knowledge. The information itself did not vanish; the model could still distinguish between a math problem and a trivia question with perfect accuracy after training. But the geometric coordinates it used to express that distinction had been completely reorganized. It is as if a map of a city remained perfectly accurate in telling you where the library is, but the compass directions on the map had been rotated so that "north" now pointed where "east" used to be. The researchers measured this shift by comparing the model's internal state before and after training, finding that the change was not just a minor adjustment but a substantial reorganization of the model's internal landscape.

The study concludes that while the ability to separate reasoning from memory is robust, the idea that this ability relies on a single, unchanging direction is incorrect. The information persists, but the specific path it takes through the model's mind is fluid and changes with learning. This suggests that for future efforts to guide artificial intelligence using these internal directions, scientists cannot simply find a path once and assume it will remain valid. Instead, they must accept that the map itself is constantly being redrawn, even as the landmarks remain in place. The stability of the knowledge does not guarantee the stability of the path used to access it.

Technical Summary: Beyond Fixed Directions

Problem Statement

Recent work in mechanistic interpretability and reinforcement learning (RL) has proposed that high-level behavioral distinctions in Large Language Models (LLMs), such as the difference between reasoning-oriented and factual-recall tasks, can be characterized by a single, fixed direction in activation space. Methods like DiRL utilize such directions as fixed geometric anchors during Group Relative Policy Optimization (GRPO) to bias exploration.

This paper investigates the validity of two underlying assumptions of the "fixed-direction" hypothesis:

Separability: Can the distinction between reasoning-oriented and factual-recall task groups be captured by a single linear direction across transformer layers?
Stability: Does the specific direction identified before RL remain geometrically aligned after the model undergoes GRPO training?

The author distinguishes between decodability (whether the information exists in the representation) and directional stability (whether the specific vector encoding that information remains invariant).

Methodology

Model and Data

Model: The study utilizes Qwen3-0.6B. For the stability analysis (RQ2), the base model is compared against a publicly released GRPO-trained checkpoint (x32/Qwen3-0.6B-GRPO-GSM8K-Think).
Dataset: A controlled dataset of 400 examples (200 reasoning-oriented, 200 factual-recall) was constructed.
- Reasoning: Sourced from MATH-500 and GSM8K.
- Factual Recall: Sourced from PopQA and TriviaQA.
- Control: To eliminate token length as a confounding variable (which previously yielded ~0.965 accuracy in uncontrolled settings), the author employed a length-matching procedure using token-length bins of width 10. This resulted in a balanced set where surface vocabulary and length are controlled, though lexical/domain correlations remain.

Experimental Design

The study addresses two Research Questions (RQs):

RQ1: Single-Direction Separability

Direction Construction: Two methods were used to define a separating direction at each of the 28 transformer layers:
1. Mean-Difference Direction: The vector difference between the class means of the hidden states.
2. Probe-Weight Direction: The weight vector of a logistic regression probe trained to classify the task groups.
Evaluation: The author compared the performance of these one-dimensional projections against a full 1024-dimensional linear probe (logistic regression) using held-out stratified cross-validation.
Controls: Experiments included a length-only classifier, a TF-IDF lexical classifier, and a random-label probe to ensure the signal was not an artifact of the pipeline.

RQ2: Directional Stability Post-GRPO

Stability Metrics: The author computed the cosine similarity between directions extracted from the base model and the RL-trained model:
- Between mean-difference directions ( $\cos(\hat{d}^B, \hat{d}^{RL})$ ).
- Between probe-weight directions ( $\cos(\hat{w}^B, \hat{w}^{RL})$ ).
Representation Drift: To measure broader geometric changes independent of task labels, they calculated the average cosine distance between matched hidden states of the same input examples across the two checkpoints.
Decodability Check: They verified whether the task groups remained decodable in the RL model using a probe trained specifically on the RL representation.

Key Results

RQ1: Single-Direction Decodability

Perfect Separability: On the controlled 400-example dataset, a single one-dimensional projection (using either the mean-difference or probe-weight direction) achieved AUROC = 1.00, matching the performance of the full 1024-dimensional probe.
Layer Performance: The best mean-difference direction occurred at layer 18, and the best probe-weight direction at layer 12. Both reached near-perfect accuracy (0.9975) and AUROC (1.0000).
Control Validation: While the length-only classifier dropped to ~~0.65 accuracy after matching, the TF-IDF baseline remained strong (~~0.96 accuracy), indicating that lexical and domain structures still correlate with task labels. However, hidden-state classifiers outperformed TF-IDF, confirming that the separation is not solely due to surface vocabulary.
Conclusion: The reasoning/factual-recall distinction is approximately single-direction separable in this model state.

RQ2: Geometric Instability

Directional Rotation: Despite perfect decodability in the RL model, the specific direction carrying the distinction did not remain stable.
- The mean-direction cosine similarity between base and RL models averaged 0.453 (median 0.492).
- The probe-direction cosine similarity averaged 0.445 (median 0.439).
- The minimum similarity dropped to ~0.17, indicating near-orthogonality in some layers.
Representation Drift: The broader representation also shifted significantly. Representation drift (1 - cosine similarity of hidden states) increased with depth, reaching 0.511 at the final layer (layer 28).
Decodability Preservation: Crucially, a probe trained on the RL representation still achieved AUROC = 1.00.
Conclusion: The information distinguishing the task groups is preserved after RL, but the geometric coordinates (the specific direction vector) used to encode that information are substantially reorganized.

Key Contributions

Controlled Protocol: The author moved from uncontrolled dataset comparisons (which were confounded by token length) to a rigorous, length-matched, multi-dataset evaluation (400 examples) to isolate semantic representation from surface features.
Direct Comparison of Estimators: The study directly compared mean-difference directions (closed-form statistics) against learned probe directions (discriminative optimization) across all 28 layers, finding they are highly aligned (~0.85 cosine) but distinct.
Separation of Information and Geometry: The paper provides empirical evidence that information preservation does not imply geometric preservation. A direction can be a valid separator in one model state but be geometrically distinct (rotated) in a post-RL state, even if the information remains perfectly decodable.

Significance and Claims

The paper makes a modest, specific claim regarding the "fixed-direction" assumption used in methods like DiRL:

Rejection of Invariance: The study challenges the assumption that a direction extracted before RL is a mechanistic invariant. The specific vector identified as the "reasoning direction" in the base model is not the same vector in the GRPO-trained model.
Validation of Decodability: The study confirms that the information regarding the reasoning/factual-recall distinction remains accessible in the hidden states after RL.
Implication for RL Methods: The author argues that interpreting a pre-RL direction as a persistent, fixed geometric anchor is an empirical assumption that should be measured, not presumed. While a fixed anchor may still be operationally useful for biasing exploration (as DiRL proposes), it should not be assumed to represent the same mechanistic axis throughout training.
Distinction: The results support the hypothesis that reasoning can be represented by a direction at a single model state, but they refute the hypothesis that this direction remains fixed across training interventions.

The author explicitly states they do not claim to have discovered a causal "reasoning neuron" or to have proven that the underlying computation has changed; rather, they demonstrate that the geometric realization of the representation changes while the decodability of the distinction persists.