Title: The Overthinker's DIET: Cutting Token Calories with DIfficulty-AwarE Training
Venue: NeurIPS 2025
Accepted: yes
Keywords: token compression

efficient reasoning

reinforcement learning

inference scaling
Abstract: Recent large language models (LLMs) exhibit impressive reasoning but often \textit{overthink}, 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 \textit{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 \textbf{inference scaling}. By maintaining high per-sample quality with fewer tokens, it enables better scaling performance via majority voting 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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the token overconsumption of reasoning models. In particular, the authors point out that as the problem difficulty increases, these models might overly increase the number of tokens uttered without meaningfully increasing the performance. To solve this problem, they propose difficulty aware reward estimation -- correctness reward plus adaptive trade-off parameter times token length penalty. First, difficulty of each problem is estimated as the average incorrectness of responses sampled for a given input prompt. Based on this definition, adaptive trade-off parameter is estimated. Combined with the penalty term estimated from either min/max normalization or target length normalization, reward is estimated. Naively using this new reward function with GRPO introduces distortion in length penalty term due to GRPO normalization. The authors propose independently normalizing both terms and additively combining them with an adaptive weighting term. Finally, the authors introduce a scheduling function for the length penalty term as a heuristic to escape local minima. Experimental results on reasoning benchmarks, such as MATH, AIME etc. show that the proposed approach, DIET, outperforms competitors while reducing the token frequency significantly. They also show that the relative increase in performance is higher than some baselines. Finally, as was the motivation at the beginning, response length and difficulty are much more correlated after training with DIET compared to RL or SFT fine-tuning.\n\nSTRENGTHS AND WEAKNESSES:\n# Strengths\n1. The paper is overall well-written with a clear motivation and well-explained methods and experiments sections.\n\n2. Combining difficulty with length penalty is intuitive and practically useful.\n\n3. Empirical results on improved correlation is a promising direction in improving inference-time scaling with better token usage.\n\n# Weaknesses\n1. **Pitfalls of using reward in GRPO is based on unjustified assumptions.** I have two problems with the explanation in Appendix B, both are based on the fact that $\\alpha_{ada}=\\alpha_{base} \\hat{C}$ is defined as weighted average correctness.\n- The authors make an independence assumption in Eq (13) but ignore the fact that $\\alpha_{ada}$ is defined as weighted average correctness which certainly correlates with the outcome reward function. For example, if $N=1$ and the reward is the indicator function whether a response is correct or not, it is identical to the estimated correctness. This suggests that, if we add an additional non-zero term to Eq (13), $\\alpha_{ada}$ will not cancel when $\\hat{C} \\approx 1$ or $0$.\n- The authors claim that when $\\hat{C}=0.5$, effective penalty scaling is minimized. This is not true in general. If you plot the penalty scaling term wrt $\\hat{C}$ for various $\\sigma_p$ and $\\alpha_{base}$ values, it shows a monotonic behavior with extreme points at $\\hat{C}=0$ or $1$.\n\nWhile empirical results show that naively using the combined reward in GRPO leads to worse performance, theoretical explanation is insufficient and requires further clarification.\n\n2. **Ablation on using penalty term as is.** In GRPO, can you use the penalty term (eq (5) or Eq (7)) as is in Eq (9) rather than using GRPO normalization? Is normalization for the penalty term important?\n\n3. **How does scaling N change the performance?** What happens if you use different $N$ which impacts estimated correctness as well as GRPO normalization?\n\n4. **Functional form of difficulty vs token budget relationship.** Can you discuss what functional form captures the relationship between difficulty and token budget the best? In Figure 2, you show somewhat a logarithmic relationship but the penalty terms you use are different.\n\n5. **Clarification on inference scaling.** Can you quantify the scaling efficiency of approaches in Figure 3? You claim that your approach shows faster convergence but TokenSkip seems to have steeper curves. Please also use log-scaling to align curves.\n\nQUESTIONS:\n1. Please explain pitfalls of using reward in GRPO in more details.\n\n2. What is the benefit of GRPO normalization on penalty term?\n\n3. How does scaling N change the performance?\n\n4. What is the correct functional form of difficulty vs token budget relationship?\n\n5. Can you quantify the scaling efficiency of approaches in Figure 3?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

comments: SUMMARY:
This paper proposes a suite of methods that reduce the token length of chains of thought maintaining quality.
This is done by considering the difficulty of the problem overall, as measured by the average accuracy, which is calculated in the process of doing group-level objectives such as GRPO.

STRENGTHS AND WEAKNESSES:
Strengths:
- Clarity: The paper is written quite clearly and I was able to understand all the methodology while I believe.
- Originality: The idea of using the average reward over the example is creative. And while it is intuitive, I don't think that I would have immediately come up with it.

Weaknesses:
- Significance: One major potential issue with the experimental results is that they are all done on very small models that have already been trained for reinforcement learning through distillation. It is unclear whether this sort of approach would scale up to the larger models that people actually use.

QUESTIONS:
- Do you think this approach will scale to larger models?

LIMITATIONS:
Yes, it was fine.
TITLE: Official Review
CLARITY: 4
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nTraining LLMs to reason more effectively using RL has shown impressive results in some challenging domains. However, while enhanced reasoning has enabled these models to tackle complex problems, they often tend to overthink, generating lengthy responses that are inefficient. This paper attempts to address this issue by proposing a training framework called Difficulty-Aware Training, or DIET. The core idea is to train LLMs to dynamically determine the number of tokens required to solve a task based on its difficulty.\n\nThe paper introduces two dynamic objectives: an adaptive trade-off parameter and a dynamic length target. It also discusses an implementation detail regarding how to compute the advantage when incorporating the dynamic objective (penalty) into the RL objective. The authors conducted experiments using the R1-Distilled model with the DeepScale R dataset, comparing their approach against several competitive baselines. The results demonstrate that the proposed framework not only reduces token length based on problem difficulty but also maintains strong performance.\n\n[1] L1: Controlling How Long A Reasoning Model Thinks With Reinforcement Learning by Aggarwal 2025\n\nSTRENGTHS AND WEAKNESSES:\nThe ability to reduce models' overthinking while maintaining performance on downstream tasks is a timely and important issue. The paper effectively outlines why this problem is important, and the proposed solutions are clearly articulated. Several ablation studies are conducted to demonstrate how the suggested objective improved the models' behavior in the expected manner.\n\nAlthough the paper has some strengths, it also has several weaknesses. It lacks important baselines that would clarify how the proposed idea improves upon existing ones in the literature. For instance, the authors cite [1] but do not compare their results to it. Some other relevant papers are mentioned, but no comparison is provided. The ablation study in section 4.4 challenges the paper's proposal, suggesting that an alternative indicator bias could potentially be used to compress tokens. Notably, Kimi shows good performance while reducing token generation length; although this criterion is not directly related to the tested metric, it still demonstrates an interesting property. Finally, the paper does not provide insights into how different dynamic variables interact with one another. Specifically, there are dynamic weights, losses, and learning rates, but no ablation studies showing why these properties matter or how the model performs if some are disabled. \n\n[1] L1: Controlling How Long A Reasoning Model Thinks With Reinforcement Learning by Aggarwal 2025\n\nQUESTIONS:\n- How does the proposed framework compare to prompt-based approaches such as [1]?\n- Why do you sample t(x,π_theta) from a uniform distribution? How sensitive is the framework's performance to the choice of L_max and delta?\n- What is the difference between the proposed advantage and GRPO? The GRPO Advantage is known to have bias issues, as pointed out in [2] and [3].\n- Do you know if your gradient estimate is actually unbiased or not?\n- Do you have any ablation studies demonstrating how the dynamic weight, function, and learning rate relate to each other? I'm unsure about the importance of each component.\n\n\n\n\n[1] L1: Controlling How Long A Reasoning Model Thinks With Reinforcement Learning by Aggarwal 2025\n\n[2] Understanding R1-Zero-Like Training: A Critical Perspective Liu 2025\n\n[3] https://x.com/leloykun/status/1903382502158500119?s=46\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-05"}