Understanding and Steering the Cognitive Behaviors of Reasoning Models at Test-Time

TL;DR

This paper tackles inefficiencies and instability in extended chain-of-thought reasoning by uncovering cognitive-attention heads that track non-linear reasoning. It introduces CREST, a training-free framework that first calibrates head-specific steering directions offline and then performs norm-preserving activation interventions at test time to suppress unproductive cognitive modes. Empirical results across multiple models and reasoning benchmarks show CREST delivers up to 17.5% accuracy gains and up to 37.6% token reductions, with robust generalization across tasks and architectures. The approach provides both practical efficiency gains and new interpretability into how internal attention mechanisms encode cognitive behaviors in LLMs.

Abstract

Large Language Models (LLMs) often rely on long chain-of-thought (CoT) reasoning to solve complex tasks. While effective, these trajectories are frequently inefficient, leading to high latency from excessive token generation, or unstable reasoning that alternates between underthinking (shallow, inconsistent steps) and overthinking (repetitive, verbose reasoning). In this work, we study the structure of reasoning trajectories and uncover specialized attention heads that correlate with distinct cognitive behaviors such as verification and backtracking. By lightly intervening on these heads at inference time, we can steer the model away from inefficient modes. Building on this insight, we propose CREST, a training-free method for Cognitive REasoning Steering at Test-time. CREST has two components: (1) an offline calibration step that identifies cognitive heads and derives head-specific steering vectors, and (2) an inference-time procedure that rotates hidden representations to suppress components along those vectors. CREST adaptively suppresses unproductive reasoning behaviors, yielding both higher accuracy and lower computational cost. Across diverse reasoning benchmarks and models, CREST improves accuracy by up to 17.5% while reducing token usage by 37.6%, offering a simple and effective pathway to faster, more reliable LLM reasoning.

Figures (13)

• Figure 1: Visualization of probing accuracy for DeepSeek-R1-Distill-Qwen-1.5B. (Left) Accuracy on linear and non-linear reasoning steps, with high-accuracy regions (i.e., larger than 85%) highlighted in gold boxes. (Right) Accuracy measured across randomly sampled tokens. See Setup in Section \ref{['sec:prob']}.
• Figure 2: Illustration of cognitive reasoning steering at test-time.
• Figure 3: A detailed example illustrating how intervening in the original reasoning by enhancing or suppressing non-linear reasoning, alters the reasoning trajectory. Results are obtained using DeepSeek-R1-Distill-Qwen-1.5B. Please zoom in for clearer visualization.
• Figure 4: Statistical analysis of the number of reasoning steps under varying levels of intervention strength $\alpha$ in equation \ref{['eq:keyfunc1']}.
• Figure 5: Cumulative Eigenvalues of the covariance matrix of head vectors in the last layer of DeepSeek-R1-Qwen-1.5B. The PCA matrix $A$ here is of dimension $d\times d$. Notably, the top 100 principal components already capture nearly all of the variance, indicating that the effective dimensionality of the head activations is much lower than the raw space.