cs.LGAug 22, 2026
SaveBeyond Fixed Directions: Adaptive Representation Analysis of Reasoning and Memorization in LLMs
Organizations: Independent Researcher
Abstract
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.
Figures & tables
| Dataset | Task type | Raw size | Final use |
|---|---|---|---|
| GSM8K | Reasoning | 7,473 | length-binned |
| MATH-500 | Reasoning | 500 | length-binned |
| PopQA | Factual recall | 14,267 | length-binned |
| TriviaQA | Factual recall | 138,384 | length-binned |
| Total (raw) | – | 160,624 | – |
| Total (deduplicated) | – | 97,477 | – |
Table 1: Composition of the source datasets used to construct the controlled RQ1/RQ2 evaluation set. The final 400-example set (200 reasoning-oriented, 200 factual-recall) is obtained by token-length-bin matching over the deduplicated pool, not by uniform random sampling.
| Method | Layer | Accuracy | AUROC |
|---|---|---|---|
| Mean-difference direction | 18 | .9975 | 1.0000 |
| Probe-weight direction | 12 | .9975 | 1.0000 |
| Full 1024-D probe | 12 | .9975 | 1.0000 |
| Random-label probe | – | .4875 | .4816 |
Table 2: Final RQ1 results. A one-dimensional direction matches the full linear probe on the controlled task groups. Random labels produce chance-level performance.
| Baseline / classifier | Accuracy | AUROC |
|---|---|---|
| Length-only classifier | .6475 | .6645 |
| TF–IDF lexical classifier | .9600 | .9911 |
| Random-label probe | .4875 | .4816 |
| Mean-difference direction (best layer) | .9975 | 1.0000 |
| Probe-weight direction (best layer) | .9975 | 1.0000 |
| Full 1024-D probe (best layer) | .9975 | 1.0000 |
Table 3: Complete baseline comparison for the RQ1 controlled task-group distinction. The length-only baseline is substantially weakened relative to the original unmatched diagnostic (Appendix A ), while TF–IDF remains a strong but strictly weaker predictor than any hidden-state-based classifier. The random-label probe establishes the chance floor.
Figure 1: Expanded length-controlled token-length distributions after matching, shown by source dataset (left) and by task type (right). The final balanced dataset is substantially more controlled for prompt length than the original dataset comparison.
Figure 2: Expanded length-controlled final-layer PCA. Reasoning-oriented and factual-recall representations form two visually separated clusters along the leading principal components of the final-layer hidden state, consistent with an approximately linear separating structure on the controlled 400-example dataset.
Figure 3: Layer-wise AUROC for the mean-difference and probe-weight directions. Both one-dimensional directions remain near-perfect across depth.
Figure 4: Layer-wise AUROC of the full 1024-dimensional linear probe. The full probe provides the high-dimensional reference against which one-dimensional directions are evaluated.
Figure 5: Cosine alignment between the class mean-difference direction and the independently learned probe-weight direction across transformer layers.
Figure 6: Held-out layer-18 projection scores along the mean-difference direction. The two task groups occupy almost disjoint scalar ranges.
Figure 7: Base-to-GRPO cosine similarity for the mean-difference and independently learned probe directions. Both measures show strong early alignment followed by substantial rotation in middle layers and partial recovery in later layers.
Figure 8: Representation drift between the base and GRPO models. Early layers remain highly aligned, whereas deeper layers undergo substantially larger geometric changes.
Figure 9: Layer-wise probe AUROC before and after GRPO. Near-perfect decoding survives despite substantial changes in the separating geometry.
Figure 10: Cosine similarity between independently learned probe directions at corresponding base and GRPO layers. The classifier remains predictive even where the learned separating vector rotates.
| Hypothesis | Verdict | |
|---|---|---|
| H1 | Reasoning-oriented and factual-recall examples are linearly separable in hidden-state space. | Supported |
| H2 | A single one-dimensional direction is sufficient to match a full 1024-D probe. | Supported |
| H3 | The mean-difference / probe direction extracted before RL remains the same direction after GRPO. | Not supported |
| H4 | The broader representation geometry (not only the task direction) is stationary under GRPO. | Not supported |
| H5 | Task-relevant information is preserved after GRPO even though its geometric encoding changes. | Supported |
Table 4: Hypothesis-level summary. H1–H2 concern RQ1 (separability and dimensionality); H3–H5 concern RQ2 (stability). The asymmetric pattern – H1, H2, and H5 supported, H3 and H4 not supported – is the central empirical result of the paper.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 11: Final-layer PCA of the original, unmatched MATH-500-versus-PopQA diagnostic. The clean linear separation visible here is consistent with – but does not on its own distinguish between – a semantic reasoning/memorization signal and a token-length confound.
Figure 12: Prompt token-length distributions for the original MATH-500-versus-PopQA diagnostic before any length control. The Memorization (PopQA) group is concentrated at very short token lengths while the Reasoning (MATH-500) group spans a much wider and longer range, motivating the length-matching procedures described in Section 3.2.
Figure 13: Prompt token-length distributions before (left) and after (right) strict pairwise length matching in the original diagnostic. The right panel shows near-exact length equalization between the Reasoning and Memorization groups, but only 20 examples remain after this strict matching procedure.
Figure 14: Final-layer PCA of the strict, pairwise length-matched diagnostic ( ). The two groups remain qualitatively separable even under strict length control, providing early motivation for the expanded, better-powered construction used in the main text.
Figure 15: Retention of one-dimensional AUROC relative to the full probe. The one-dimensional directions retain essentially all probe performance.
Figure 16: Expanded deduplicated token-length distributions by source dataset (left) and by task type (right), prior to length-bin matching. This diagnostic motivates explicit length control before interpreting hidden-state separation.
Figure 17: Token-length distributions by source dataset (left) and by task type (right) before expanded length matching. The strong mismatch illustrates why the uncorrected dataset cannot support a clean representation claim.
More in this topic
Representation Geometry in Language Models
Representation geometry in language models studies how embeddings and hidden states organize concepts, computation, and behavior.
See all 421 papers