PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Shaheen Nabi

arXiv:2608.21919v1cs.LG

TL;DR

This paper tests whether reasoning-oriented and factual-recall task groups are approximately single-direction separable and whether that geometry remains stable after GRPO. Using Qwen3-0.6B and controlled task groups, it finds perfect single-direction decodability but substantial directional reorganization after GRPO, while decodability persists.

  • Problem

    The paper asks whether reasoning and factual-recall task groups are approximately single-direction separable and whether that geometry remains stable after reinforcement learning.

  • Method

    The study analyzes Qwen3-0.6B hidden representations using expanded, deduplicated task groups matched by token-length distributions and compares base with GRPO-trained states.

  • Results

    The best one-dimensional directions reach AUROC = 1.00, matching the full 1024-dimensional probe, while after GRPO mean-direction cosine averages 0.453, probe-direction cosine 0.445, and final-layer representation drift reaches approximately 0.511; probe AUROC remains 1.00.

  • Takeaways & Limitations

    Single-direction decodability holds for the studied task groups, but fixed-direction stability is challenged because representational information persists while its geometry changes.

  • Takeaways & Limitations

    The findings come from one model family and do not establish a changed reasoning mechanism or a universal semantic reasoning axis.

Abstract

from arXiv · show

Recent work has proposed that reasoning and memorization in language models can be characterized by a single representation direction, including methods that keep this direction fixed during reinforcement learning. We test two assumptions behind this view. First, are reasoning-oriented and factual-recall task groups approximately single-direction separable? Second, does the resulting geometry remain stable after GRPO? Using Qwen3-0.6B and a controlled 400-example dataset, we find that a one-dimensional projection can match a full 1024-dimensional linear probe with AUROC = 1.00 on the studied task groups. However, after GRPO, the corresponding direction is substantially reorganized: mean-direction cosine averages 0.453, probe-direction cosine 0.445, while direct representation drift reaches 0.511 at the final layer. Probe AUROC nevertheless remains 1.00. The evidence therefore supports single-direction decodability for the studied task groups but challenges fixed-direction stability: the information persists while its geometric realization changes.

1 Introduction

This study tests whether reasoning-versus-factual-recall representations are single-direction separable and whether that geometry remains aligned after GRPO. It finds strong separability but treats post-training stability as a distinct empirical question.

  • Linear separability and directional stability are distinct empirical properties.
  • The study tests single-direction separability and post-GRPO geometric stability in Qwen3-0.6B.
  • The analysis focuses on whether reasoning-oriented and factual-recall groups can be captured by one linear direction across transformer layers.
  • The study asks whether the separating representation remains geometrically aligned after GRPO or is reorganized by reinforcement learning.
  • The contribution is a controlled representation-level test of the fixed-direction assumption, without claims about cross-task transfer, sparse-autoencoder decomposition, or causal mechanistic replacement.

2 Prior Work and the Fixed-Direction Assumption

Prior work motivates single-direction accounts of behavioral distinctions and fixed anchors for reinforcement learning. This paper separates decodability, directional stability, and mechanistic stability, testing only the first two.

  • The linear representation hypothesis predicts that behavioral properties may correspond approximately to directions whose projections track those properties.
  • The paper provides a controlled test of reasoning-versus-factual-recall separability and whether it persists across a training intervention.
  • Prior behavioral studies motivate distinguishing generalizable computation from retrieval-like pattern matching, but do not establish representation geometry.
  • Probe accuracy can overstate representational information, so the study uses a random-label control alongside its full linear probe.
  • DiRL computes a reasoning-versus-memorization direction from the base model and keeps it fixed as a geometric anchor during GRPO.
  • The analysis distinguishes decodability, direction stability, and mechanistic stability, then tests only the first two.

3 Experimental Design

The experiments compare one-dimensional representation directions with a full probe on a controlled, length-matched dataset, then track directional alignment and representation drift from base to GRPO.

  • Representation Extraction: The study extracts 1024-dimensional hidden states at the final formatted prompt-token position across 28 transformer layers.
  • Direction Estimators: A mean-difference direction is computed from class means, providing a non-optimized estimator based on first-order class statistics.
  • Classifier Comparison: Held-out evaluation compares one-dimensional projections and normalized probe directions with a full 1024-dimensional logistic-regression probe.
  • Base-to-RL Comparison: Direction stability compares corresponding base and RL directions, whereas representation drift measures matched-example movement independently of task labels.
  • Base-to-RL Comparison: The RQ2 comparison uses the custom base model and a publicly released GRPO-trained checkpoint that was not trained in this study.
  • Controlled Dataset: The final evaluation contains 400 examples: 200 reasoning-oriented and 200 factual-recall examples from four datasets.
  • Controlled Dataset: The original MATH-500-versus-PopQA comparison was confounded by prompt length, while strict pairwise matching yielded only 20 examples.
  • Controlled Dataset: Token-length-bin matching over a deduplicated four-dataset pool produces the principal balanced evaluation set.

4 Results

On the controlled 400-example dataset, one-dimensional directions match the full 1024-dimensional probe, supporting approximately single-direction separability. However, strong lexical controls and task-group correlations limit this conclusion to the studied setting.

  • RQ1: Single-Direction Separability: AUROC = 1.0000±0.0000 for both the best mean-difference direction and the full 1024-dimensional probe.Both methods also report Accuracy = 0.9975±0.0050 at their respective best layers.
  • RQ1: Single-Direction Separability: Adding the remaining 1023 dimensions provides essentially no classification improvement over one-dimensional projections.The comparison is made at the respective best layers using the same controlled task groups.
  • Controls: 0.6475 ± 0.0533 accuracy and 0.6645 ± 0.0585 AUROC were achieved by length-only controls after matching.This fell substantially from the original unmatched diagnostic, reducing but not eliminating the length confound.
  • Controls: 0.9600 ± 0.0146 accuracy and 0.9911 ± 0.0113 AUROC were achieved by TF–IDF, which remained weaker than hidden-state classifiers.Lexical and domain structure therefore persists despite token-length equalization.
  • Layer-wise Geometry: Mean-difference and probe-weight directions have average cosine similarity of approximately 0.8495 across layers.Alignment exceeds 0.95 in the earliest layers and remains above 0.75 at its minimum near layer 15.
  • Layer-wise Geometry: Near-perfect AUROC can coexist with different separating directions, especially in middle layers where alignment falls as low as ≈0.73 near layer 15.Multiple directions can achieve ceiling classification performance in this regime.
  • RQ1 Interpretation: The studied task groups admit an approximately single-direction representation, but this is not proof of a causal reasoning neuron or universal semantic axis.The scope is limited to Qwen3-0.6B and the controlled dataset, with lexical and domain correlations remaining.

5 RQ2: Representation Stability After GRPO

RQ2 finds that GRPO substantially reorganizes the geometry of the reasoning–factual-recall distinction while preserving near-perfect decodability. Direction stability and representation drift both change across layers, showing that post-RL information preservation does not imply preservation of the original direction.

  • 5.1 GRPO Reorganizes the Separating Direction: Layer-wise direction similarity is highest early, drops through the middle, reaches a minimum near layer 15, and partially recovers later.The profile indicates a layer-dependent geometric reorganization rather than uniform scaling.
  • 5.2 Representation Drift: 0.5106 final-layer representation drift accompanies the directional change, while early-layer drift is only 0.0128.Mean drift across layers is approximately 0.2409; at layer 28, cosine similarity is approximately 0.4894.
  • 5.3 Decodability Survives the Reorganization: 1.0000 AUROC remains after GRPO despite geometric change, with RL accuracy reaching 0.9950.The best base-model layer also reaches 1.0000 AUROC and 0.9975 accuracy.
  • 5.4 Information Preservation Does Not Imply Geometric Preservation: The result supports reparameterization: GRPO preserves information for distinguishing the task groups but changes the coordinates expressing it.Held-out AUROC evaluates whether some direction separates the RL representations, not whether it matches the base-model direction.
  • 5.4 Information Preservation Does Not Imply Geometric Preservation: Mean-difference and independently trained probe directions show similar layer-wise behavior, while direct drift corroborates broader representational change.This agreement indicates the reorganization is not confined to one classifier vector or an artifact of direction estimation.
  • 5.4 Information Preservation Does Not Imply Geometric Preservation: A fixed direction is an empirical anchor, not a guaranteed invariant after post-training.RQ2 does not support treating a direction extracted before RL as training-invariant.

6 Joint Interpretation

The paper’s two research questions separate low-dimensional decodability from geometric stability. Single-direction classifiers work for the controlled task groups, but GRPO moves the representation while preserving decodability; broader claims and downstream method effects remain outside scope.

  • Phase I – representation exists: A single scalar projection matches a full 1024-dimensional probe for the controlled task distinction.RQ1 therefore supports compatibility with a low-dimensional representation of the task groups.
  • Phase II – representation moves: Direction cosine drops to roughly 0.45 on average while late-layer representation drift exceeds 0.50.RQ2 shows that the compact representation is not geometrically stationary under GRPO.
  • The combined result: The combined result supports single-direction effectiveness for the studied groups but not persistence of the same direction throughout post-training.The stronger fixed-direction claim is not supported by the observed cross-model geometry.
  • Scope and implications: The study does not show that a fixed direction cannot improve RL or evaluate DiRL’s downstream reward or accuracy.It instead argues that representational stability assumed by direction-aware RL methods should be measured rather than presumed.

7 Discussion

The experiments preserve near-perfect decodability after GRPO but challenge interpreting the original direction as a stable mechanistic object. The findings motivate tracking representation geometry during training without claiming causal mechanism or downstream RL effects.

  • What the Results Support: AUROC remains 1.00 after GRPO, showing that the reasoning-oriented versus factual-recall distinction remains decodable.The result does not require the original coordinate direction to remain unchanged.
  • What the Results Support: Mean-direction and probe-direction comparisons show substantial rotation, particularly in middle layers, challenging fixed-direction stability.The evidence concerns geometric alignment across corresponding layers, not the underlying computation.
  • Interpretation: A discriminative property can survive a change of coordinates, so preserved information and changed representational basis are compatible.High AUROC alongside low-to-moderate directional cosine is therefore not contradictory.
  • Interpretation: Cosine similarity measures geometry rather than computation, so rotation alone cannot establish a changed reasoning mechanism.Causal claims would require interventions such as activation patching, targeted ablation, or causal mediation.
  • Implications: Future representation-aware post-training may benefit from re-estimating or tracking the relevant subspace instead of treating one pre-training direction as permanently valid.The paper presents this as a design principle and leaves adaptive tracking unevaluated.
  • Implications: The experiment tests mechanistic invariance for the studied base-to-GRPO transition, not DiRL’s downstream reward or reasoning performance.It therefore does not show that DiRL’s fixed-anchor design is inferior or ineffective.

8 Limitations •

The study’s limitations constrain its scope across models, training trajectories, task labels, representation locations, geometric interpretations, and probe classes. It also does not establish causality or evaluate downstream RL performance.

  • Scope: One model family is studied, so generalization to larger Qwen3 variants or other architectures remains untested.Direction-rotation speed under RL could depend on model scale or architecture.
  • Training setting: One publicly released GRPO checkpoint is compared with the base model, preventing conclusions about when rotation occurs or how it evolves during training.The comparison is not a controlled trajectory.
  • Training setting: The external RL checkpoint was not trained under the study’s protocol, so geometry cannot be attributed to particular hyperparameters or training trajectories.Potentially relevant factors include reward design, KL penalty, learning rate, and step count.
  • Data and labels: TF–IDF remains highly predictive despite length matching, leaving lexical and domain structure as a task-group confound.Additional TF–IDF matching could reduce the confound but would cost data.
  • Data and labels: Factual recall is an operational comparison class rather than an assumption of pure memorization, because questions may involve inference and reasoning tasks may use memorized templates.The labels should not be read as definitive cognitive mechanisms.
  • Measurement: The main analysis uses final formatted prompt-token representations and does not examine intermediate positions, generated-token states, or attention patterns.Those unexamined signals could provide different evidence about the distinction.
  • Measurement: Cosine similarity cannot distinguish task-subspace rotation from a global basis change without a null distribution over task-irrelevant directions.Thus, its stability interpretation is limited.
  • Measurement: The linear direction estimators assume approximately linear encoding, although near-perfect linear AUROC leaves little practical headroom for improvement.A nonlinear probe could still reveal additional structure.

9 Reproducibility

The analysis uses deterministic, inspectable artifacts and a fixed 400-example protocol to support reproducible representation, probing, stability, and drift measurements. Saved intermediate outputs allow figures and numerical results to be regenerated without repeating extraction.

  • Artifacts and protocol: [400, 28, 1024] tensors store representations for 400 examples across 28 layers and 1024 dimensions.Base and RL tensors use the same example order, enabling per-example drift computation.
  • Artifacts and protocol: The RQ2 comparison uses the public x32/Qwen3-0.6B-GRPO-GSM8K-Think checkpoint without claiming to have trained it.Both base and RL comparisons use Qwen3-0.6B representations.
  • Data preparation: A single common formatting template is applied before tokenization so raw dataset formatting does not introduce spurious separability.Token lengths are measured after formatting.
  • Evaluation: 5-fold stratified cross-validation preserves the 200/200 class balance in every fold for held-out accuracy and AUROC.Labels are the reasoning/factual-recall distinction.
  • Analysis pipeline: Mean-difference and probe directions, cosine stability, representation drift, and probing are computed as separate analysis artifacts.The saved outputs include per-layer direction vectors, cosine and drift scalars, and cross-validated metrics.
  • Reproducibility: Intermediate metadata, indices, hidden states, directions, scalars, and cross-validated results are saved before figure rendering.Figures can therefore be regenerated without rerunning representation extraction or probing.

10 Conclusion

The study finds an asymmetric result: the studied distinction is captured by a single direction, but that direction is not stable after GRPO. Information remains decodable while its geometric realization changes.

  • Conclusion: AUROC = 1.00 for both one-dimensional directions matches the full 1024-dimensional probe on the controlled task groups.This supports single-direction decodability for the studied reasoning-oriented and factual-recall groups.
  • Conclusion: After GRPO, mean-direction cosine averages 0.453 and probe-direction cosine 0.445, while final-layer representation drift reaches approximately 0.511.Probe AUROC remains 1.00 despite this geometric reorganization.
  • Conclusion: A direction can be compactly useful at one model state without remaining invariant under post-training.The conclusion distinguishes representational decodability from geometric stability.

A Original MATH-500 versus PopQA Diagnostic

The appendix documents two earlier dataset-construction stages that motivated, but were not used as evidence for, the final RQ1/RQ2 protocol.

  • The appendix covers two earlier stages of dataset construction described in Section 3.2.
  • These stages motivated the final RQ1/RQ2 protocol.
  • The stages were explicitly excluded as evidence for the final RQ1/RQ2 protocol reported in the main text.

A.1 Stage 1: The Unmatched Diagnostic

The initial MATH-500-versus-PopQA diagnostic showed clean representation separation, but token-length imbalance confounded interpretation. Strict matching preserved qualitative separation while leaving only 20 examples, motivating a larger controlled construction.

  • Stage 1: The Unmatched Diagnostic: The unmatched MATH-500-versus-PopQA comparison separated cleanly along final-layer PCA’s first principal component, but token length was mismatched.The groups were labeled “Reasoning” and “Memorization” in the initial exploratory analysis.
  • Stage 1: The Unmatched Diagnostic: Near-exact pairwise length matching left only 20 usable examples, while the groups remained qualitatively separable in final-layer PCA.This result motivated the expanded, better-powered construction used in the main text.
  • Stage 1: The Unmatched Diagnostic: Token length strongly predicted the label in the initial dataset, so the comparison was not used as final scientific evidence.The expanded control reduces this confound but does not eliminate lexical/domain confounding.
  • Stage 1: The Unmatched Diagnostic: The original diagnostic had short PopQA prompts and a wider, longer MATH-500 range, motivating length matching.
Loading 2608.21919v1…