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

TL;DR

DIET tackles the problem of overthinking in reasoning-focused LLMs by introducing difficulty-aware token compression integrated into reinforcement learning. It jointly optimizes performance and token usage, guided by on-the-fly difficulty estimates and stable training via Advantage Weighting, plus a cyclical pressure schedule. The approach yields substantial token reductions (up to ~45%) while preserving or enhancing P@1 and enabling better inference scaling through majority voting. A core contribution is the principled handling of weighting in group-normalized RL, plus methods to preserve the natural length–difficulty relationship essential for adaptive reasoning. The results demonstrate practical gains in efficiency and scalability, offering a concrete pathway to more efficient, high-performing LLMs for complex reasoning tasks.

Abstract

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.

Figures (7)

• Figure 1: An overview of DIET by mitigating LLM verbosity using difficulty-aware training.
• Figure 2: LLM's response length relative to problem difficulty.
• Figure 3: Micro average of majority voting Pass@1 on all the benchmarks.
• Figure 4: Pearson correlation between problem difficulty and average response length of different methods. All p-values are lower than 0.01. Methods that are not included have lower correlation.
• Figure 5: Advantage Weighting vs. Reward Weighting analysis. (Left) Training curves (Pass@1 vs. Response Length) demonstrate better performance with Advantage Weighting. (Right) Final evaluation results show Advantage Weighting yields superior performance-efficiency points.

Theorems & Definitions (1)

• Remark 1: Pitfall of Naive Reward Weighting