PaperlayerView exact revision on arXiv

Source-linked AI summary

The Reasoning-Memorization Interplay in Language Models Is Mediated by a Single Direction

Yihuai Hong, Dian Zhou, Meng Cao, Lei Yu, Zhijing Jin

arXiv:2503.23084v1cs.CLcs.AI

TL;DR

LLMs can struggle to generalize when reasoning questions differ substantially from training examples. This paper identifies linear reasoning features in model activations and reports that intervening on them improves reasoning accuracy and robustness to perturbations.

  • Problem

    LLMs sometimes struggle with reasoning questions that differ substantially from those encountered during training.

    Source: §1 · p. 1

  • Method

    The paper identifies linear reasoning features by contrasting representations of reasoning-intensive and memory-intensive questions, then studies interventions on those features during inference.

    Source: §3.1 · pp. 3–4

  • Results

    The same features explain and mediate reasoning ability across domains and languages, and interventions improve reasoning accuracy and robustness to perturbations.

    Source: §1 · p. 2, §5 · p. 7

  • Takeaways & Limitations

    The findings provide a mechanistic understanding of how reasoning and memorization interact in LLMs.

    Source: §5 · p. 7

    The study covers relatively small LLMs and mostly short-answer reasoning; whether LiReFs support deliberate reasoning remains unclear.

    Source: Limitations · p. 8

Abstract

from arXiv · show

Large language models (LLMs) excel on a variety of reasoning benchmarks, but previous studies suggest they sometimes struggle to generalize to unseen questions, potentially due to over-reliance on memorized training examples. However, the precise conditions under which LLMs switch between reasoning and memorization during text generation remain unclear. In this work, we provide a mechanistic understanding of LLMs' reasoning-memorization dynamics by identifying a set of linear features in the model's residual stream that govern the balance between genuine reasoning and memory recall. These features not only distinguish reasoning tasks from memory-intensive ones but can also be manipulated to causally influence model performance on reasoning tasks. Additionally, we show that intervening in these reasoning features helps the model more accurately activate the most relevant problem-solving capabilities during answer generation. Our findings offer new insights into the underlying mechanisms of reasoning and memory in LLMs and pave the way for the development of more robust and interpretable generative AI systems.

1 Introduction

The paper investigates how language models balance reasoning against memorization when generalizing beyond familiar questions. It proposes that linear directions in model activations mediate this balance and tests whether intervening on them affects reasoning.

1 Introduction

  • The authors frame the problem as LLMs struggling with reasoning questions that differ substantially from training examples, amid concern about reliance on memorized patterns.They argue that understanding the interplay matters for developing more reliable language-based reasoning systems.Source: §1 · p. 1
  • The authors propose that a linear feature in activation space enables systematic generalization when active, while its absence leaves the model in memorization mode.They describe this as a hypothesis about the reasoning-memorization dynamic.Source: §1 · p. 2
  • LiReFs are identified by contrasting hidden representations of reasoning-intensive and memory-intensive questions, which become linearly separable in activation space.The paper says inference-time enhancement can shift the model toward a more generalizable “thinking mode.”Source: §1 · p. 2
  • Across four LLMs and six datasets, the same reasoning features explain and mediate reasoning ability across knowledge domains and languages.The authors present these features as a shared control mechanism for switching between reasoning and memorization.Source: §1 · p. 2
  • Intervening on LiReFs at inference time is reported to enhance reasoning generalizability.The paper presents this as a validation of the features’ functionality.Source: §1 · p. 2
  • Case analyses report that LiReF mediation reduces reasoning errors and misapplication of reasoning or memorization ability.The introduction also describes reasoning and memorization as extremes on a generalizability spectrum.Source: §1 · pp. 1–2

2 Related work

Prior work distinguishes reasoning from memorization through generalization behavior and studies internal reasoning mechanisms, while linear-feature research offers a way to analyze and intervene on model behavior.

2 Related work

  • Memorization has several meanings across research, while this paper adopts poor reasoning generalizability as its definition in reasoning tasks.Other uses include verbatim reproduction of training data and counterfactual effects of omitting examples.Source: §2 · p. 2
  • Reasoning-memorization studies examine whether models infer new conclusions or reconstruct pretraining patterns, including performance changes under controlled input perturbations.The cited examples include work sequence reversal and alphabet shifting [McCoy et al., 2023; Prabhakar et al., 2024], and perturbations [Wu et al., 2024; Xie et al., 2024].Source: §2 · p. 2
  • Prior mechanistic work describes ordered reasoning stages and identifiable residual-stream traces, while attention heads contribute to both knowledge recall and latent reasoning.These processes are characterized as distinct yet interconnected [Zheng et al., 2024].Source: §2 · p. 3
  • Linear semantic features have been identified by contrasting inputs and manipulated to control behavior; this paper extends that approach to reasoning versus memory recall.The cited work discusses semantic directions and targeted interventions [Park et al., 2024; Marks and Tegmark, 2023; Stickland et al., 2024].Source: §2 · p. 3

3 Linear reasoning features (LiReFs)

The section defines LiReFs as residual-stream directions contrasting reasoning- and memory-intensive questions, then examines their separation and association with reasoning demand across models and tasks.

3.1 Background

  • LiReFs are difference-in-means vectors from last-token residual activations for reasoning-intensive versus memory-intensive questions.Adding the vector shifts activations toward the mean for reasoning-intensive questions; ablation removes the component along its direction.Source: §3.1 · pp. 3–4

3.2 Datasets and Models

  • The study uses MMLU-Pro, GSM-8K/MGSM, PopQA, and C-Eval Chinese, with four model families in base and instruction-tuned forms.GPT-4o scores split MMLU-Pro questions into reasoning and memory subsets; the other benchmarks are assigned to those subsets by dataset.Source: §3.2 · p. 4

3.3 Analysis results

  • The same LiReF linearly separates reasoning- from memory-intensive questions across task formats, domains, and languages.Questions requiring the same capability also cluster together, with task domain and language having little influence within each category.Source: §3.3 · p. 4, Figure 2 · p. 5
  • Across all eight base and instruction-tuned models, reasoning questions mostly activate LiReFs positively and memory questions negatively, especially in middle layers.Base and instruction-tuned profiles are highly consistent in three of four model families.Source: §3.3 · pp. 4–5, Figure 3 · p. 5

3.4 The gradient nature of reasoning-memorization interplay

  • 0.840 (LLaMA3-8B-base) and 0.752 (Mistral-7B-v0.3-base) are the Spearman correlations between GPT-4o reasoning scores and LiReF projections.Higher reasoning scores correspond to progressively higher projections along the LiReF direction.Source: Figure 4 · p. 5, §3.4 · p. 5
  • MBPP and HumanEval activations lie near the reasoning-memory boundary, consistent with coding tasks engaging both reasoning and memory.The comparison uses two coding datasets alongside MMLU-Pro reasoning and memory subsets.Source: Figure 5 · p. 6, §3.4 · p. 6

4 Causal validation of LiReFs

Inference-time LiReF intervention improves performance across memory-intensive and reasoning-intensive tasks, while case analyses and perturbation tests suggest it can correct capability misapplication and distinguish generalizable reasoning from memory recall.

4.1 Inference-time LiReF intervention

  • Intervening on LiReFs at inference time improves four LLMs across memory-intensive and reasoning-intensive tasks, with gains consistent across task types, domains, and languages.Table 1 presents performance for four base models on six benchmarks before and after intervention.Source: §4.1 · p. 6, Table 1 · p. 7

4.2 Cases Study

  • On 184 MMLU-Pro cases, shifting LiReF activation to match GPT-4o reasoning scores raises LLaMA3-8B-base accuracy from 0 to 0.21.The cases were selected where reasoning labels contradicted feature-subspace positions and the model answered incorrectly.Source: §4.2 · p. 6
  • The authors suggest some reasoning errors reflect insufficient activation of acquired generalizable thinking capabilities rather than missing relevant knowledge.They report that targeted inference-time feature intervention can alleviate these errors.Source: §4.2 · p. 7

4.3 Reasoning Generalization Effects

  • GSM-Symbolic performance rises as LiReF enhancement increases and drops more under suppression than GSM-8k, consistent with enhanced generalizable reasoning rather than case-based memory.GSM-Symbolic contains 5,000 instances generated from 100 GSM-8k templates with varied numerical conditions and other factors.Source: §4.3 · p. 7
  • GSM-8k's smaller suppression-related drop, including a slight gain at α = −0.05, suggests possible data leakage and memory-intensive answering; MMLU-Pro-M improves as LiReFs are suppressed, consistent with greater memory recall.The authors interpret GSM-8k's response as possible prior exposure, while the MMLU-Pro-M pattern supports their observation about memory recall.Source: §4.3 · p. 7

5 Conclusion

The paper concludes that shared linear reasoning features mediate reasoning and memorization across domains and languages, and that intervention improves reasoning-task accuracy and robustness. Example outputs illustrate correction of errors after features are shifted to the appropriate region.

5 Conclusion

  • Shared linear reasoning features in activation space mediate LLM reasoning and memorization across domains and languages, and intervention improves accuracy and robustness on reasoning tasks.The authors report that manipulating these features helps activate more relevant problem-solving capabilities during answer generation.Source: §5 · p. 7
  • Table 2 examples show errors corrected after memory-intensive questions with features in the reasoning region are shifted to the appropriate region.The examples use the LLaMA3-8B-base model.Source: Table 2 · p. 8
  • Figure 6 suggests all four base models may rely on memory for GSM-8k performance, indicating potential data-leakage risk on that reasoning dataset.The figure compares GSM-8k, GSM-Symbolic, and MMLU-Pro-M across intervention intensities.Source: Figure 6 · p. 8

Limitations

The study's scope is limited by model size, task format, and its operationalization of memorization.

Limitations

  • The study examines reasoning features only in relatively small LLMs, although scaling model size and inference-time computation can improve reasoning capability.The authors cite Hoffmann et al. (2022) and OpenAI (2024) on scaling.Source: Limitations · p. 8
  • The study focuses mostly on short-answer reasoning, leaving LiReF use for deliberate reasoning with chain-of-thought, self-reflection, or tree-of-thought prompting unclear..Source: Limitations · p. 8
  • Memorization is formulated as performance inconsistency under reasoning-question perturbation, and its mechanistic equivalence to counterfactual memorization remains an open question.Counterfactual memorization is described as answer changes after removing a similar training example.Source: Limitations · p. 8

A Prompts

The appendix identifies the prompt used to obtain GPT-4o reasoning scores for questions.

A Prompts

  • Table 3 presents the prompt used to query GPT-4o for a Reasoning Score for each question..Source: A Prompts · p. 16

B Details of Datasets

The appendix documents datasets used in Sections 3 and 4, spanning MMLU-Pro subsets, factual question answering, Chinese humanities questions, and grade-school mathematics.

MMLU-Pro-M and MMLU-Pro-R

  • MMLU-Pro questions are split into reasoning and memory subsets using GPT-4o Reasoning Scores: above 0.5 for MMLU-Pro-R and 0.5 or below for MMLU-Pro-M.The benchmark spans 14 domains, has over 12,000 questions, and offers 10 options per question.Source: B Details of Datasets · p. 16

PopQA

  • PopQA evaluates factual entity knowledge through open-domain questions without supplied answer paragraphs, using few-shot prompting and accuracy based on substring matches.Entities are represented as subject–relationship–object triplets.Source: B Details of Datasets · p. 16

C-Eval-H

  • C-Eval-H uses the humanities portion of C-EVAL, a Chinese assessment spanning 13,948 questions, 52 disciplines, and four difficulty levels.C-EVAL covers subjects from humanities to science and engineering; this work focuses on humanities.Source: B Details of Datasets · p. 16

GSM8k

  • GSM8k contains 8.5K grade-school math problems paired with natural-language solutions to assess multi-step mathematical reasoning.Source: B Details of Datasets · p. 16

MGSM

  • MGSM extends GSM8K to 10 typologically diverse languages through manual translations for evaluating multilingual reasoning.Source: B Details of Datasets · p. 16

C.1 Detailed Plot of the PCA results

The supplementary PCA analyses examine results across model layers and MMLU-Pro subject domains, and report that the first principal component captures most of the mean activation difference between memory- and reasoning-intensive questions.

C.1 Detailed Plot of the PCA results

  • The first principal component captures the majority of the mean activation difference between DMemory and DReasoning.Source: §C.1 · p. 16
  • Additional PCA results cover various layers of the LLaMA3-8B-base and Gemma2-9B-base models.Source: §C.1 · p. 16
  • The appendix also provides fine-grained PCA visualizations of questions across MMLU-Pro subject domains.Source: §C.1 · p. 16

D Details of the Intervention Experiments

The intervention-experiment appendix specifies question scoring, inference and tuning settings, and supplementary PCA visualizations.

D Details of the Intervention Experiments

Loading 2503.23084v1…