Source-linked AI summary

ThoughtFold: Folding Reasoning Chains via Introspective Preference Learning

Ziyan Liu, Xueda Shen, Yuzhe Gu, Songyang Gao, Kuikun Liu, Guangran Cheng, Chengqi Lyu, Dahua Lin, Wenwei Zhang, Kai Chen

arXiv:2606.03503v1cs.AI

TL;DR

LRMs can overthink by memorizing redundant explorations in otherwise correct reasoning chains, while trajectory-level length penalties lack step-level credit assignment. ThoughtFold uses introspective redundancy identification and fine-grained masked preference learning to fold reasoning chains, reducing DeepSeek-R1-Distill-Qwen-7B’s token usage by approximately 56% while maintaining state-of-the-art accuracy.

  • Problem

    Outcome-based RLVR reinforces redundant explorations alongside necessary deductions in correct CoTs, while trajectory-level length penalties lack step-level credit assignment for mitigating overthinking.

  • Method

    ThoughtFold identifies redundancy within correct trajectories through introspective pruning and applies masked fine-grained preference learning to favor concise reasoning paths.

  • Results

    Approximately 56% lower average token consumption for DeepSeek-R1-Distill-Qwen-7B while maintaining state-of-the-art accuracy, with gains across reasoning models and benchmarks.

  • Takeaways & Limitations

    ThoughtFold improves efficiency–accuracy trade-offs and generalizes beyond in-domain math tasks, producing more direct solution paths rather than merely truncating verbose reasoning.

Abstract

from arXiv · show

Large Reasoning Models (LRMs) have achieved remarkable progress thanks to Reinforcement Learning with Verifiable Rewards (RLVR) on Chain-of-Thoughts (CoTs). However, since long CoTs naturally contain trial and errors and mainstream RLVR approaches choose outcome-correct CoT trajectories for memorization, the redundant explorations in long CoTs are inevitably reinforced, which results in the over-thinking issues of LRMs. Previous attempts to resolve this issue mainly give more advantage to shorter trajectories, yet their learning signals are still outcome-based and cannot reduce the memorization of redundant explorations in long CoTs. Therefore, we propose ThoughtFold, a framework that leverages fine-grained preference learning to mitigate redundant explorations for efficient reasoning. ThoughtFold employs an introspective strategy to identify redundancy within each correct trajectory, which yields a spectrum of candidate sub-trajectories. Leveraging this spectrum, we introduce a masked preference optimization objective that explicitly penalizes redundant explorations and encourages the model to directly bridge essential reasoning segments, effectively folding its reasoning chains into a more concise path. Extensive experiments show that ThoughtFold significantly enhances efficiency. It reduces the token usage of DeepSeek-R1-Distill-Qwen-7B by approximately 56% while maintaining state-of-the-art accuracy.

1. Introduction

ThoughtFold addresses overthinking in LRMs by combining outcome-based RLVR with fine-grained preference learning that identifies and penalizes redundant reasoning. Experiments show approximately 56% lower average token consumption for DeepSeek-R1-Distill-Qwen-7B while maintaining state-of-the-art accuracy.

  • Motivation: RLVR’s final-correctness supervision indiscriminately reinforces every step in correct trajectories, causing models to memorize necessary deductions and redundant explorations.This redundant memorization produces the overthinking phenomenon as an algorithmic artifact of RLVR.
  • Method: ThoughtFold integrates outcome-based RLVR with fine-grained preference learning to identify and explicitly fold redundant thoughts.Its introspective strategy starts from an outcome-correct trajectory and identifies redundancy within the reasoning chain.
  • Method: ThoughtFold explicitly penalizes redundant explorations and encourages direct bridging between essential logical steps.This folds the reasoning chain instead of uniformly reinforcing the entire correct CoT.
  • Results: Approximately 56% reduction in average token consumption was achieved for DeepSeek-R1-Distill-Qwen-7B while maintaining state-of-the-art accuracy.Experiments covered Qwen3 and DeepSeek-series reasoning models on benchmarks ranging from GSM8K to AIME and GPQA.

2. Preliminaries

The preliminaries formalize reasoning trajectories as discrete chains of thought plus final answers, then contrast outcome-based RLVR with fine-grained preference objectives. They motivate masked supervision as a way to reinforce correct reasoning steps while avoiding redundant or erroneous content.

  • Trajectory representation: An LRM πθ maps query x to trajectory τ = (z, y), where z is chain-of-thought reasoning and y is the final answer.The reasoning trajectory z is decomposed into discrete steps si using predefined rules such as blank-line delimiters.
  • Trajectory representation: Decomposing z into discrete logical units improves the management and cognition of the reasoning process.Each si represents a distinct logical unit within the reasoning path.
  • RLVR: RLVR maximizes expected trajectory rewards, but its binary final-answer signal assigns the same advantage to every CoT token.This coarse-grained assignment reinforces redundant steps contained within correct trajectories.
  • Direct Preference Optimization: DPO learns from preference pairs (x, yw, yl), treating yw as preferred over yl through an implicit policy-based reward.The preference relation is yw ≻ yl, with the winner having higher quality than the loser.
  • Direct Preference Optimization: Mask-DPO decomposes sequences into step-level annotations and uses masked divergences to reinforce correct winner steps and penalize hallucinations in loser steps.Its objective maximizes the margin between these masked terms, preventing reinforcement of errors in preferred responses or suppression of correct rejected segments.

3. Method

ThoughtFold combines outcome-based RLVR with introspective, fine-grained preference learning to identify redundant reasoning and optimize concise reasoning paths. Its dynamic masking penalizes redundant explorations while preserving essential logic, with GRPO maintaining trajectory-level accuracy supervision.

  • Framework Overview: ThoughtFold integrates fine-grained preference learning with outcome-based RLVR for efficient reasoning.The framework uses introspection to identify redundant explorations and dynamic masking to provide step-level optimization signals.
  • Introspective Redundancy Identification: An introspective prune-and-verify process generates candidate sub-trajectories by truncating tails and folding low-utility internal steps.Tail truncation uses binary search over prefix length, while internal folding ranks steps by attention-based importance and bridges retained high-importance steps.
  • Preference Pair Construction: The candidate spectrum dynamically constructs preference pairs, rewarding shorter candidates that remain correct and recording failures caused by over-aggressive pruning as negative constraints.The shortest valid trajectory discovered so far is maintained as the anchor zbest.
  • Dynamic Mask Strategy: Dynamic masks apply step-level signals that encourage Fold Anchors, mask shared reasoning, and penalize redundant or insufficiently grounded steps.For concise successes, the winner mask activates Fold Anchors and the loser mask penalizes redundant reasoning; for failures, masking penalizes the Fold Anchor and incorrect answer.
  • Joint Optimization Objective: The total objective combines fine-grained preference optimization with GRPO, balancing redundancy removal and trajectory-level accuracy through coefficient λ.GRPO supplies supervision for essential steps that receive no direct redundancy-removal signal.

4. Experiments

Experiments show that ThoughtFold improves accuracy while substantially shortening reasoning across four models and five benchmarks. Ablations and structural analyses attribute these gains to attention-guided pruning and dynamic masking, while ML@k results demonstrate increasingly efficient rollout selection.

  • Main Results: ThoughtFold improves absolute accuracy by 0.98%–2.82% and compresses sequence length by 39.4%–56.1% versus vanilla models across four models and five benchmarks.Against standard GRPO, it gains 1.04%–2.24% accuracy while reducing length by 40.2%–56.2%.
  • Main Results: ThoughtFold outperforms S-GRPO in both accuracy and efficiency and generalizes from in-domain mathematics to out-of-domain scientific tasks such as GPQA.It adaptively shortens chains more aggressively on simpler GSM8K than on complex benchmarks.
  • Ablation Study: The full ThoughtFold configuration achieves 79.00% accuracy and 5,874 tokens, outperforming random attention selection and coarse-only pruning, especially on complex reasoning benchmarks.The results support attention-guided pruning as a mechanism for distinguishing essential logic from off-target attempts.
  • Ablation Study: Disabling dynamic masking drops overall accuracy to 75.80%, below standard GRPO at 77.08%, because shared reasoning steps receive conflicting gradient updates.This identifies step-level masking as important for resolving credit-assignment ambiguity between concise and redundant trajectories.
  • Minimum Average Length@k: On AIME, ThoughtFold has a steeper ML@k decay and lower bound than Short-RL as k increases, despite comparable average lengths at ML@1.This indicates stronger efficiency gains when multiple rollouts are available.
  • Visualizing Reasoning Topology: Concept graphs show ThoughtFold eliminating the redundant loops and backtracking seen in vanilla Qwen3-8B, while Short-RL’s outcome-based reward can induce hallucinations and wrong answers.ThoughtFold instead directly derives key insights while maintaining accurate and efficient reasoning.

5. Related Works

Related work centers on RLVR for training large reasoning models and on efficiency methods that either modify inference without parameter updates or internalize concise reasoning through training. RLVR uses sparse ground-truth rewards to encourage complex exploration and self-correction, while efficient-reasoning research spans prompting, pruning, early exit, and concise-CoT fine-tuning.

  • RLVR: RLVR has shifted LRM training from supervised fine-tuning toward reinforcement learning with verifiable, ground-truth-based rewards.The rewards are sparse and encourage models to explore complex reasoning paths and self-correcting behaviors.
  • Efficient Reasoning: Efficient-reasoning methods generally divide into training-free approaches and training-based approaches.Training-free methods reduce computation without updating model parameters, whereas training-based methods aim to internalize efficiency in the model.
  • Efficient Reasoning: Training-free methods use dynamic prompting, adaptive sampling and pruning, or inference-time early-exit mechanisms to reduce computation without parameter updates.These approaches alter inference behavior rather than retraining the model.
  • Efficient Reasoning: Early training-based efforts focused on supervised fine-tuning with concise Chain-of-Thought data.These methods seek to make reasoning efficiency intrinsic to the model itself.

6. Conclusion

ThoughtFold combines introspective redundancy detection with fine-grained preference learning to enable efficient reasoning. Experiments show strong efficiency-accuracy trade-offs and generalization, while analysis attributes these gains to direct solution paths and compact trajectory representations.

  • Conclusion: ThoughtFold integrates introspection with fine-grained preference learning to explicitly penalize redundancy.The framework is designed for efficient reasoning.
  • Conclusion: ThoughtFold achieves superior efficiency-accuracy trade-offs with strong generalization.These findings come from the reported experiments.
  • Conclusion: Analysis indicates that ThoughtFold reshapes reasoning topology by creating direct solution paths rather than merely truncating verbose ones.This distinguishes structural folding from simple shortening.
  • Conclusion: Mechanistically, ThoughtFold acts as a stabilizing operator that yields more compact trajectory representations in embedding space.The compact representations are part of the proposed mechanism for efficient reasoning.

A. Quantile radius analysis

The quantile-radius analysis uses Mahalanobis distance over jointly scaled principal components to characterize reasoning-step representation distributions. ThoughtFold preserves diversity in the distributional core while compressing tail outliers, whereas Short-RL leaves the core largely unchanged but produces extreme tail instability.

  • Method: Effective Radius uses Mahalanobis distance around a global centroid with jointly computed principal components to measure statistical variation.The analysis first derives centroids for individual reasoning steps and a global mean reference.
  • ThoughtFold: ThoughtFold shows negative ΔR across q ∈ [20, 75], indicating a slightly larger radius than the Base model in the distributional core.The expanded core volume is described as preserving semantic diversity and preventing mode collapse.
  • Short-RL: For Short-RL, ΔR ≈ 0 at q < 95, suggesting that standard RL leaves the Base model’s semantic core structurally unchanged.This contrasts with the model’s behavior in the extreme tail.
  • Short-RL: At q ≈ 100, Short-RL reaches ΔR ≈ −4, indicating extreme tail expansion and instability.The analysis attributes this tail explosion to outliers moving farther from the centroid and erratic exploration away from the valid-reasoning manifold.
  • ThoughtFold: At quantiles > 80, ThoughtFold exhibits a sharp positive ΔR spike, indicating compression of the Base model’s heavy-tailed outliers.The figure characterizes this as a dual geometric phenomenon: core diversity is preserved while extreme-tail outliers are compressed.

B. Derivation of expected minimum length

The section derives the expected minimum length among k uniformly sampled, without-replacement rollouts from n sorted average lengths. It computes the minimum’s probability by counting total and favorable subsets, then restricts the expectation to feasible indices.

  • Setup: The derivation models n independent rollout lengths as sorted averages l1 ≤ l2 ≤ ··· ≤ ln and samples a uniformly random subset S of size k ≤ n.The goal is the expected value of the minimum element in S.
  • Probability derivation: The probability that li is the subset minimum is obtained as favorable subset count divided by the total number of size-k subsets.The total subset count is represented by a binomial coefficient.
  • Probability derivation: For li to be minimum, it must be included and the remaining k − 1 elements must be selected from larger-or-equal lengths, yielding n − i eligible elements.This counts subsets completed using elements after li in the sorted list.
  • Expected minimum estimator: The minimum-probability term is zero when n − i < k − 1, so only indices satisfying i ≤ n − k + 1 contribute to the expectation.This feasibility condition limits the estimator’s summation because insufficient larger elements cannot complete the subset.

C. Implementation Details · C.1. Algorithm Details

The paper specifies ThoughtFold’s implementation through separate procedures for constructing preference pairs and conducting training.

  • C.1. Algorithm Details: ThoughtFold’s implementation includes an algorithm for constructing preference pairs.This procedure is presented as Algorithm 1.
  • C.1. Algorithm Details: A second algorithm defines the training process of ThoughtFold.The training procedure is presented as Algorithm 2.
  • C.1. Algorithm Details: Algorithm 1 and Algorithm 2 together describe the core implementation workflow.The first constructs preference pairs, while the second specifies training.
  • C.1. Algorithm Details: Preference-pair construction is treated as a distinct implementation stage.The paper explicitly presents this stage separately from training.
  • C.1. Algorithm Details: ThoughtFold training is treated as a distinct implementation stage following preference-pair construction.The paper presents separate algorithms for these two stages.
  • C.1. Algorithm Details: The implementation details are organized around preference construction and model training.These are the two procedures identified in the algorithm description.

C.2. Prune-and-Verify Protocol … D.2. Reinforcement Learning

ThoughtFold constructs reliable folding targets by pruning reasoning trajectories and verifying that shortened contexts still yield correct answers. It then integrates these preference pairs with group-relative reinforcement learning, alongside related process-reward-model and reinforcement-learning foundations.

  • C.2. Prune-and-Verify Protocol: K = 4 parallel verification rollouts reduce noise when testing whether a pruned candidate preserves correct answer derivation.Verification uses four stochastic continuations instead of one.
  • C.2. Prune-and-Verify Protocol: A pruned candidate is accepted only when ˆpverify ≥0.75, requiring at least three of four correct continuations.Abnormally formatted or excessively long continuations are invalidated.
  • C.2. Prune-and-Verify Protocol: Algorithm 1 identifies introspective redundancy and constructs a preference dataset by pairing concise successful candidates with the current best trajectory.The procedure begins with tail truncation using binary search on trajectory length.
  • C.2. Prune-and-Verify Protocol: The training procedure samples trajectory groups, retains correct trajectories, and applies self-reflective preference construction to form Dpair.This data-construction stage feeds the joint policy update.
  • D.1. Process Reward Models: Process reward models provide step-level supervision by identifying incorrect or suboptimal intermediate reasoning steps.OmegaPRM uses Monte Carlo tree search and binary search to locate the first erroneous mathematical-reasoning step.
  • D.2. Reinforcement Learning: Reinforcement learning optimizes sequential decision-making policies through reward feedback and has supported control, planning, game playing, and policy optimization.The passage situates RL as a general optimization framework with deep neural networks.

E. Justification of Attention-based Pruning

ThoughtFold uses answer-token attention to rank reasoning-step importance, motivated by redundant explorations and prior evidence that attention identifies important tokens and sparse reasoning dependencies. Middle-layer attention is the default and performs best in the reported Qwen3-8B sensitivity analysis, while prune-and-verify preserves correctness.

  • Why attention can indicate reasoning-step importance: Answer-generation attention provides a natural proxy for reasoning-step importance because long reasoning chains contain redundant explorations, repeated verification, and unnecessary detours.The proxy measures how strongly answer-generation tokens attend to each reasoning step.
  • Why attention can indicate reasoning-step importance: Prior efficient-inference studies support attention-based importance estimation, including H2O’s heavy-hitter token identification and FROST’s observation of sparse attention over reasoning steps.FROST is described as directly relevant to long reasoning models, while H2O provides evidence from generative inference.
  • Pruning and verification: Lower-attention steps are prioritized for pruning and verification, but a step is foldable only if the pruned trajectory still reaches a correct answer.Attention ranks candidates; it is not used as a standalone deletion rule.
  • Why middle-layer attention is used: The default middle transformer layer is chosen because shallow layers can be noisy and lexical, whereas intermediate layers tend to encode more informative and transferable representations.The rationale is also consistent with studies finding stronger embeddings and more useful hidden representations in intermediate layers.
  • Sensitivity to layer choice: On Qwen3-8B, the middle layer achieves the best accuracy and shortest average response length among first-, middle-, last-, and all-layer attention choices.The other layer choices cause only mild degradation, indicating that performance is not highly fragile to layer selection.
Loading 2606.03503v1…