# The Overthinker's DIET: Cutting Token Calories with DIfficulty-AwarE Training

arXiv: https://arxiv.org/abs/2505.19217 · HTML: https://mlsanity.com/paper/2505.19217/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 23 of its 32 references that are in the map.
- First posted on arXiv 25 May 2025.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)
- Efficient LLM inference and serving / [Efficient reasoning and agent inference](https://mlsanity.com/inference/efficient-reasoning-and-agent-inference/)

## Builds on (most cited first, 23 of 23)

- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2501.12948/) (2025)
- [HybridFlow: A Flexible and Efficient RLHF Framework](https://mlsanity.com/paper/2409.19256/) (2024)
- [Kimi k1.5: Scaling Reinforcement Learning with LLMs](https://mlsanity.com/paper/2501.12599/) (2025)
- [Training Large Language Models to Reason in a Continuous Latent Space](https://mlsanity.com/paper/2412.06769/) (2024)
- [Do NOT Think That Much for 2+3=? On the Overthinking of o1-Like LLMs](https://mlsanity.com/paper/2412.21187/) (2024)
- [L1: Controlling How Long A Reasoning Model Thinks With Reinforcement Learning](https://mlsanity.com/paper/2503.04697/) (2025)
- [Cognitive Behaviors that Enable Self-Improving Reasoners, or, Four Habits of Highly Effective STaRs](https://mlsanity.com/paper/2503.01307/) (2025)
- [Demystifying Long Chain-of-Thought Reasoning in LLMs](https://mlsanity.com/paper/2502.03373/) (2025)
- [Chain of Draft: Thinking Faster by Writing Less](https://mlsanity.com/paper/2502.18600/) (2025)
- [O1-Pruner: Length-Harmonizing Fine-Tuning for O1-Like Reasoning Pruning](https://mlsanity.com/paper/2501.12570/) (2025)
- [TokenSkip: Controllable Chain-of-Thought Compression in LLMs](https://mlsanity.com/paper/2502.12067/) (2025)
- [Training Language Models to Reason Efficiently](https://mlsanity.com/paper/2502.04463/) (2025)
- [From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step](https://mlsanity.com/paper/2405.14838/) (2024)
- [Compressed Chain of Thought: Efficient Reasoning Through Dense Representations](https://mlsanity.com/paper/2412.13171/) (2024)
- [When More is Less: Understanding Chain-of-Thought Length in LLMs](https://mlsanity.com/paper/2502.07266/) (2025)
- [Distilling System 2 into System 1](https://mlsanity.com/paper/2407.06023/) (2024)
- [C3oT: Generating Shorter Chain-of-Thought without Compromising Effectiveness](https://mlsanity.com/paper/2412.11664/) (2024)
- [Sketch-of-Thought: Efficient LLM Reasoning with Adaptive Cognitive-Inspired Sketching](https://mlsanity.com/paper/2503.05179/) (2025)
- [Concise Thoughts: Impact of Output Length on LLM Reasoning and Cost](https://mlsanity.com/paper/2407.19825/) (2024)
- [Self-Training Elicits Concise Reasoning in Large Language Models](https://mlsanity.com/paper/2502.20122/) (2025)
- [Stepwise Perplexity-Guided Refinement for Efficient Chain-of-Thought Reasoning in Large Language Models](https://mlsanity.com/paper/2502.13260/) (2025)
- [Reasoning Effort and Problem Complexity: A Scaling Analysis in LLMs](https://mlsanity.com/paper/2503.15113/) (2025)

## Built on by (most cited first, 9 of 9)

- [CRISP: Compressed Reasoning via Iterative Self-Policy Distillation](https://mlsanity.com/paper/2603.05433/) (2026)
- [Towards Concise and Adaptive Thinking in Large Reasoning Models: A Survey](https://mlsanity.com/paper/2507.09662/) (2025)
- [ARES: Multimodal Adaptive Reasoning via Difficulty-Aware Token-Level Entropy Shaping](https://mlsanity.com/paper/2510.08457/) (2025)
- [OckBench: Measuring the Efficiency of LLM Reasoning](https://mlsanity.com/paper/2511.05722/) (2025)
- [ATTNPO: Attention-Guided Process Supervision for Efficient Reasoning](https://mlsanity.com/paper/2602.09953/) (2026)
- [Self-Aligned Reward: Towards Effective and Efficient Reasoners](https://mlsanity.com/paper/2509.05489/) (2025)
- [Concise Reasoning in the Lens of Lagrangian Optimization](https://mlsanity.com/paper/2510.10168/) (2025)
- [EvoThink: Evolving Thinking in Large Reasoning Models via Self-Pruning and Aha-Moment Preference Optimization](https://mlsanity.com/paper/2607.19962/) (2026)
- [AdaThinking-E: One-Token Entropy Regulation for Adaptive Thinking](https://mlsanity.com/paper/2608.26141/) (2026)

## Abstract (arXiv)

Recent large language models (LLMs) exhibit impressive reasoning but often over-think, generating excessively long responses that hinder efficiency. We introduce DIET ( DIfficulty-AwarE Training), a framework that systematically cuts these "token calories" by integrating on-the-fly problem difficulty into the reinforcement learning (RL) process. DIET dynamically adapts token compression strategies by modulating token penalty strength and conditioning target lengths on estimated task difficulty, to optimize the performance-efficiency trade-off. We also theoretically analyze the pitfalls of naive reward weighting in group-normalized RL algorithms like GRPO, and propose Advantage Weighting technique, which enables stable and effective implementation of these difficulty-aware objectives. Experimental results demonstrate that DIET significantly reduces token counts while simultaneously improving reasoning performance. Beyond raw token reduction, we show two crucial benefits largely overlooked by prior work: (1) DIET leads to superior inference scaling. By maintaining high per-sample quality with fewer tokens, it enables better scaling performance via majority voting with more samples under fixed computational budgets, an area where other methods falter. (2) DIET enhances the natural positive correlation between response length and problem difficulty, ensuring verbosity is appropriately allocated, unlike many existing compression methods that disrupt this relationship. Our analyses provide a principled and effective framework for developing more efficient, practical, and high-performing LLMs.
