SteadyThought: Mitigating LLM Under-Thinking via Thought-Level Preference Optimization

AI summary

Overview

Research area: Large Reasoning Models (LRMs) — inference efficiency, chain-of-thought behavior, and preference optimization.

Technical level: Advanced. The paper assumes familiarity with Direct Preference Optimization (DPO), SimPO, logit-level decoding control, and token entropy.

Scope: A thought-level preference optimization framework (Steady Thought, ST) that identifies where a reasoning model abandons a promising line of thought and trains it to commit to that line instead, cutting token usage while preserving or improving accuracy.

What This Paper Is About

Large Reasoning Models often find a correct line of reasoning early but keep switching to new ones ("wait", "alternatively"), producing long, wasteful responses — a failure mode the authors call under-thinking. Existing fixes suppress switching indiscriminately, which also removes the model's ability to abandon genuinely wrong reasoning. ST instead treats under-thinking as a preference problem at specific switching points: it forces the model to continue a single path from a promising thought, keeps the continuation only if it reaches a correct answer, and trains the model to prefer that persistent path over its original switch-heavy one.

Key Contributions

1. A formalization of under-thinking as preference optimization. The paper defines a commit trajectory versus a switch trajectory at each decision point and frames the goal as making the commit trajectory preferred under a Bradley-Terry-style scoring model instantiated by the model's own log-probabilities.

2. A dual-channel switching-point detector. Thought boundaries are located using either explicit linguistic markers (e.g., "alternatively", "another way") or concurrent predictive-entropy spikes at step boundaries, catching both explicit and implicit pivots.

3. A trajectory synthesis procedure. From each detected switch point, decoding is constrained by sharply lowering the logits of transition tokens ("wait", "but", "however") so the model must continue the current thought; only continuations that reach a correct answer are retained.

4. Steady Thought Preference Optimization (STPO). A reference-free, length-normalized preference objective (inspired by SimPO) applied only to the continuation after a shared prefix, so the learning signal lands exactly at the point of divergence rather than over the whole response.

Main Findings

• Consistent token savings with accuracy gains: Across four base models (DeepSeek-R1-Distill-Qwen-1.5B, Qwen3-8B, DeepSeek-R1-Distill-Qwen-14B, Qwen3-32B), ST reduced average token consumption by 30.0%, 34.5%, 19.3%, and 20.4% respectively, while average accuracy rose by 1.85, 4.00, 2.30, and 2.93 percentage points. The headline reduction cited in the abstract and conclusion is up to 63.8%.

• Per-benchmark examples: On MATH-500, DeepSeek-R1-Distill-Qwen-1.5B moved from 82.0% accuracy / 4385 tokens to 83.6% / 2691; on AIME 2024, Qwen3-8B moved from 62.1% / 10895 tokens to 64.6% / 7640.

• Generalization to unseen coding tasks: On LiveCode, used as an out-of-distribution benchmark because training data was mathematical only, token usage fell by 23.8%, 20.8%, 21.0%, and 15.9% across the four models, with accuracy gains of 2.9, 11.4, 4.7, and 3.3 points.

• Generalization to non-mathematical reasoning: On MuSR, StrategyQA and MedQA-USMLE, DeepSeek-R1-Distill-Qwen-1.5B cut average tokens by 60.2% with a 3.07-point accuracy gain; Qwen3-8B cut tokens by 34.3% with a 1.34-point gain.

• Structural change in reasoning: Comparing base and ST models, both total response length and the number of thought switches dropped, and this held even on the subset of problems solved correctly both before and after training — evidence that the savings come from more direct trajectories rather than from giving up on hard problems.

• Recovery ability is preserved: On a controlled set built from problems where the base model's first thought was judged incorrect (using DeepSeek-V4-Pro as judge), DeepSeek-R1-Distill-Qwen-1.5B improved from 13.2% to 20.3% accuracy with continuation length falling from 7838 to 5468 tokens; Qwen3-8B went from 23.5% to 26.6% with length falling from 9830 to 6783.

• Cross-lingual transfer: Trained only on English math data and tested on the Chinese GAOKAO-bench, DeepSeek-R1-Distill-Qwen-1.5B cut tokens by 38.3% (2829 to 1746) while accuracy rose from 80.2% to 83.7%; Qwen3-8B cut tokens by 57.6% (4420 to 1872) while accuracy rose from 86.0% to 88.4%.

• STPO beats alternative training schemes: Compared against SFT-standard, SFT-thought, DPO-standard, DPO-thought and SimPO-standard, STPO gave the best accuracy/efficiency combination — for example, Qwen3-8B on MATH-500 reached 93.4% at 2139 tokens versus SimPO-standard's 93.6% at 2908 tokens and DPO-thought's 92.8% at 4594 tokens.

Methodology in Plain English

1. Split responses into thoughts. Each model response is divided at structural delimiters ("\n\n") into candidate steps.

2. Detect switching points two ways. A boundary counts as a switch if the next step opens with a phrase from a fixed list of transition markers, or if predictive entropy simultaneously exceeds a threshold on both sides of the boundary. The threshold is calibrated per model by choosing the lowest value at which at least 80% of detected boundaries are genuine switches; the final thresholds used were 2.2, 1.3, 1.9, and 1.3 for the 1.5B, 8B, 14B, and 32B models.

3. Force a single-path continuation. From the question plus the current thought, the model generates a continuation while the logits of switch-triggering tokens are pushed toward zero. If that continuation reaches the correct answer, the thought is kept as "promising"; otherwise it is discarded. Working from the thought rather than the full prefix also shortens context during synthesis.

4. Build preference pairs. The chosen response is the forced persistent continuation; the rejected response is the original sequence of subsequent thoughts.

5. Train with STPO. The loss is a SimPO-style sigmoid objective with length-normalized log-probabilities, a temperature β and a margin γ, conditioned on the question and the promising thought. Because the chosen continuation is often much shorter than the rejected one, length normalization matters; and because supervision starts at the divergence point, correct early reasoning is not discarded as it would be in whole-response preference tuning.

6. Training setup. Training data came from omni-math; batch size 256; learning rate 5×10⁻⁷ for the 1.5B model with γ=1.0, β=2.0; LoRA for the 8B, 14B and 32B models with learning rates of 3×10⁻⁵ (8B and 32B, β=2.0) and 13×10⁻⁵ (14B, β=3.0); context window 6144; DeepSpeed ZeRO-3 with CPU offload; warm-up ratio 0.1 with a cosine schedule; inference at max_tokens=14336, temperature=0.6, top_p=0.95.

7. Evaluation. MATH-500 (500 problems), AIME 2024 (30 problems, averaged over eight runs), GSM8K (1,319 problems), LiveCode (400 problems, two runs), plus MuSR, StrategyQA and MedQA-USMLE for further out-of-distribution checks. Baselines are NoThink, NOWAIT, SEAL and TIP.

Why This Matters

Research impact: The paper reframes "too much thinking" not as a length problem but as a commitment problem at identifiable decision points. It argues against blanket suppression of switching and offers a targeted alternative, plus evidence that the learned behavior transfers across domains and languages rather than memorizing English transition tokens.

Real-world applications:

• Math and science tutoring systems, where long-winded model traces are costly and confusing to students.
• Competitive programming and code assistants, where the LiveCode results suggest shorter solution traces without accuracy loss.
• Medical and domain question answering, where the MedQA-USMLE results indicate shorter reasoning in high-stakes settings.
• Agentic and multi-step pipelines, where each reasoning call consumes latency and budget that shorter traces directly reduce.

Industry relevance: Token count is the dominant cost driver for reasoning-model inference. Reductions of roughly 19% to 63.8% on the same hardware translate directly into serving cost and latency savings, and the method is a post-training step that does not require architectural changes or a reward model.

Future Directions

• Distinguishing necessary from unnecessary switches more explicitly. ST reduces switching statistically and preserves recovery ability on a controlled set, but a mechanism that decides per-instance when switching is warranted remains open.

• Reducing dependence on hand-specified markers and token lists. The marker list and the suppressed transition tokens ("wait", "but", "however") are manually defined; how much performance depends on them, and whether they can be learned or extended to more languages, is not settled in the main text.

• Extending beyond math-trained data. All training used omni-math, with gains shown on coding and other OOD sets; training on mixed-domain data is an obvious next step.

• Combining with adaptive compute and RL. ST is a static post-training intervention; pairing it with difficulty-aware budget allocation or reinforcement learning over reasoning length is a natural extension the paper does not test.

Target Audience

Researchers and engineers working on reasoning-model efficiency, preference optimization, or post-training of LRMs. It is most useful to readers already comfortable with DPO/SimPO-style objectives and decoding-time interventions; practitioners focused on inference cost will find the benchmark tables directly actionable, while newcomers will need background in chain-of-thought and preference tuning to follow the method.

Authors’ abstract

Flexible switching between reasoning trajectories (i.e., thought switching) has significantly enhanced the reasoning capabilities of Large Reasoning Models (LRMs). However, existing models often switch excessively yet fail to sustain promising reasoning thoughts, a phenomenon termed "under-thinking". While recent efforts suppress switching to mitigate this, such over-correction may discard valuable trajectories. To address this challenge, we propose Steady Thought (ST), a novel thought-level preference optimization framework. ST formalizes under-thinking as a preference issue at switching points, forcing single-path continuations from identified points to yield high-quality reasoning. Then, ST performs thought-level preference optimization by treating the newly generated response as preferred and the original one as dis-preferred. Experiments across multiple models and datasets show that ST effectively reduces token consumption, achieving reductions of up to 63.8%, while maintaining strong reasoning accuracy. Further analysis suggests that ST helps models acquire a generalizable ability to reduce redundant switching across domains and languages. Code is available at https://github.com/MIRALab-USTC/SteadyThought.