Table of Contents
Fetching ...

ThinkEdit: Interpretable Weight Editing to Mitigate Overly Short Thinking in Reasoning Models

Chung-En Sun, Ge Yan, Tsui-Wei Weng

TL;DR

The paper identifies overly short chain-of-thought reasoning as a systematic failure in reasoning-enabled LLMs and reveals a linear reasoning-length direction in transformer residual representations, with middle layers playing a key role. It introduces ThinkEdit, a weight-editing method that targets a small subset (~4%) of attention heads to remove the short-reasoning component, requiring only about 0.2% of the model parameters. Through experiments on GSM8K, MATH-Level5, and MMLU Elementary Math (and non-math domains), ThinkEdit yields a substantial improvement for short-reasoning outputs (+6.39%) and a solid overall gain (+3.34%), while preserving most of the model’s behavior. The work provides mechanistic insights into how reasoning length is controlled in LLMs and demonstrates a practical, training-free intervention with broad relevance to reasoning quality and interpretability.

Abstract

Recent studies have shown that Large Language Models (LLMs) augmented with chain-of-thought (CoT) reasoning demonstrate impressive problem-solving abilities. However, in this work, we identify a recurring issue where these models occasionally generate overly short reasoning, leading to degraded performance on even simple mathematical problems. Specifically, we investigate how reasoning length is embedded in the hidden representations of reasoning models and its impact on accuracy. Our analysis reveals that reasoning length is governed by a linear direction in the representation space, allowing us to induce overly short reasoning by steering the model along this direction. Building on this insight, we introduce ThinkEdit, a simple yet effective weight-editing approach to mitigate the issue of overly short reasoning. We first identify a small subset of attention heads (approximately 4%) that predominantly drive short reasoning behavior. We then edit the output projection weights of these heads to remove the short reasoning direction. With changes to only 0.2% of the model's parameters, ThinkEdit effectively reduces overly short reasoning and yields notable accuracy gains for short reasoning outputs (+6.39%), along with an overall improvement across multiple math benchmarks (+3.34%). Our findings provide new mechanistic insights into how reasoning length is controlled within LLMs and highlight the potential of fine-grained model interventions to improve reasoning quality. Our code is available at: https://github.com/Trustworthy-ML-Lab/ThinkEdit

ThinkEdit: Interpretable Weight Editing to Mitigate Overly Short Thinking in Reasoning Models

TL;DR

The paper identifies overly short chain-of-thought reasoning as a systematic failure in reasoning-enabled LLMs and reveals a linear reasoning-length direction in transformer residual representations, with middle layers playing a key role. It introduces ThinkEdit, a weight-editing method that targets a small subset (~4%) of attention heads to remove the short-reasoning component, requiring only about 0.2% of the model parameters. Through experiments on GSM8K, MATH-Level5, and MMLU Elementary Math (and non-math domains), ThinkEdit yields a substantial improvement for short-reasoning outputs (+6.39%) and a solid overall gain (+3.34%), while preserving most of the model’s behavior. The work provides mechanistic insights into how reasoning length is controlled in LLMs and demonstrates a practical, training-free intervention with broad relevance to reasoning quality and interpretability.

Abstract

Recent studies have shown that Large Language Models (LLMs) augmented with chain-of-thought (CoT) reasoning demonstrate impressive problem-solving abilities. However, in this work, we identify a recurring issue where these models occasionally generate overly short reasoning, leading to degraded performance on even simple mathematical problems. Specifically, we investigate how reasoning length is embedded in the hidden representations of reasoning models and its impact on accuracy. Our analysis reveals that reasoning length is governed by a linear direction in the representation space, allowing us to induce overly short reasoning by steering the model along this direction. Building on this insight, we introduce ThinkEdit, a simple yet effective weight-editing approach to mitigate the issue of overly short reasoning. We first identify a small subset of attention heads (approximately 4%) that predominantly drive short reasoning behavior. We then edit the output projection weights of these heads to remove the short reasoning direction. With changes to only 0.2% of the model's parameters, ThinkEdit effectively reduces overly short reasoning and yields notable accuracy gains for short reasoning outputs (+6.39%), along with an overall improvement across multiple math benchmarks (+3.34%). Our findings provide new mechanistic insights into how reasoning length is controlled within LLMs and highlight the potential of fine-grained model interventions to improve reasoning quality. Our code is available at: https://github.com/Trustworthy-ML-Lab/ThinkEdit

Paper Structure

This paper contains 38 sections, 7 equations, 10 figures, 15 tables.

Figures (10)

  • Figure 1: Cumulative accuracy as a function of the reasoning length threshold. The x-axis represents the cutoff threshold on reasoning length, and the y-axis shows the average accuracy of all responses with reasoning length below that threshold. Models consistently exhibit lower accuracy for overly short reasoning (e.g. length <1000).
  • Figure 2: The overview of ThinkEdit framework. We first identify that there exist linear directions for controlling reasoning length in the hidden space, and then perform weight editing on the key attention heads.
  • Figure 3: Global steering results. Top: On GSM8K, positive $\alpha$ extends reasoning length and improves accuracy in the 8B and 14B models, while negative $\alpha$ shortens reasoning and lowers accuracy. Bottom: On MATH-Level5, negative $\alpha$ similarly shortens reasoning and reduces accuracy.
  • Figure 4: Heatmap illustrating the short reasoning contribution $\overline{C}^{h}_{\textrm{short}}$ for each attention head $h$. Heads with higher values (in red) show stronger alignment with short reasoning behavior.
  • Figure 5: Global steering with the reasoning length direction extracted from MLPs. The trend is similar as steering with attention-based directions.
  • ...and 5 more figures