The Overthinker's DIET: Cutting Token Calories with DIfficulty-AwarE Training
Authors: Weize Chen
Jiarui Yuan
Tailin Jin
Ning Ding
Huimin Chen
Zhiyuan Liu
Maosong Sun
Venue: NeurIPS 2025 poster ()
Keywords: token compression
efficient reasoning
reinforcement learning
inference scaling
Primary area: deep_learning
TL;DR: DIET makes LLMs more token-efficient by using problem difficulty to dynamically guide compression during reinforcement learning, boosting reasoning performance and enabling superior inference scaling under fixed budgets.

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.

The Overthinker's DIET: Cutting Token Calories with DIfficult-AwarrE Training
Weize Chen, Jiarui Yuan, Jin, Ning Ding, Huimin Chen, Zhiyuan Liuf Maosong Sun
Tsinghua University
{chenwz21, puanjr22)@mmills tsinghua. cn, liuzy@tsinghua.co edu, en
Abstract
Recent large language models (LLMs) exhibit impressive reasoning but often over- think, generating excessively long responses that hinder efficiency. We introduce DIET (Dfficulty-AwarF) Training), 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, optimize the performance--fficceenee trade-off. We also theoretically analyze the pitfalls naive reward weighting group-normalized RL algorithms like GRPO, and propose Advantage Weighting technique, which enables stable and effective implementation of these difficulty-aww aware objectives. Experimental results demonstrate that DIET significantly reduces token counts while simultaneously im- proving 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 principled and ef- fective framework for developing more efficient, practical, and high- beigh-eeffrming LLMs. Our code available at thttps://gitae com tttunl//lllIT
Introduction
Recent breakthroughs in large language models (LLMs) have yielded remarkable reasoning ca- pabilities, particularly when enhanced through reinforcement learning (RL) from outcome-based rewards (OpenAl| 2024 DeepSeek-Al et al. 2025: Team 2025). These models excel in complex domains like mathematics and coding, often generating sophisticated reasoning chains (Gandhi et al.) 2025 Pan et al., 22225] Luo et al. 20255). However, this enhanced reasoning frequently comes with a significant side effect: dramatic increase in response length compared base or instruction-tuner models. While some verbosity can facilitate complex thought, it often leads to overthinking: models produce excessively long responses, sometimes thousands of tokens, even for simple queries (e.g.. "2+3=?") (Chen et al., 2024| Chang et al.) 2025 Luo et al., 2025a). This verbosity severely impacts inference latency and computational costs, hindering the practical deployment these powerful reasoning models. Initial attempts to mitigate overthinking via supervised fine-tuning (SFT), direct preference optimization (DPO), or simple length penalties in RL objectives (Team et al. 2025 Xia et al. 2025 Aggarwal & Welleck 2025) often struggle, leading to performance degradation, I.C., models that are concise but inaccurate.
"Equal Contributions Corresponding Author
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

trade-off and 'Dverthinking" Problem
Outcome: Effcient Model
B-3uf
Ovarter hoid
Partorm mannar
Over Overthveiing
Model Taa
El Elicart Model
Overtirking Mooel
@V Soww
Token
Telen Mann-Bedge.
Coally
Effccent Model
Batter Token-Perfomance
Bather Inference Scoling
Batter Dangth-rfffcuty Correlation
Diffcully Esitimation
Diffcuuty-Awers Training
Question
Objective:
Implementation:
X Reward Weighling:
Adaptive Weighting
Adaptive Target
Advantage Weighting
Comprei
Austonss
Ayeuilty
Correctness: City
Difficuity: Disuns)
quattions
Horder Lightr porety
Herder pashtom Badget
Lightr penclty
Maat
Figure 1: An overview of DIET by mitigating LLM verbosity using difficulty- aware training.
A crucial dimension often overlooked in token compression is the intrinsic link between problem difficulty and the appropriate level verbosity. contend that a "one-size-fits" strategy is fundamentally flawed. Complex problems may necessitate longer, more detailed reasoning, whereas simpler ones should elicit direct, concise answers. Indeed, as we later show ($2.3) LLMs often naturally exhibit tendency use more tokens for problems they find more challenging. Token compression methods that ignore this inherent relationship could force a suboptimal compromise across the difficulty spectrum. We argue that explicitly incorporating on on-the-fly difficulty estimation is
compression
essential for developing token compression strategies that effectively preserve reasoning capabilities.
To address these challenges, introduce DIET (Difficult--waar) Training). framework to systematically "cut token calories" from overthinking LLMs. DIET integrates on- on-the-ffy -lyy problem difficulty into the RL training process, dynamically adapting token compression based estimated task difficulty. This enables DIET to selectively encourage conciseness for simpler problems while preserving necessary verbosity for complex ones, improving the achieving better reasoning with significantly fewer tokens than the base model.
The effective implementation of such difficulty- objectives within popular RL algorithms like GRPO (Shao et al. 2024) presents its own challenges. We theoretically analyze the pitfalls of naively applying weighted rewards in these settings and propose robust Advantage Weighting technique that ensures stable and effective token compression training. Beyond standard performance gains, our work reveals two often-overlooked benefits of difficulty-aware compression: (1) Critically for practical deployment, we show that DIET leads to superior inference scaling. By maintaining high per-sample quality with fewer tokens, it enables significantly better inference scaling performance under small computational budgets, an area where most prior token compression efforts falter or show degradation. (2) We demonstrate that DIET not only compresses, but also enhances the natural positive correlation between an LLM's ' response length and problem difficulty: a desirable characteristic for adaptive reasoning that many existing compression methods inadvertently disrupt.
This paper thus formalizes the difficulty-aware token compression problem and presents DIET as comprehensive solution. Through the core adaptive mechanisms, the enabling Advantage Weighting technique, and by demonstrating unique benefits in enhancing inference scaling and preserving adaptive verbosity, DIET offers principled and highly effective approach for developing more efficient, practical, and powerful reasoning LLMs.
Problem Formulation and Preliminaries
2.1 Token Efficiency in LLMs with Reasoning Capabilities
Formally, we define the token efficiency problem studied in this paper as a multi-objective optimiza- tion challenge. Given capable reasoning LLM policy ma and distribution of reasoning problems

---

D, we aim to find policy Ti that optimizes both performance and token efficiency:
z arg max (P x>] Bx-Dlf(L(T,
(1)
where [[T,, .)) represents the expected response length (in tokens) for prompt x under policy a, P(T, 2) represents the performance (c.g., accuracy math problems) task > 0 a coefficient that controls the trade-off between performance and token efficiency, and f is a monotonically increasing transformation function with different implementations in prior work.
The current paradigm of RL from outcome rewards typically focuses solely on maximizing per- formance without considering token efficiency (a 0), which leads to overthinking. Previous approaches addressing this issue have generally attempted to incorporate token length into the reward function in an intuitive manner:
r(z, v) Toutcome (x, (2,2) a (L(())).
where Tourcome (z, y) represents the outcome reward. L(y) is the length of response y in tokens. However, the uniform penalty a treats all problems equally, failing to account for varying difficulty levels, and with inappropriate a, it quickly leads to suboptimal tradeoffs, where performance degrades substantially as response length is reduced.
2.2 Model-Based Difficulty Estimation
critical insight of our work is that the optimal response length for a task should vary according to its difficulty. Intuitively, challenging problems may benefit from extended reasoning, while simpler questions can be answered concisely without sacrificing accuracy. We estimate the difficulty of given problem based the performance of the policy model itself during training. Formally, define the estimated correctness for a given prompt æ under policy my based on sampled responses {ys )
N
Ö(z, TO) = N Σ T(ya is correct) where y: ~ no(-(xr).
The estimated difficulty D(z,To) can then be defined as 1 - Ö(z no). This formulation captures an intuitive notion: tasks that the current policy consistently fails on (C' ss 0) are considered difficult, while those consistently answered correctly (C == 1) considered easy.
Notably, popular RI. algorithms such as GRPO (Shao et al., 2024) and RLOO (Ahmadian et al.) 2024) already require sampling multiple responses (N > 1) per prompt within each training batch to estimate advantages. Therefore, computing Ö(z, T&) or D(x, тT) incurs computational overhead, as the necessary samples and correctness evaluations are already part the core RL algorithm. This makes on-the-fly difficulty estimation highly practical for integration into the training process, we explore in subsequent sections.
2.3 Preliminary Analysis: Intrinsic Correlation of Response Length and Problem Difficulty
Before introducing our compression methods, we analyze the LLM's natural verbosity relative to task difficulty. Using R RI-Distill-Qwen- 1.5B and difficulty estimated via Eq. ( we find that solution length increases with complexity, even without difficulty- aware training.
15000
12500
Token
10000
7500
5000
Fig. [ shows a strong positive correlation: av- erage response length increases with estimated problem difficulty, similar findings have been observed in Estermann & Wattenhofer (2025); Wu al. (2025). This suggests LLMss inherently allocate more tokens to harder problems, likely for more elaborate reasoning. This pivotal observation implies that common cone-size-fits-lle compression strategies, which ignore difficulty, risk either truncating vital reasoning on complex tasks or under-compisssing simple ones. Indeed, many prior
2500
0.4
Estimated Difficultty
Figure 2: LLM's response length relative to prob- lem difficulty.
3

---

compression techniques overlook this, and as shown in $4.4. can distort this natural relationship post-training.
This baseline behavior is cornerstone of our motivation. Since LLMs instinctively use more tokens for harder problems, effective compression should leverage this. We hypothesize that integrating on- the-fly difficulty awareness into compression training enables models to compress more intelligently: compress aggressively on easier problems while preserving token budgets for complex ones, which is the key to achieving superior performance-efficieneyy.
3 Difficulty-Aware Reinforcement Learning for Token Compression
The baseline analysis in 2..3 shows that LLMs naturally adjust response length based on problem difficulty. This motivates our approach: explicitly incorporating on-the-fly difficulty estimation into the RL objective. Our goal is train models that compress responses for simpler tasks while retaining reasoning for harder ones. This section presents our methods: we first formulate how to integrate difficulty into the optimization objective, then introduce key technique for stable implementation within policy gradients, and finally propose a strategy to refine training dynamics. An overview is shown Fig.
3.1 Formulating Diffieulty-Aware Optimization Objectives
The standard RL objective for token compression (Eq. (T) often focuses on task performance P(Tr, with simple, uniform trade-off for response length (T, To instill difficulty awareness, propose modifying this objective so that the pressure to be concise adapts to the estimated difficulty of prompt & under the current policy vy Eq. (3). We explore two primary strategies for this:
3.1.1 Adaptive Trade-off Parameter Cada(2; mT)
The most direct way to instill difficulty awareness into the objective from Eq. L is to make its trade-off coefficient adaptive to the estimated problem difficulty. The intuition is to apply stronger pressure for conciseness larger a) when the model finds the problem easy (high C') and relax it when the model struggles (low C).
We replace the constant a with an adaptive trade-off parameter, O'da (z, mo), calculated using the correctness estimate C(., To):
Orada (I,, no) = Obbase w((( mn));
where Otbase > is a hyperparameter controlling the overall maximum strength of the efficiency objective, and w(-) is a monotonically increasing function mapping correctness Ô E [] to non-negative weight. A simple and effective choice is the identity function, w(()) Č, resulting in the efficiency pressure scaling linearly from 0 for the hardest problems (Č 0) up to Obass for the easiest ones (C 1). For the penalty function fL(()) itself, we adopt the formulation from Team et al. 2025) and denote it as /Kimi
Ts: (iin(0, 7;),
if T(y is correct) 1
fkimi (L(y:)) =
T(y is correct) = O,
where 74 = 0.5 L(y.) 11)--ming min, L(y,)
max, L(v,) 1- L(us D(B))te The overall optimization objective effectively becomes
maximizing [p=T[P((TT x) 3 penalty strength based on difficulty, prioritizing performance for difficult problems and conciseness for easy ones, while using fixed form for f.
Eyy-r(-lu) [Qad (x. ffimi (L))(((()))) This approach directly modifies
3.1.2 Dynamic Length Target [aya(I.
An orthogonal strategy is to make the penalty function f(L(y)) itself in Eq. ( inherently adaptive to problem difficulty. Instead of relying on fixed functional form for / that only considers L(y) (or L(ys) relative to peers), we now define f to be explicitly conditioned on dynamic target length, (æ,, mo). This target length depends on the estimated problem difficulty, allowing a larger verbosity "budget" for harder problems.

---

First, we define how the target length t is determined. Based on the estimated difficulty (ax, To), we can sample target from a distribution that assigns higher values for higher difficulty:
(æ, ma) ~ Uniform(maxx(, Lmax (Ō(x, ma) - -)), Lmax D(a, na)),
(6)
where Lmax and 0 are hyperparameters defining the maximum potential target length scale, and buffer range (e.g., 0.1). This sampling procedure assigns shorter targets for harder problems (high D) and longer targets for easier ones (low D).
Next, we define our difficulty-adapiiee adaptive penalty function, (dyn 2; ne. (1,)). This function quantifies the extent to which the generated length L(y.) for response y/s <-|12) exceeds its specific target t( mT). normalized across N sampled responses {y, }= for the same prompt æ. Let p(y., t) = max(0, L(y:) - t(z, To)) be the raw exceedance. Our adaptive penalty function is:
Sayn(Ui, V, no. {ui}) p(y,, () 1p
Op +
where Up and Op are the mean and standard deviation of {p(y), /)) The overall optimization objective thus becomes maximizing Eg (~ (1,))(())))(() (P(T)) - a fayn t(z, To). {uy D]]- This approach shifts the efficiency goal from minimizing absolute length to minimizing length relative to difficulty-arn aware budget. These two strategies, adaptive a and adaptive length targets, can be explored as alternatives or potentially combined.
3.2 Implementing Weighted Objectives with Policy Gradients: The Advantage Weighting
Following the conceptual formulation of difficulty-aware objectives ($3.1, we address their RL implementation. We focus on policy gradient (PG) methods, particularly GRPO (Shao et al. 2024). which uses per-prompt, multi-sample (N > 1) advantage normalization for stability. Integrating our weighted aifficulty-aware penalties requires careful analysis of weiigh-norrmalizttinn interactions.
We show that naively combining task rewards (Foutcome) with weighted difficult- penalties into single reward signal before GRPO normalization causes problematic interactions. The combined reward's normalization factor (e.g., its standard deviation) then depends on both penalty and outcome variances. Since outcome variance often correlates with problem difficulty (being highest for medium- difficulty problems and lowest for very easy or very hard ones), the intended effect of the penalty weight in the final normalized advantage becomes distorted. Specifically, derived in Appendix [ the token penalty is wnexpectedly weakened for problems of modest difficulty where outcome variance is high. Conversely, it can be unexpectedly exacerbuted for problems where outcome variance is very low (i.e., those the model consistently gets right or wrong), undermining the desired difficulty-awant adaptation.
Remark (Pitfall of Naive Reward Weighting). In group-normalize PG algorithms like GRPO, combining reward components before normalization causes the penalty weight's 's effect be distorted by the task's 's outcome variance, undermining the intended difficulty aware adaptation.
To prevent this distortion and ensure correct weighting, we propose Advantage Weighting. This method normalizes advantages for the outcome reward (T (Toutcome.i) and raw penalty magnitude (p.) separately, before combining them using the difficulty-aware penalty weight. Specifically, advantages for task outcome and the penalty term are normalized independently using their per-prompt means (Houtoome: /bp) and standard deviations (outcome: Op)) from N rollouts:
Fo Moutcume
Joutcome + E Ap.s Pi H4p
Agutcome.
(8)
op +
Here, p; is the raw token penalty magnitude (c.g., Eqs. (5) and 77). The final policy gradient
advantage A combines the normalized components weighted by a' (e.g.. anda or a from $3.1
A Avutcomes - a'' AA..
(9)
Thus, each component is scaled by its own variance before adaptive weighting. This allows a to correctly modulate the normalized penalty's 's influence relative to that of the normalized outcome, faithfully reflecting the intended difficulty-aware trade-off and resolving the distortion. Experiments in 4.5 validate this improvement, showing significant benefits over naive reward weighting.
5

---

3.3 Refining Training Dynamies: Cyclical Compression Pressure
implemented with Advantage Weighting ($3.2) provide robust udaptive compression, constant pressure may risk premature convergence on brevity or hinder exploration. To address this and potentially enhance the perfornance-rffiiinnyy trade-off, we explore temporally modulating the compression intensity during training.
Although the difficulty-wwar objectives ($3.)
Inspired by annealing schedules, eyclically vary the difficulty-wware penalty strength using time-varying cosine modulation factor c(t) = 0.5 (1 + cos ((()) where T is the cycle period. This factor c(t) smoothly oscillates between 1 (maximum pressure) and 0 (minimum pressure) and scales the difficulty-awwy awarre penalty component within the final advantage calculation (Eq. (9):
Ái(!) Aoutcome, - c(t) (T TT) Ap.i-
(10)
This temporal variation aims to improve robustness and the final performance-effccency; tra potentially allowing the model to escape local optima related excessive brevity and better balance reasoning consolidation with conciseness pressure.
3.4 The DIET: Difficulty-AwarE Training Method
Our approach, DIET, effectively integrates all the elements discussed in this section. Specifically, DIET employs an objective that synergizes the principles of adaptive penalty strength ($3.1.) and dynamic length target ($3.1.)). This combined objective is implemented using Advantage Weighting ($3.2) and its training is optimized with Cyclical Compression Pressure ($3.3). This holistic strategy, DIET, underpins the best performance-efficienee results presented in our later experiments.
Experimental Validation
4.1 Experimental Setup
Base Model and Algorithm. Our experiments use the RI-Distilled Qwen 1.5 model (DeepSeek-A] al. 2025). For RL algorithm, employ GRPO (Shao et al.
2024). All proposed RL methods
built on GRPO and use the Advantage Weighting technique
3.2p.
Training. We use the DeepScaleR dataset (Luo et al.
2025b
featuring high-quality mathematical
problems of diverse complexities. We use veRL. (Sheng et al. 2024) as the training framework, and train models on 8 A100 GPUs. For more details of training, please refer to Appendix D
Baselines. We compare against: (1) The Base Model without any compression. (2) SFT- & DPO- Based Methods: Kimi 1.5 SFT (Team et al. 20255 (fine-tuning on shortest correct responses), Kimi
1.5 DPO (Tam al. 2025) (using shortest correct as positive DPO examples), and TokenSkip (Xia
et al. 2025) (SFT on responses with redundant tokens removed). (3) Other RL-Based Methods:
CosEn (Chang et 2025). O1-Pruner (Luo al. 2025a). and Kimi 1.5 RL Team al. 2025).
Evaluation. We assess Pass @ (P@ 1) and average response length (Tokens, Tok) on MATH
500 (Hendrycks et al. 2 AIME 2024, AMC 2023, Olympiad Bench (He al. 2024). and
Minerva (Lewkowycz et al., 2022). We sample 32 samples for each question in AIME24 and 10 samples for others to estimate the P@ 1. Details for evaluation are presented in Appendix
4.2 DIET Achieves Improved Performance with Reduced Tokens
Table
! summmrrzess the performance and token efficiency of our methods against basclines. The Base Model establishes a strong performance benchmark but with substantial verbosity, underscoring the overthinking issue.
Among existing approaches, Kimi 1.5 DPO is strong baseline that slightly reduces tokens and improves the performance over the base model. TokenSkip achieves extreme token reduction, but its average P@ drops significantly. demonstrating that aggressive, non-nuanced compression severely degrades reasoning. Standard RL baselines generally achieve more significant token reduction, however, the performances are slightly inferior to the base model. These methods highlight the difficulty of achieving both high performance and low token counts simultaneously without more sophisticated adaptation.

---

Table Average performance @ 1. %) and token length of RI-Distilled Qwen 1.5B trained with different token compression methods. For each benchmark, highest P@] I is bolded. For the "Macro Average" columns, bolding indicates the best P@1 and a favorable performance- trade-off.
MATH 500
AIME 2024
AMC 2023
Olymplud.
Minerva
Macro Average
Method
P@1
Tok
P@1
Tok
P
Tok
P681
Tok
P61 1
Tok
P60 1
Tok
Base Model
82.1
5534
28.5
16590
62.7
10615
43.5
11587
26.0
7076
48.6
10250
SFT- & DPO-Based
Kimi 1.5 SFT
68.5
6761
22.0
17400
60.4
9323
39.4
10036
23.6
2804
42.7-19.13 19.7-1
9865
-55
Kimi DPO
83.3
13389
8678
9604
49.9 -2.7%
8441
-17.0%
TokenSkip
64.1
1120
6.8
2231
37.3
1401
25.8
2061
20.7
1674
30.9
1697
RL-Based CosEn
75.6
2735
27.5
12492
61.1
6970
42.9
8307
27.1
3485
46.8
6798-231.9%
Ol-Pruner
79.1
2531
25.0
8961
62.5
5010
39.0
5242
23.7
2400
45.9
4829 1339--.3.0
Kimi 1.5 RL
66.3
1552
18.8
9109
44.7
3808
28.5
4774
16.7
1009
35.0
4050 4050--01.6%
Our Difficuly-Awarre Methods
Dynamic Target (33.1.2
2792
10288
63.4
6017
6490
48.6 .0.0%
5657
-5.10%
Adoptiive Weighting
N2.7
2876
32.2
10255
64.4
5819
43.7
6:494
26.6
3170
49.9-2.85
S723
DIET
83.0
3061
31.8
10578
65.4
6425
43.7
6917
26.9
3505
50.2.3..8 u3.
6097
-00779
Our Dificculty-Aware Methods consistently yield a better performance-fficcenccy frontier. The
Dynamic Target approach maintains the Base Model's Macro Average P@ I while reducing average tokens by substantial 45.0%. The Adaptive Weighting method further improves the Macro Average P@I 1 to 49.9% (+2.8% over Base), matching Kimi 1.5 DPO's performance but with considerably fewer tokens. Significantly, DIET, which leverages both Dynamic Target and Adaptive Weighting, achieves the best P@I I. This performance is achieved with an average token count of only 6097, significant 40.7% reduction compared to the Base Model. Beyond these quantitative gains, qualitative analysis (Appendix FF) further reveals that DIET training progressively refines the model's reasoning style, leading more structured language, concise calculations, and marked reduction in unnecessary self-doubt and redundant post-solution exploration.
These results underscore the effectiveness of incorporating nuanced difficulty awareness. While methods like TopenSkip can produce very short answers, they do so at an unacceptable performance drop. Our difficulty-wware methods, especially the combined strategy, demonstrate that it is possible to achieve substantial token reductions while maintaining. and even enhancing the sophisticated reasoning capabilities of the LLM, leading to superior performance- -efermaanc-eficcinney trade-off.
4.3 DIET's Advantage in Inference Scaling
An overlooked benefit of token compression its potential to enhance inference scaling performance under a fixed total token budget. Shorter responses al- low for more samples to be drawn for techniques like majority voting, which can improve overall aceuracy if per-sample quality is maintained. While previous token compression work typically focused on single- sample token counts and Pass @ I, we investigate this
Majortty Votiig Accuracy
0.55
0.50
Kimi DPO
TokenSkip
CoiIi
0.45
01 Pruner
Kimi RL
0.40
Adaptive Weighting
inference scaling behavior. We show that many ex- isting compression methods fail to translate token savings into improved scaling performance, whereas DIET achieve superior majority voting accuracy, par- ticularly at practical, lower token budgets.
Dynamic Target DIST
10000
20000
60800 70000
Figure Micro average of majority voting Pass @ on all the benchmarks.
Fig. 3 plots the micro average of majority voting Pass@1 1 across all math benchmarks against the total token budget. The results highlight the advantages of our approaches. Methods like TokenSkip and Kimi SFT. despite allowing many samples due extreme compression, exhibit low and quickly stagnating majority voting accuracy, confirming that their severe per-sample quality degradation undermines scaling benefits. While stronger baselines like Kimi DPO show more respectable scaling, their increase in the scaling performance is slow. The Base Model itself requires a substantial token budget before multiple samples yield significant gains.
In contrast, our difficulty-awar methods demonstrate advantages. Notably, DIET and Dynamic Target achieve significantly higher majority voting accuracy at low token budgets. This demonstrates that DIET effectively preserves per-sample quality, allowing the benefits of increased sampling to manifest carly. Although the final accuracies are similar to other basclines, the faster convergence
7

---

Setting
1930
Kini 1.
Adupiiee Neight
3455
Battor Performance
Loss Tokan
0.48
Tarper
Waighingg Mathod
Training
0.44
Rewwrr Weighting
Advartage Weighting
Pass
1.40
Advantage Weighting
0.42
1420
0.40
Weightigg
0.38
Acvantage Weighting improvee Performance Same Comprisstio Rodoo
2330
Kim 1.
Adaplis Moagring
1325
Token
Training
Aduuntage Weigting
Whipiting
3270
2000
2500
3000
4000
4500
5000
3500
Average
Length
Average Responee Length
Figure 5: Advantage Weighting vs. Reward Weighting analysis. (Left) Training curves (Pass vs. Response Length) demonstrate better performance with Advantage Weighting. (Right) Final evaluation results show Advantage Weighting yields superior performance-efficienny points.
makes its scaling more practical. Overall, the results show that our methods effectively translate token savings into improved inference scaling performance, which is crucial for practical applications.
4.4 DIET Enhances Length-Difficulty Correlation
An ideal token compression method should not only reduce verbosity but also preserve, or even enhance, intelligent allocation tokens based on problem difficulty, a natural tendency observed in base LLMSs ($2 2.33.. Uniform compression risks disrupting this by being overly terse on complex problems or insufficienlyy concise on simple ones. We investigate this by measuring the Pearson correlation between estimated problem difficulty and generated response length across test benchmarks; higher positive correlation indicates more appropriately scaled verbosity.
Fig. 4 demonstrates that our difficulty-awas aware methods ex- cel at maintaining and enhancing this crucial correlation. Our method, DIET, achieves the highest correlation, sur- passing the Base Model's inherent correlation (dashed horizontal line) and other baselines. This indicates that DIET successfully learns modulate verbosity tight alignment with problem difficulty, consistent with our de- sign goals. Notably, another strong baseline, CosFn, also shows high correlation, nearly matching DIET, but shown in $4..2 it fails to preserve performance when re- ducing tokens. In contrast, other compression techniques significantly degrade this adaptive characteristic. For in- stance, Kimi DPO and Kimi SFT show markedly weaker correlation than the Base Model, implying their compres- sion mechanisms are less sensitive to problem difficulty. O1 Pruner also shows a reduced correlation. This suggests that while these methods reduce tokens, they do in more uniform or digficult--anostic manner.
0.500
0.900
8.875
0.25
0.800
0.750
DMET
Adauttvv Heighttin
AL
af apo
Aoni SPF
Dynansic
Settings
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.
These results underscore that intelligent compression, achieved by DIET, is not merely about token reduction but about doing so in a way that respects problem complexity. By strengthening the positive relationship between difficulty and response length, DIET ensures a more rational allocation of computational budget during inference, contributing to its robust performance-efficiene
4.5 Ensuring Stable Difficulty-Aware RL Training with Advantage Weighting
We empirically validate our proposed Advantage Weighting method ($3.2). designed to overcome the signal distortion pitfalls of naive reward weighting within normalized policy gradient algorithms like GRPO. As highlighted in our analysis ($3.2), naive weighting can interact poorly with outcome variance, hindering effective penalty application.
Fig. 5 (left) plots the training dynamics (Average Training Pass @ vs. Average Response Length). Advantage Weighting (solid lines) consistently maintains higher performance than Reward Weighting
8

---

(dashed lines) as models compress responses (moving right-to-lefi during training). This demonstrates more effective learning during token reduction across all penalty settings when using Advantage Weighting. The final evaluation results (Fig. 5 right) further reinforce this. Models trained with Advantage Weighting achieve superior performance-effiienn
trade-offs, occupying the desirable top-left region (high performance, low token count). The DIET without advantage weighting fails to deliver reasonable performance, we therefore omit it from the visualization.
These empirical results strongly support our theoretical analysis ( Advantage Weighting is crucial for the stable and effective implementation of weighted objectives (like difficulty-awer aware penalties) in normalized PG algorithms. Naive reward weighting leads to suboptimal training and significantly poorer final performance-effcient, outcomes.
5 Related Work
Efficient reasoning in LLMs has recently attracted significant attention, as methods that boost reasoning performance, often via RL with outcome- based rewards, unintentionally induce verbose, overthought outputs. A growing body of work has therefore aimed at compressing the reasoning to reduce inference costs without greatly compromising accuracy.
Prompt-Based Approaches. Initial efforts primarily explore prompt engineering to reduce ver- bosity. For example, Chain-of-Draft (Xu et al., 2225)) and Sketch-of-Thouge (Aytes et al.) (025) restructure reasoning by having the model first draft a concise outline before finalizing the answer, while Constrained-CoT (Nayab et al. (024) imposes length limits via prompts. Though these methods allow quick, zero-shot adjustments, their effectiveness is limited since they do not alter the model's s internal parameters.
Training-Basee Based Compression via Supervised and RL Methods. More fundamental approaches modify the model's training process to inherently produce more concise reasoning chains. Supervised fine-tuning (SFT) techniques aim to internalize efficient reasoning patterns by training on compressed or optimized CoT data. SPIRIT-FT (Cui et al.) 2025) and Skip-Steps (Liu et al., 2024) train models on reasoning steps deemed crucial. Other SFT approaches focus on distilling longer reasoning chains from capable models into shorter, equivalent ones (Yu et al.] 2024 Kang et al.] 2024 Munkhbat al. 2025). Some methods even train models to reason implicitly in latent space, generating concise outputs without explicit step-by-step textual reasoning, such as Coconut (Hao et al. 2024).
CCoT (Ceng & Durme, 2024), and Implicit-CoT (Deng al. 2024).
Reinforcement learning (RL) offers another avenue, often by incorporating length penalties directly into the reward function. This typically involves combining the primary autcome-based reward (e.g., correctness) with a secondary reward term that penalizes longer sequences. Examples include approaches by Arora & Zanette 2025). L.1 (Aggarwal & Welleck 2025). Kimi-1.5 (Tem al. 2025
and work exploring how length penalties can be used to stabilize training (Chang et al.) 20255 Our work builds these efforts by comprehensively introducing difficulty awareness into the compression process, analyzing the impact of data difficulty on compression training, proposing adaptive reward shaping techniques that dynamically adjust token penalties based on on-the-fly difficulty estimation, and addressing methodological issues in applying such rewards within group-normalized RL methods.
Conclusion
To combat LLM "overthinking" and its inherent inefficiencies, we introduced DIET (DIfficulty- AwarE Training), a framework that intelligently "cuts token calories" " by integrating on-the-fly problem difficulty into the RL process for adaptive compression. DIET achieves satisfactory reasoning performance while significantly reducing token counts. Beyond these primary gains, DIET uniquely preserves and enhances the natural positive correlation between response length and problem difficulty. ensuring appropriate verbosity. Furthermore, it translates these efficiencies into superior inference scaling, delivering better performance under fixed computational budgets, which is a crucial advantage over prior methods that often falter in this regard. DIET thus offers principled, effective, and thoroughly-validate strategy for developing more practical, efficient, and ultimately more capable large language models.

---

Acknowledgement
This work is supported by the National Key R&D Program of China (No.2022ZD01 16312), the National Natural Science Foundation of China (No, 62236004), the high-quality development project of MIIT and a grant from the Guoqiang Institute, Tsinghua University. This work is also supported by the A19Stars community.
References
Pranjal Aggarwal and Sean Welleck. L1: Controlling how long reasoning model thinks with reinforcement learning. arXiv preprint Miiv:2553300699) 2025.
Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzich Fadace. Julia Kreutzer, Olivier Pietquin, Ahmet Üstün, and Sara Hooker. Back basics: Revisiting reinforce-style optimization for learning from human feedback in llms. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), ACL 2024, Bangkok, T'aallan, August 11-16, 2024, pp. 12248-12267. Association for Computational Linguistics, 2024. doi: (0.18663N112222.ACLL-L https://doi.on org 10. 18653/v1/2024, acl- waw.-oog..666
URL
Daman Arora and Andrea Zanette. Training language models to reason efficiently. CoRR. 04442, 2025, doi: 1048550/ARXIV22500.000
URL httpss://oii...o org/ 111/0.00.88///
2r02v.2502.0 04466
Simon A Aytes, Jinheon Baek, and Sung Ju Hwang. thoouht Efficient Ilm reasoning with
adaptive cognitive-inspired sketching. arXiv preprini arXiv:2503.05/79 2025.
Edward Y. Chang. Yuxuan Tong, Morry Niu, Graham Neubig, and Xiang Yue. Demystifying long chain-of-thought reasoning in 1lms. CoRR, abs/2502.00373. 2025, doi: 10.48550/ARXIV.2500. 03373, URL htttp:://dee org 48550/ ariiv. 2502. 03373
Xingyu Chen, Jiahao Xu, Tian Liang, Zhiwei He, Jianhui Pang, Dian Yu, Linfeng Song, Qiuzhi Liu, Mengfei Zhou, Zhuosheng Zhang, Rui Wang, Zhaopeng Tu, Haitao Mi, and Dong Yu. Do NOT think that much for 2+3=? on the overthinking of ol-like llms. CoRR, abs/2412.21187 2024. doi: 10.45550/ARXI.2412222111
URL /ttppsss:///ooi wrweei...or11.11 00.444550/aaaa1 2412. 21187
Jeffrey Cheng and Benjamin Van Durme, Compressed chain of thought: Efficient reasoning through dense representations. CoRR, abs/2412.13171, 2024, doi: 10445500AARXIV.24221311 URL
https://doi. org/10. 100 48550/arXiv 2412. 13171
Yingqian Cui, Pengfei He, Jingying Zeng. Hui Liu, Xianfeng Tang, Zhenwei Dai. Yan Han, Chen Luo, Jing Huang, Zheng Li, Suhang Wang, Yue Xing, Jiliang Tang, and Qi He. Stepwise perplexity- guided refinement for efficient chain-of-thought reasoning in large language models. CoRR, abs/2502. 13260, 2025. doi: 10,48550/ARXIV.55222 13260. URL https://o/// org/ arXiv.2502. 13260
hrg///..100 1////104455555 /
DeepSeek-AI, Daya Guo, Dejian Yang. Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, Xiaokang Zhang, Xingkai Yu, Yu Wu, Z. F. Wu, Zhibin Gou, Zhihong Shao, Zhuoshu Li, Ziyi Gao, Aixin Liu, Bing Xue, Bingxuan Wang. Bochao Wu, Bei Feng, Chengda Lu, Chenggang Zhao Chengqi Deng. Chenyu Zhang, Chong Ruan, Damai Dai, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang, Han Bao, Hanwei Xu, Haocheng Wang, Honghui Ding, Huajian Xin, Huazuo Gao, Hui Qu, Hui Li, Jianzhong Guo. Jiashi Li, Jiawei Wang. Jingchang Chen. Jingyang Yuan, Junjie Qiu, Junlong Li. J. L. Cai, Jiaqi Ni, Jian Liang, Jin Chen, Kai Dong, Kai Hu, Kaige Gao, Kang Guan, Kexin Huang, Kuai Yu, Lean Wang. Lecong Zhang, Liang Zhao, Litong Wang. Liyue Zhang, Lei Xu, Leyi Xia, Mingchuan Zhang, Minghua Zhang, Minghui Tang, Meng Li, Miaojun Wang. Mingming Li. Ning Tian, Panpan Huang, Peng Zhang, Qiancheng Wang. Qinyu Chen, Qiushi Du. Ruiqi Ge, Ruisong Zhang, Ruizhe Pan, Runji Wang, R. J. Chen, R. L. Jin, Ruyi Chen, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shengfeng Ye, Shiyu Wang, Shuiping Yu, Shunfeng Zhou, Shuting Pan, and S. S. Li. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. CoRR, abs/2501 12948, 2025. doi: 10.48550/ARXIIV25000
URL https: : //doi. org//rr./.. 10. 48550/arXio 2501 12948
10

---

Yuntian Deng, Yejin Choi, and Stuart M. Shieber. From explicit cot to implicit cot: Learning to internalize cot step by step. CoRR, abs/2405. 14888, 2024. doi: 10.48550AAR2X122005..555 URL https://doi org/10 48550/ arxiv. 2405. 14838
Benjamin Estermann and Roger Wattenhofer. Reasoning effort and problem complexity: A scaling analysis in llms. arXiv preprint arXiv:2503.15113, 2025.
Kanishk Gandhi, Ayush Chakravarthy, Anikait Singh, Nathan Lile, and Noah D Goodman. Cognitive behaviors that enable self-improving reasoners, or, four habits of highly effective stars. arXiv preprint arXiv:2503.01307 2025.
Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu, Jason Weston, and Yuandong Tian. Training large language models to reason in a continuous latent space. CoRR, abs/2412.06769, 2024, URL https:///doo
doi: (0.4850/ARXI22412.6076
088.404555aiii 2412. 06769
Chaoqun He, Renjie Luo, Yuzhuo Bai, Shengding Hu, Zhen Leng Thai, Junhao Shen, Jinyi Hu, Xu Han, Yujie Huang, Yuxiang Zhang, Jie Liu, Lei Qi, Zhiyuan Liu, and Maosong Sun. Olympiadbench: A challenging benchmark for promoting AGI with olympiad-level bilingual multimodal scientific problems. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2024, Bangkok, Thailand, August 11-16, 2024, pp. ciation for Computational Linguistics, 2024. doi: 211
Asso- URL
https://doi org/10 10. 11653//1//0000
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the MATH dataset. In Joaquin Vansehoren and Sai- Kit Yeung (eds.), Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks 1, NewrIPS Datasets and Benchmarks 2021, December 2021, virtual, 2021. URL hactpp///dassssst-eeeecccaaarspp
nwwrips wwwp:/www...pppppeeeeeeeeee ccppa/ html
Yu Kang, Xianghui Sun, Liangyu Chen, and Wei Zou, C3ot: Generating shorter chain bhain-offthoouee thothtt without compromising effectiveness. CoRR, 116122222.1116.4, 2024. doi: 10.48500/RXXI.2412..
11664. https:////// htt.://..00..00100660//////11111 2412. 11664
Aitor Lewkowycz, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay V. Ramasesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman Solo, Yuhuai Wu, Behnam Neyshabur, Guy Gur-Ari, and Vedant Misra. Solving quantitative reasoning problems with language models. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NewrIPS 2022, New Orleans, LA, USA, November www.:iinssscc...o..e.o cc/ wwtpcc/pppppeeeeeeeeeeeeeeeeeeee files/ wwtpee/spppprreeee (0aper/0022] 2022/
28 December 9, 2022, 2022. Baahh//1abbeefffffo (aaaaaefffccee2000ff000033333333 wrt.ee/eeSeeeoooeceeeeeeeeeaeeae
html
Tengxiao Liu, Qipeng Guo, Xiangkun Hu, Cheng Jiayang, Yue Zhang, Xipeng Qiu, and Zheng Zhang. Can language models learn to skip steps? In Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet, Jakub M. Tomezak, and Cheng Zhang (eds.), Advances in Neural Information Processing Systems 38: Annual Conference on Neural In- formation Processing Systems 2024, NeurIPS 2024, Vancouver, BC, Canada, December cc.
10 15, 2024, 2024. URL /tttp://papers.t pttp:/ppppprsnnipppccppaaaeppree files/ (t1lles/paper//22// .html
Haotian Luo, Li Shen, Haiying He, Yibo Wang, Shiwei Liu, Wei Li, Naiqiang Tan, Xiaochun Cao, and Dacheng Tao. -1-pruner: bength-harronnizing fine-tuning for ol-like reasoning pruning. CoRR, abs/2501. 11000.1127000. 12500, 2025a. doi: 10448550ARKIV.2501.2555 18850/arXiss 2501. 12570
https:///
Michael Luo, Sijun Tan, Justin Wong, Xiaoxiang Shi, William Y. Tang, Manan Roongta, Colin Cai, Jeffrey Luo, Tianjun Zhang, Li Erran Li, Raluca Ada Popa, and Ion Stoica. Deepscaler: Surpassing ol ol-preview with 1.5b model by scaling rl. http:://ppeetyyrrrriiii.S....... sitt
2025b. Notion Blog.

---

Tergel Munkhbat, Namgyu Ho, Seo Hyun Kim, Yongjin Yang, Yujin Kim, and Se- Young Yun. Self-training elicits concise reasoning in large language models. CoRR, 2bbs/2502.20122.. 2025. doi: 1044550AARXIV220022022
org/10, /0/ 48550/arXiv 2502. 20122
Sania Nayab, Giulio Rossolini, Giorgio C. Buttazzo, Nicolamaria Manes, and Fabrizio Giacomelli. Concise thoughts: Impact of output length on LLM reasoning and cost. CoRR, abx/240..19825,
2024. doi: 10.4855/AARXIV.240..1999 19825
URL
OpenAl. Learning to reason with 1mms, 2024. Accessed: 2025-04-06.
URL httts://openaa.coommimg
www.-rassn-wiiee 11ms/
Jiayi Pan, Junjie Zhang, Xingyao Wang. Lifan Yuan, Hao Peng, and Alane Suhr. Tinyzero.
2025. Accessed: 2025-01-24.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Mingchuan Zhang, K. Li,
Y. Wu, and Daya Guo. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. CoRR, abs/2402.03300 2024. doi: 10.4850/ARXXI.2402.63333
URL
https:://oi.... 20./10. 48550/arXi> 2400 2402. 03300
Guangming Sheng, Chi Zhang, Zilingfeng Ye, Xibin Wu, Wang Zhang, Ru Zhang, Yanghua Peng, Haibin Lin, and Chuan Wu. Hybridflow: A flexible and efficient rlhf framework. arXiv preprint arXiv: 2409. 19266, 2024.
Kimi Team, Angang Du, Bofei Gao, Bowei Xing, Changjiu Jiang, Cheng Chen, Cheng Li, Chenjun Xiao, Chenzhuang Du, Chonghua Liao, et al. Kimi k1.5: Scaling reinforcement learning with llms. arXiv preprint 125592 2025.
Qwen Team. Qwq-32b: Embracing the power of reinforcement learning, March 2025. URL https://qwenlnn www.lllm..ittbub.. www.ttt...obooooooeeoeeeeoooeeee wwwpp//ggggg/////////////
Yuyang Wu, Yifei Wang, Tianqi Du, Stefanie Jegelka, and Yisen Wang. When more is less: Under- standing chhii--ff-thoughe length in llms. arXiv preprint arXiv:2502.0724 07200000072662222 2025.
Heming Xia, Yongqi Li, Chak Tou Leong, Wenjie Wang, and Wenjie Li. Tokenskip: Controllable chain-of-thought compression in Ilms. CoRR. abs/2502 12067, 2025, doi: 1044850/ARXIV.2500.
12067. URL https://doi.orrg 48850//ra.85//// 2502. 12067
Silei Xu, Wenhao Xie, Lingxiao Zhao, and Pengcheng He. Chain of draft: Thinking faster by writing less. CoRR, abs/2502. 18600. 2025. doi: 10.4850/ARXXV.2502..888
URL https
//doi.org/10. 48850//5455// 2502. 18600
Ping Yu, Jing Xu, Jason Weston, and Ilia Kulikov. Distilling system 2 into system 1. CoRR. abs/2407.06023, 2024. doi: 10.45550/RXXV.2407..006
htt/htsss:///oi org/10g 1etg//0.485500 48550/
arXiv.2407.060 06023
A Limitations
While DIET demonstrates significant advancements in creating more token-efficient and performant reasoning LLMs, we acknowledge certain limitations.
Our empirical validation has primarily focused on mathematical reasoning benchmarks. Although these tasks robustly test complex reasoning and verbosity patterns, the generalization of DIET's benefits to other diverse domains warrants more extensive investigation. The optimal balance between conciseness and necessary verbosity might vary across these different applications. Still, we think that mathematical reasoning is representative reasoning task that can be used to validate the effectiveness of our methods.
Furthermore, while the principles of our difficulty- framework are conceptually orthogonal to many existing RL-based token compression techniques, suggesting potential for synergistic combinations, this work did not investigate such hybrid approaches. The empirical exploration of combining DIET with other methods potentially achieve further performance enhancements remains an avenue for future research.
Addressing these areas could lead to even more versatile and efficient large language models.
12

---

B Derivation of Advantage Distortion under Naive Reward Weighting
This appendix provides the mathematical details supporting Remark I in 3.2, demonstrating how naive reward weighting before normalization in algorithms like GRPO distorts the intended effect of penalty term. To isolate the distortion caused by the normalization procedure itself, analyze the case with a constant penalty trade-off parameter, Obusse
The naive approach combines the outcome reward Toutcome, and the penalty term p; for response ys to prompt æ into a single reward:
Toutcome, Obase Ps-
(11)
GRPO then computes the normalized ad advantage based the empirical mean Hr and standard deviation o, these combined rewards {ry})= for the responses sampled for prompt æ:
A -
(12)
Let Moutcumee O Vultoome and Mp: Op be the empirical means and standard deviations of the outcome rewards and penalty terms, respectively, within the batch for prompt æ. Then f4, Moutoome Obase/pp Assuming Tourcome and p are approximately independent given the prompt within the batch, the variance the combined reward is:
of Var(r;) Var(T i ObaseP)) 2 Var(Touteome,, à) ta Var(p,) o tame (13)
Substituting the mean and standard deviation into the advantage calculation:
A
(T ('outcome,
ObussPs)
(Moutcomee
Cbbssepp
6 Catcome + ansc +
(T l'oucoome
Houtcome)
Olbase (ps Hp)
Qutcome + O of +
O'base
(r (Foutcome, Moutcome)
(ps - Ap). (14)
outcome + base op +
outcome
+ a
2
outcome scaling
Effective Penally Scaling: fp
The crucial term is the effecttve penalty scaling factor fp. This factor dictates how strongly the
centered penalty term (pi- - -1p) contributes to the advantage signal and the subsequent policy gradient update VoJ x E[[D log A').
Critically, ip depends the task outcome variance 2 outcome. For binary outcome rewards (correc-
Vincorrect), o² outcome Ć(1 - C), where Č' the estimated correctness (Eq.3) 3 This introduces an unintended dependency on problem difficulty, distorting the effect of the constant hyperparameter
Obase:
Easy/Hard Problems (Ò or 0): In these cases, the outcome variance 2 Cutcome R 0. The
effective penalty scaling becomes fp A3 Vo 1 By (assuming Obase 0). The
intended constant penalty weight Obase is effectively removed by the normalization, and the penalty's influence is scaled only by its own standard deviation, Tp. The hyperparameter no longer controls the penalty strength.
Intermediate Difficulty Problems (Ò n 0.5): Here, the outcome variance o2 outcome is maximal (<< 0.25 for binary rewards). The denominator / ooucome + oo is larger, meaning the effective penalty scaling 9p is minimized. The influence of the constant penalty Qbase is most strongly suppressed precisely when the task difficulty is intermediate.
This analysis reveals that normalizing the combined reward r' distorts the effect of the penalty weight Obwee The interaction with the sffficulty-deendenn outcome variance o 2utcome prevents Obuse from applying consistent pressure. The policy gradient updates do not accurately reflect the intended penalty strength, potentially hindering convergence.
The situation becomes even more complex when using the adaptive weight Pada(Z; (&) Obaas w( C)) from the main paper. If we substitute Qada into Eq. [4 and assume a linear weighting function w(C) = Ć, the effective penalty scaling factor +p becomes a non-linear function of correctness Č
13

---

To visualize this, we can analyze its squared form, p.
0.3
base C
(15)
C(1 Õ) + Buse C
Assuming C + for now, this simplifies to:
0.6
22
base
buse base
(16)
0.4
1-C
o ?
0.2
This relationship highly non-linear. To illustrate, we can take the case where Obase 0.1 and oz = (a plausible assumption the penalty term also nor-
0.0
0.
0.4
0.6
0.8
1.0
malized, as we have done in Eq. 77). The resulting
function for Pp is plotted in Fig. 6 The curve is far from the linear relationship influence is negligible for most of the difficulty range
Figure 6: An example plot of Eq. (16) with
desire; the penalty's
Olbase = 0.1 and o =
(Ô < 0.8) and then increases sharply only as the problem becomes extremely easy. This interaction obscures the difficulty-aware trade-off that would like to intervene.
The Advantage Weighting approach presented in $3.2 avoids all these issues by normalizing the outcome and penalty advantages separately before applying the weight, thus preserving the intended adaptive scaling and ensuring stable, effective training.
C Unbiased Gradient Estimation with Advantage Weighting
In this section, provide formal derivation to show that the core structure of Advantage Weighting method provides an unbiased estimate of the true policy gradient. Our goal is to first derive the unbiased estimator and then explain its connection our practical implementation.
First, we define the optimization objective (0), which combines the task performance reward with the difficulty-awe length penalty:
( EE P.wDDyyrte((ll) [Fivaa (Ir, y)], where Fisall (x, Foutcume (x. y)- a' (x,
mo) p(v)- (17)
Here, Toutcome is the task reward, p(y) is the penalty magnitude, and a' az(x, TT) is the adaptive trade-off parameter, which is treated constant for each group of samples generated for a prompt æ.
According to the Policy Gradient Theorem, the true gradient of this objective is:
( E.,y V log my (y/z) Tioaa (z. y)]
(18)
To reduce variance, one can introduce baseline (zx) that depends on the prompt x but not the specific response y. The gradient estimate remains unbiased because the expectation of the baseline term is zero:
E (yymo(-la) V ())l)) B(x)] = Bxx) (yy-na(lla) Ve log Ta (ylx)]
=B(x)
To To(y 1)) V log To ( (y|z)) dy
VoTo(u(>) dy
B(x)
To(y)
Te(
(19)
B(x)
VoTo(ul>) dy
B(x) V Te Te(ul=)) dy
B(x) (o(1)
14

---

Thus, the true gradient can be equivalently written as:
( = Eu.y (Vg log (o(y) (Zoua( y) Bxx))
(20)
The term A(x, 1/) Vooull (I, 3) (x)) the advantage function. Our Advantage Weighting method constructs specific baseline for this multi-component reward setting, For group of N responses {ys) 1 { to prompt z, we define the empirical means /Muntcame (z) and /1p()) We then construct the
baseline as:
(zz) Moutcome (x) al TO) (2(x))
(21)
Since this baseline only depends group-level statisties for given prompt x it is valid baseline. Substituting this into Eq. (20), the advantage for a sample y; becomes:
Aax y1)
(r Touteome,
a' P+) - (Houttcome al 1p)
(I (Tottcome, - [Moutcome)
a' (x,, TO) (pc 10--14 1bp)
(22)
This is precisely the structure of the advantage used in our proposed Advantage Weighting method, is only that in our practice, we adopt GRPO-style normalization to divide the two terms in the advantage with their standard deviation respectively. The stochastic gradient estimator (:)
Va log ( (z) Axx, yt) is an unbiased estimator of the true policy gradient (f) While the core structure of our method unbiased, our practical implementation follows the common practice of GRPO and normalizes each advantage component by its standard deviation. This final normalization step introduces known bias the gradient's magnitude, but follows the more common practice in LLM RL. One can simply remove the standard deviation denominator in Eq. (8) to obtain an unbiased estimator.
D Training Details
For all baseline methods, we follow the hyperparameter settings reported in their original implemen- tations. For our RL-based methods, in the rollout phase, set the number of rollouts to 8, with top-p value of 0.95, temperature of 0.6, and a maximum response length of 8192 tokens. During the training phase, we Oboss in Eq. 44) to 0.5, half-cycle of Cyclical Compression Pressure 100, kl loss coefficient to 0.001, the learning rate to lc-6, and the batch size to 128.
E Evaluation Detailes
This section of the appendix provides details regarding the evaluation model.
Parameters used for Evaluation: During the evaluation, we employed a temperature of 0.6, a top-p value of 0.95, and a maximum response length of 32,768.
Sample Count for Different Datasets: For the MATH 500, AMC 2023, Olympiad, and Minerva datasets, we adopted a sample count of 10; for the AIME 2024 dataset, we adopted a sample count of 32.
Method of Calculating Pass@1. 1, and Token: For each question, we considered the average accuracy of every sample as Pass@1 1, the average response length of every sample as Token.
Prompt Used in Evaluation: We utilized the prompt "<Question> Let's think step by step and output the final answer within Wboxed }" during the evaluation.
During the Inference Scaling evaluation of mathematical problems, we utilized Python's 's Sympy"mmodul to ascertain the equivalence of two mathematical formulas in LaTeX format. We group the responses with equivalent answers, and select the largest group as the majority voting result. We adopted the LaTeX mathematical formula that appeared most frequently in the k samples after transformation by sympy as the result of Majority Voting. The Inference Scaling Accuracy was computed based on whether the Majority Voting result was equivalent to the Ground Truth.
Inference Scaling:
htttps:/www.... https:///ww.syypy.... org/
15

---

Property Satisfaction over Base Across Training Steps
2.90
rediced, sell doult
loss, exploration
Reletive Setisfaction
0.03
concise colculatiion condis__olcuultiio
structured, Janguage
0.70
0.89
0.50
Step
Figure 7: Evolution of qualitative characteristics during DIET training, showing the average rate at which model checkpoints satisfy each property better than the Base Model.
Qualitative Analysis of Behavioral Changes During Training
To understand how our DIET training qualitatively refines the model's verbosity and reasoning style beyond aggregate token counts, we conducted case study. We aimed to assess specific behavioral changes related to token reduction by defining four qualitative metrics:
The model exhibits less hesitation or redundant self- correction once correct reasoning path is identified.
Reduced Unnecessary Self-Doubt:
2. Reduced Past-Solution Exploration: The model curtails exploration of alternative methods or further elaboration after correct answer has already been found.
3. Improved I Language Structure: The model's output is more organized and flows logically, with fewer digressions or poorly structured sentences.
4. Concise Calculation Process: Mathematical or logical steps are presented more directly and with less intermediate clutter.
For this analysis, we use questions from AMC 2023, AIME 2024, and MATH 500, and compare responses from the base model against those from various checkpoints of our DIET model during its training process. To evaluate the relative improvement on the aforementioned characteristics, we utilized Gemini 2.5 Pro as a judge to determine each checkpoint response demonstrated the targeted behavior when compared with response from the base model.
Fig- 7 illustrates the average "Property Satisfaction Rate" for these four characteristies training progresses. Each curve represents the proportion of cases where the DIET checkpoint was judged superior to the Base Model for that specific characteristic, averaged across datasets.
The trends in Fig. Lindicate that as DIET training advances, the model progressively improves across these qualitative dimensions. We observe clear learning curve where the model becomes more adept at producing concise and well-structured language. Notably, the "Reduced Self-Doubt" characteristic shows the most significant impact, with satisfaction rates reaching approximately 90% by the later stages of training. "Structured Language" and "Concise Calculation" also demonstrate substantial gains, with satisfaction rates plateauing around 85% and 83%, respectively. "Less Exploration" after finding solution also improves steadily, reaching around 80%. These qualitative improvements suggest how DIET shortens the reasoning trajectory.
G Ablation Studies
G.1 Impact of Cyclical Compression Pressure
We analyze the effect of the cyclical compression pressure strategy, detailed in $3.3 by comparing our primary difficulty-awan approaches: Adaptive Weighting, Dynamic Target, and our combined DIET method with and without this temporal modulation.
16

---

MATH 500
AIME 2024
AMC 2023
Adtatio Milgghins nbr Syrl
Aduate Mecghting wihe Ovrl
Maqohe srddi wis Cye!
DSST Cyel
und Oye
wis
Tagul
Spaare Tirget
Syarres Target
MNT
0.6-
0.07
1800
3830
2003
22t0
2403
2000
2000
1230
1933
8300
2803
side 33ice
31500
1888
4883
2000
0043
1853
2003
Olymplad
3.300-
Heighting
0385
Oyd
Sysrres Sogelt
Spurrs Segeet
Dysures Torel
0.49
0.455
0.47
2300
0.10
02555
0.833
1590
4000
5.00
5049
Takan
Figure 8: Pass@1 versus average token count for Adaptive Weighting, Dynamic Target, and DIET configurations with and without eyclical compression, across various benchmarks and their Macro Average. The dashed red line indicates the Base Model's s Pass@ I performance. Cyclical training generally trades slight increase tokens for improved Pass@
Table 2: Macro Average Pass @1 1 (%) and Token Count for the Dynamic Target method with varying trade-off parameter Otasse.
MATH 500
AIME 2024
AMC 2023
Olympiad.
Minerva
Maero Average
Obase
P@I
Tok
P@1 1
Tok
P@1 I
Tok
P@]
Tok
P@1
Tok
P@1
Tok
0.1
83.1
2616
29.7
8852
65.1
5198
35.0
10326
27.1
2951
48.0
5988
0.5
80.1
1546
27.0
6794
63.6
3467
41.7
3814
24.9
1126
47.4
3349
1.0
73.5
1298
16.4
4217
49.5
1966
34.6
2475
22.2
1016
39,2
2194
Fig. 8 iilluttrate these comparisons. Consistently across the benchmarks, particularly in the Macro Average (bottom right), applying cyclical compression shifts the methods higher Pass@ com- pared to their non-cyclical counterparts. This performance improvement often allows the cyclical variants meet or exceed the Base Model's Pass @ (dashed red line).
This gain in reasoning accuracy typically corresponds to a moderate increase in average token length. For example, in the Macro Average plot, the DIET configuration achieves the highest Pass@) 1, while DIET w/o Cycl uses fewer tokens but results the lowest Pass This pattern suggests that while non-cyclical versions offer more aggressive token reduction, the "relax" phases in cyclical training are beneficial for achieving peak performance, justifying its inclusion in our best-performing DIET configurations presented in Table Thus, cyclical compression aids difficulty- methods in achieving superior balance of high performance and significant token savings relative to the Base Model.
G.2 Impact of Trade-off Parameter Chass
To determine appropriate value for the trade-off parameter Obase in Eq. I we conducted an ablation study on the Dynamic Target as a preliminary experiment. This preliminary experiment aimed to find a balance between maintaining reasoning performance and achieving significant token reduction. We tested Oiase E {0.1, 0.5, 1.0), with results shown in Table
As shown in Table increasing Chase leads more aggressive token reduction but also correspond- ing decrease Pass performance. Specifically, on Macro Average, increasing Ohass from 0.1 to 1.0 reduces tokens significantly from 5988 to 2194, but Pass @ @ drops from 48.0% to 39.2%. The
17

---

Table Macro Average Pass @ (%) and Token Count for the ablation of GRPO-style normalization
MATH 500
AIME 2024
AMC 2023
Olympiad.
Minerva
Macro Average
P@ 1
Tok
P@ 1
Tok
P@1 1
Tok
P@1
Tok
P@ 1
Tok
P@1
Tok
W Norm
83.0
3061
31.8
10578
64.5
6425
43.7
6917
26.9
3505
50.2
6097
W/o Norm
82.5
2856
29.0
9994
66.5
5724
43.8
6384
26.9
3070
49.7
5605
Table 4: Macro Average Pass @ 1 (%) and Token Count for the Dynamic Target method with varying rollout parameter N.
MATH 500
AIME 2024
AMC 2023
Olympiad.
Minerva
Macro Average
N
P@1 1
Tok
P@]
Tok
P@
Tok
P@
Tok
P@1
Tok
P@1 1
Tok
4
78.1
2104
29.1
8709
63.8
4993
42.7
5298
24.8
2270
47.7
4674
8
83.0
3061
31.8
10578
64.5
6425
43.7
6917
26.9
3505
50.2
6097
16
82.2
2613
28.3
9899
66.6
5328
44.0
5970
26.3
2834
49.5
5328
setting of Obose 0.5 yields satisfactory trade-off. Based on this balance, we selected Otbose 0.5 the default for all the difficulty-awy aware training in main experiments.
G.3 Impact of Normalization Within Advantage Weighting
To test the importance of normalizing the penalty term within our Advantage Weighting framework. we ran an ablation where we removed the GRPO-style normalization from the penalty advantage. The results shown on Table
As the results show, removing normalization leads substantial token reduction but also causes significant degradation in performance. Some careful tuning of the weighting parameter might mitigate the issue, but normalization is an effective and simple approach. Therefore, we adopt the normalization. We hypothesize that the un-normalized penalty can have extreme values, which overly biases the policy gradient towards generating shorter sequences at the expense of correctness.
G.4 Impact of Rollout Parameter N
an important hyperparameter of GRPO, to verify its impact on DIET, we conducted an ablation study, and the results are shown Table
The results show a clear trend. Decreasing the sample count to N=4 results in noticeable drop in performance. We attribute this to the less stable estimation of problem difficulty and reward baselines, which can introduce noise into the policy gradient updates. Increasing the sample count from N=8 to N=16 does not yield further performance improvements. However, it lowers the number of tokens, potentially because the difficulty estimation more accurate, allowing for more token reduction for problems with suitable Cefficulty.Cosiddee the trade-off between performance and efficiency, we select N=8 in our main experiments.
G.5 Impact of Dynamic Length Target Paramter Tymax and
We also conducted ablation study for Lmax and 8, and the results are shown in Table and Table 6 respectively.
Reducing Iynax decreases average response length with modest performance degradation, similar to decreasing 8. We chose the current configuration as the optimal balance between performance and token efficiency.
H Scale to Larger Models
To verify the scalability of the DIET, we also conducted experiments on the R1-Distilled Qwen 7B. The experimental results are shown in Table From the results, we draw two key conclusions:
18

---

Table Macro Average @ I (%) and Token Count for the Dynamic Target method with varying
Lmaz.
MATH 500
AIME 2024
AMC 2023
Olympiad.
Minerva
Macro Average
Lmaz
P@1
Tok
P@1 1
Tok
P@1
Tok
P@1 1
Tok
P@1
Tok
P@
Tok
2048
82.9
2851
28.8
9762
64.8
5597
43.0
5742
29.0
2973
49.7
5385
8192
83.0
3061
31.8
10578
64.5
6425
43.7
6917
26.9
3505
50.2
6097
Table 6: Macro Average Pass 1 (%) and Token Count for the Dynamic Target method with varying
MATH 500
AIME 2024
AMC 2023
Olympiad.
Minerva
Macro Average
8
P@1 1
Tok
P@ 1
Tok
P@1
Tok
P@1
Tok
P@1
Tok
P@
Tok
0.05
82.6
2903
28.3
7630
65.0
5907
43.4
6710
26.3
3303
49.1
5290
0.1
83.0
3061
31.8
10578
64.5
6425
43.7
6917
26.9
3505
50.2
6097
DIET Improves Performance while Reducing Tokens: Our DIET framework achieves the highest average Pass score (65.4%), improving upon the already strong 7B base model (64.4%). Crucially, simultaneously reduces the average token count This demonstrates that our method effectively enhances both performance and efficiency at larger scale.
Superior Cefformanne-ffficiinnn Trade-off: Most other compression methods exhibit clear trade-off, where significant token reduction leads to noticeable drop in performance. DIET breaks this trend.
I Functional Form of Dynamic Length Target
As illustrated in Fig. 2 the relationship between difficulty and token count follows more logarithmic trend rather than linear one. To explore this, we conducted this experiment. We first fitted logarithmic curve to the data in Fig. 2 then scaled and moved the curve to ranges from 0 to 8192 when is in [0, 1], resulting in = 3116.9 9 (n.21.332 + 1.66) 1579.7. We then used this function to sample target lengths in our Dynamic Target method, instead of sampling from a linear range. The experimental results are shown in Table g
The results show that the method is robust to the selection of the target function, the two choices do differ significantly, and the simple linear function already works well.
19

---

Table 7: Average performance and token length of RI-Distilled Qwen 7B trained with RL-based
methods
MATH 500
AIME 2024
AMC 2023
Olympiad.
Minerva
Macro Average
Method
P@1
Tok
P61
Tok
P 1
Tok
P61
Tok
P@1
Tok
P@1
Tok
Base Model
92.1
3921
53.5
13389
82.5
7730
56.4
8890
4930
64.4
7772
RL-Based
CosPn
87.1
1397
50.5
8659
78.3
3578
54.0
4593
36.4
1345
61.3
3914
O1-Pruner
71.1
4958
39.2
12054
79.1
6014
50.9
6852
30.5
5277
54.2
Kimi 1.5
64.0
1124
44.3
7903
70.6
2783
46.1
3225
32.5
677
51.5
3142
Our Difleuuty-Aware Methods
Dynamic Target (...222
90.4
2121
49.6
8186
79.8
4416
54.3
4677
34.5
1819
61.7
4244
Adaptive Weighting 5
90.6
1782
52.1
8313
80.2
4050
54.9
4838
36.9
1638
62.9
4124
DIET ( 3.1.1 3.1.2
92.1
3187
57.9
10124
82.6
6075
56.5
7026
37.9
3695
65.4
6021
Table 8: Macro Average Pass @1 (%) and Token Count for different length penalty functions of Dynamic Target
MATH 500
AIME 2024
AMC 2023
Olympiad.
Minerva
Macro Average
P@
Tok
P@1
Tok
P@1
Tok
P@1 1
Tok
P@0
Tok
P@
Tok
Logarithmic
83.5
2971
29.8
10890
61.5
6526
43.7
6693
26.5
3319
49.0
6079
Linear
83.0
3061
31.8
10578
64.5
6425
43.7
6917
26.9
3505
50.2
6097
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We have outlined our main contributions in point-by-point manner in both the abstract and the instruction sections.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
In Appendix we provide an analysis of the current limitations of our work.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
20

---

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We provide derivation of advantage distortion under naive reward weighting
in Appendix B
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide test details Appendix E We will also open-source our evaluation code after proper organization.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same
21

---

dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis-
sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (c.g.. ,, a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-souree models, it may be that access to the model is limited in some way (e.g.. registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification:
We will open-source our code after proper organization.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
CooeeSubmissionooliiee
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips.
ssinnoliccy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: We provide detailed training details in Appendix D and detailed evaluation details in Appendix
22

---

Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We conducted a p-value test when computing the Pearson correlation between problem difficulty and average response length and ensured that the p- p-value is less than 0.01. And we reported error bars in Fig.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots. The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We detail the computational resources used in our experiments in 4.1
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics | http:///wwwappppppppppp/pcaaaaaa
23

---

Answer: [Yes]
Justification: Our research adheres to the NeurIPS Code of Ethics. It does not involve human subjects, personally identifiable information, sensitive data.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our experiments are conducted using the already widely-used open-source model, which will not induce new societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We use open-source reasoning model for our experiments, which will pose no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
24

---

We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All external assets used in our work, including code(e.g.. verl, sympy), datasets(e.g.. deepscaler), and models(e.g. RI-Distilled Qwen), are properly credited. We ensure that the licenses and terms of use for these resources are explicitly stated and strictly followed.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, papppswitthhaoooo thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is available online, the authors encouraged to reach out the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We will open-source code and provide detailed documentation and com- ments for ease of use.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Crowdsourcing and human subjects are not involved in our research.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
25

---

Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Crowdsourcing and human subjects are not involved in our research. Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing. editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: We only used LLM for writing and editing which does not impact the core methodology.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy htappp///eerrpp.cc/ooeeffeeeeee for what should or should not be described.
26


Reviews:
[{"review_id": "KtAhzdTH5b", "reviewer": "Reviewer_zKPN", "summary": "Training 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", "questions": "- 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", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "The 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", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the valuable feedback and for acknowledging the importance of our work, the clarity of our proposed solutions, and the effectiveness of our ablation studies.\n\nBelow, we respond to each of the questions and weaknesses identified.\n\n---\n\n### **1. On the Comparison with L1**\nWe view our work and L1 as addressing different and complementary aspects of the token compression problem.\n\n- **Different Focus:** L1's primary contribution is a novel RL reward designed to control response length. In contrast, our work, DIET, focuses on a different axis: proposing a framework for integrating on-the-fly problem difficulty into the length-controlled training process. Our main contributions are methodologies for how to apply **any** length-related penalty, theoretically including L1’s penalty, in a more nuanced, difficulty-aware manner.\n- **Potential for Synergy:** This distinction makes our approaches highly compatible. Our DIET framework is designed to be modular: actually, the current implementation of Adaptive Weighting can partly be viewed as applying our difficulty-aware techniques to the penalty function from Kimi RL to demonstrate DIET's effectiveness. **One could just as easily replace the penalty term in DIET with the L1 reward**. DIET would then modulate this L1 penalty based on task difficulty, and our Advantage Weighting technique would ensure its stable incorporation into the GRPO update step. **This would effectively create a \"Difficulty-Aware L1\" with minimal modification.**\n\nAs L1 is a **concurrent and complementary** work to DIET, it was not included in our original set of baselines. We chose to demonstrate DIET's effectiveness on the Kimi penalty as a representative example of modern length-control rewards, and compare with methods that do not rely on designing a length-related RL reward. We agree that formally exploring the synergy between DIET and L1 is a compelling direction for future work and will add this discussion to our paper. Also, we will make the above claim more explicit in our revision.\n\n---\n\n### **2. Difference from GRPO and Unbiased Gradient Estimation**\n\nThank you for these critical questions on the nature of our advantage calculation.\n\n- **Clarification of our Contribution:** To be clear, we do not propose a new advantage function to replace GRPO's. The Advantage Weighting technique is a method for correctly structuring a multi-component reward (i.e., outcome + penalty) *before* it is fed into an RL algorithm like GRPO. Our analysis in Appendix B shows that naively combining rewards leads to signal distortion during GRPO's normalization step. Advantage Weighting prevents this by normalizing each component's advantage separately, ensuring the final weighted combination is faithful to the intended objective.\n- **Unbiased Gradient Estimation:** This is an excellent point. While GRPO's advantage normalization biases gradient magnitude, we add a proof that our Advantage Weighting (without std normalization) yields unbiased policy gradient estimates. The remaining bias matches GRPO's inherent limitation and can be eliminated by removing its std normalization.\n\nHere's the major part of the added proof in our manuscript (refresh the page if it is not correctly rendered):\n\n> [...]\n> \n> First, we define the optimization objective $J(\\theta)$, which combines the task performance reward with the difficulty-aware length penalty:\n> \n> $$\n> J(\\theta) = \\mathbb{E}\\_{x \\sim \\mathcal{D}, y \\sim \\pi\\_{\\theta}(\\cdot|x)} \\left[ r\\_{\\text{total}}(x, y) \\right],\n> \\quad \\text{where} \\quad\n> r\\_{\\text{total}}(x, y) = r\\_{\\text{outcome}}(x, y) - \\alpha'(x, \\pi\\_{\\theta}) \\cdot p(y).\n> $$\n> \n> Here, $r\\_{\\text{outcome}}$ is the task reward, $p(y)$ is the penalty magnitude, and $\\alpha'(x, \\pi\\_{\\theta})$ is the adaptive trade-off parameter, which is treated as a constant for each group of samples generated for a prompt $x$.\n> \n> According to the Policy Gradient Theorem, the true gradient of this objective is:\n> \n> $$\n> \\nabla\\_{\\theta} J(\\theta) = \\mathbb{E}\\_{x, y} \\left[ \\nabla\\_{\\theta} \\log \\pi\\_{\\theta}(y|x) \\cdot r\\_{\\text{total}}(x, y) \\right].\n> $$\n> \n> To reduce variance, one can introduce a baseline $B(x)$ that depends on the prompt $x$ but not the specific response $y$. The gradient estimate remains unbiased because the expectation of the baseline term is zero:\n> \n> $$\n> \\mathbb{E}\\_{y \\sim \\pi\\_{\\theta}(\\cdot|x)} \\left[ \\nabla\\_{\\theta} \\log \\pi\\_{\\theta}(y|x) \\cdot B(x) \\right]\n> = B(x) \\cdot \\int\\_{y} \\pi\\_{\\theta}(y|x) \\nabla\\_{\\theta} \\log \\pi\\_{\\theta}(y|x) dy\n> = B(x) \\cdot \\int\\_{y} \\nabla\\_{\\theta} \\pi\\_{\\theta}(y|x) dy\n> = B(x) \\cdot \\nabla\\_{\\theta} \\int\\_{y} \\pi\\_{\\theta}(y|x) dy\n> = B(x) \\cdot \\nabla\\_{\\theta}(1)\n> = 0\n> $$\n> \n> Thus, the true gradient can be equivalently written as:\n> \n> $$\n> \\nabla_{\\theta}  J(\\theta) = \\mathbb{E}\\_{x, y} \\left[ \\nabla_{\\theta} \\log \\pi_{\\theta}(y|x) \\cdot (r_{\\text{total}}(x, y) - B(x)) \\right].\n> $$\n> \n> The term $A(x,y) = r_{\\text{total}}(x, y) - B(x)$ is the advantage function. Our Advantage Weighting method constructs a specific baseline for this multi-component reward setting. For a group of $N$ responses $\\\\{y\\_i\\\\}\\_{i=1}^N$ to a prompt $x$, we define the empirical means $\\mu\\_{\\text{outcome}}(x)$ and $\\mu\\_p(x)$. We then construct the baseline as:\n> \n> $$\n> B(x) = \\mu_{\\text{outcome}}(x) - \\alpha'(x, \\pi_{\\theta}) \\cdot \\mu_p(x).\n> $$\n> \n> Since this baseline only depends on group-level statistics for a given prompt $x$, it is a valid baseline. Substituting this into the gradient equation, the advantage for a sample $y_i$ becomes:\n> \n> $$\n> \\begin{aligned}\n> A(x, y_i) = \\left( r_{\\text{outcome}, i} - \\alpha' p_i \\right) - \\left( \\mu_{\\text{outcome}} - \\alpha' \\mu_p \\right)\n> = (r_{\\text{outcome},i} - \\mu_{\\text{outcome}}) - \\alpha'(x, \\pi_{\\theta}) \\cdot (p_i - \\mu_{p}).\n> \\end{aligned}\n> $$\n> \n> This is precisely the structure of the advantage used in our proposed Advantage Weighting method. The only difference is that in our practice, we adopt GRPO-style normalization to divide the two terms in the advantage with their standard deviation respectively (Eq.8), which follows common practice, but introduces bias. Theoretecally, the stochastic gradient estimator $\\hat{g}\\_i(\\theta) = \\nabla\\_{\\theta} \\log \\pi\\_{\\theta}(y\\_i|x) \\cdot A(x, y\\_i)$ is an unbiased estimator of the true policy gradient $\\nabla\\_{\\theta} J(\\theta)$ if we remove the standard deviation denominator in Eq.8.\n\nWe additionally conducted an experiment using the unbiased version of our method, which removes the std and sequence length normalization, as suggested in Dr.GRPO (we only report the averaged performance due to length limit in the rebuttal):\n\n| Setting | Avg P@1 | Avg Tok. |\n| --- | --- | --- |\n| DIET (in our paper) | 50.2% | 6097.2 |\n| Unbiased DIET | 49.6% | 5033.6 |\n\nWhile the unbiased variant exhibits a modest performance decrease, it achieves a reduced average token count. We will incorporate these experimental results in our discussion and propose this modification as a valuable enhancement to our method. Thank you for your suggestion!\n\n---\n\n### **3. Ablation Studies on Dynamic Variables**\n\nWe appreciate the request for more insight into the ablation study, as you have proposed in your questions 2 and 5.\n\n**Ablation on Dynamic Weight ($\\alpha_{base}$):** We tested several values for the base trade-off parameter under the Adaptive Weighting setting.\n\n| $\\alpha_{base}$ | Avg P@1 | Avg Tok |\n| --- | --- | --- |\n| 0.1 | 48.0% | 5988.6 |\n| 0.5 | 49.9% | 5723.2 |\n| 1.0 | 39.2% | 2194.4 |\n\nWe chose $\\alpha_{base}=0.5$ for our main experiments as it provided the best balance between significant token reduction and high performance.\n\n**Ablation on Dynamic Function:** Indeed, as shown in Figure 2, the relationship follows a logarithmic pattern rather than a linear one. Building on this observation, we conducted additional experiments where we sampled target token lengths from a fitted logarithmic function instead of the original linear function.\n\nTo derive an appropriate logarithmic sampling function, we first fitted the model $y = a*ln(bx + c)$ to the data in Figure 2, then scaled and added a bias term to normalize the output range to $[0, 8192]$ for $x \\in [0, 1]$, resulting in: $y = 3116.9ln(21.33x + 1.66) - 1579.7$. During training, we randomly sample x from $[\\hat{D}-\\delta, \\hat{D}]$ and use the corresponding y value as the target token length. The results are as follows:\n\n| Setting | Avg P@1 | Avg Tok |\n| --- | --- | --- |\n| Linear (in our paper) | 50.2% | 6097.2 |\n| Logarithmic | 49.0% | 6079.8 |\n\nThe results show that the method is robust to the selection of target function, the two choices do not differ significantly. This is a valuable ablation that highlights the impact of the target function's form, and we will include it in our appendix.\n\n**Abaltion on $\\delta$ and $L_{max}$ in $t(x,\\pi_{theta})$**\n\nWe added ablation study for $\\delta$ and $L_{max}$, the results are as follows:\n\n| Max Tokens | Avg P@1 | Avg Tok. |\n| --- | --- | --- |\n| 4096 | 49.7% | 5385.0 |\n| 8192 (in our paper) | 50.2% | 6097.2 |\n\n| delta | Avg P@1 | Avg Tok. |\n| --- | --- | --- |\n| 0.05 | 49.1% | 5290.6 |\n| 0.1 (in our paper) | 50.2% | 6097.2 |\n\nReducing $L_{max}$ decreases average response length with modest performance degradation, similar to decreasing $\\delta$. We chose the current configuration as the optimal balance between performance and token efficiency.\n\n**Learning Rate Selection:** Following standard practices from recent RL work at comparable scales, we fixed the learning rate to isolate our method's contributions. This ensures fair comparison by eliminating optimizer configuration as a confounding factor.\n\nWe'll include these results in the final version. Your feedback has been invaluable in strengthening our paper's completeness."}, {"review_id": "Xy1wioqO9S", "reviewer": "Reviewer_gh6m", "summary": "This paper proposes a suite of methods that reduce the token length of chains of thought maintaining quality.\nThis 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.", "questions": "- Do you think this approach will scale to larger models?", "limitations": "Yes, it was fine.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- Clarity: The paper is written quite clearly and I was able to understand all the methodology while I believe.\n- 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.\n\nWeaknesses:\n- 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.", "was_revised": true, "final_justification": "With the additional experiments on larger models I think that the paper deserves publication.", "rebuttal": "We sincerely thank the reviewer for your positive feedback on the clarity and originality of our work. We agree that the question of scalability is one of the most important aspects for the practical significance of any new method in this field.\n\n**Our answer to this concern is: Yes, our approach scales effectively to larger models.** To demonstrate this, we have conducted a new set of experiments on the **R1-distilled Qwen 2.5 7B model**. Due to the rebuttal time limit, we present the results from competitive RL-based baselines. DIET not only maintains its effectiveness but excels at this larger scale, achieving state-of-the-art performance while improving token efficiency.\n\n| Model | MATH500 P@1 | MATH500 Tok. | AIME24 P@1 | AIME24 Tok. | AMC23 P@1 | AMC23 Tok. | Olympiad P@1 | Olympiad Tok. | Minerva P@1 | Minerva Tok. | **Macro Avg P@1** | **Macro Avg Tok.** |\n| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n| Base Model (7B) | 92.1% | 3921 | 53.5% | 13389 | 82.5% | 7730 | 56.4% | 8890 | 37.4% | 4930 | 64.4% | 7772 |\n| CosFn | 87.1% | 1397 | 50.5% | 8659 | 78.3% | 3578 | 54.0% | 4593 | 36.4% | 1345 | 61.3% | 3914 |\n| O1-pruner | 71.1% | 4958 | 39.2% | 12054 | 79.1% | 6014 | 50.9% | 6852 | 30.5% | 5277 | 54.2% | 7031 |\n| Kimi 1.5 RL | 64.0% | 1124 | 44.3% | 7903 | 70.6% | 2783 | 46.1% | 3225 | 32.5% | 677 | 51.5% | 3142 |\n| Adaptive W. | 90.6% | 1782 | 52.1% | 8313 | 80.2% | 4050 | 54.9% | 4838 | 36.9% | 1638 | 62.9% | 4124 |\n| Dynamic T. | 90.4% | 2121 | 49.6% | 8186 | 79.8% | 4416 | 54.3% | 4677 | 34.5% | 1819 | 61.7% | 4244 |\n| **DIET** | **92.1%** | **3187** | **57.9%** | **10124** | **82.6%** | **6075** | **56.5%** | **7026** | **37.9%** | **3695** | **65.4%** | **6021** |\n\nFrom these new results, we draw two key conclusions:\n\n1. **DIET Improves Performance while Reducing Tokens:** Our DIET framework achieves the highest average Pass@1 score (**65.4%**), improving upon the already strong 7B base model (64.4%). Crucially, it simultaneously reduces the average token count This demonstrates that our method effectively enhances both performance and efficiency at a larger scale.\n2. **Superior Performance-Efficiency Trade-off:** Most other compression methods exhibit a clear trade-off, where significant token reduction leads to a noticeable drop in performance. DIET breaks this trend.\n\nThese results confirms that DIET is not an artifact of small-scale models but is a robust and scalable framework. We will add this new experimental section to the final version of the manuscript. Thank you again for encouraging us to strengthen this crucial aspect of our paper. We hope you find our responses satisfactory and we would be grateful for your consideration in raising the score."}, {"review_id": "deq1PxONYd", "reviewer": "Reviewer_vKd6", "summary": "This 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.", "questions": "1. 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?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "# 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.", "was_revised": true, "final_justification": "Clarifications and additional results. But one weakness is still not very clear.", "rebuttal": "We sincerely thank the reviewer for the thorough and insightful review. The deep dive into our theoretical analysis is valuable and demonstrates a careful reading of our work.\n\nWe have addressed the identified weaknesses below and have conducted new experiments and revised our manuscript to incorporate this valuable feedback.\n\n---\n\n### **1. On the Pitfalls of Naive Reward Weighting in GRPO**\n\nWe thank the reviewer for raising these two important points regarding the analysis in Appendix B.\n\n- **Correlation between $\\alpha_{ada}$ and $r_{outcome}$:** You are correct that $\\alpha_{ada}=\\alpha_{base}\\hat{C}$ correlates with the outcome reward, especially for small N. **However**, our independence assumption in Eq. (13) remains valid within the context of the variance calculation. When computing the variance for a group of N responses to a prompt x, the value of $\\hat{C}$ (and thus $\\alpha_{ada}$) is calculated once for the entire group. It becomes a **fixed scalar constant** for that specific calculation. The variance is then computed over the random variables $r_{outcome,i}$ and $p_i$ within the group. Since $\\alpha_{ada}$ is not a random variable in this context, the identity $Var(r_{outcome,i}−\\alpha_{ada}p_i)\\approx Var(r_{outcome,i})+\\alpha_{ada}^2Var(p_i)$ still holds.\n- **Relationship between $\\hat{C}$ and Effective Penalty Scaling:** You have correctly identified a mistake in our analysis regarding the case where α is adaptive ($\\alpha_{ada}$). **We sincerely thank you for this careful catch**. Our original proof **held for a constant trade-off**, $\\alpha_{base}$, but we directly generalized it to $\\alpha_{ada}$ in our manuscript, which is incorrect.\n    \n    **We have now revised our manuscript** to first present the proof with the constant $\\alpha_{base}$ (where the conclusion that penalty is minimized at $\\hat{C}\\approx 0.5$ holds) and then add a detailed discussion for the more complex $\\alpha_{ada}$ case, which aligns with your observation of monotonic behavior. We have added the following discussion to our manuscript and will include it in the final version:\n    \n    > ...[content from Line 523 to Line 552, with $\\alpha_{ada}$ replaced with $\\alpha_{base}$]\n    > \n    \n    > The situation becomes even more complex when using the adaptive weight $\\alpha\\{\\text{ada}}(x, \\pi\\\\theta) = \\alpha\\{\\text{base}} \\cdot w(\\hat{C})$ from the main paper. If we substitute $\\alpha\\{\\text{ada}}$ into Eq.14 and assume a linear weighting function $w(\\hat{C})=\\hat{C}$, the effective penalty scaling factor $\\hat{\\tau}\\_p$ becomes a non-linear function of correctness $\\hat{C}$.\n    > \n    > \n    > To visualize this, we can analyze its squared form, $\\hat{\\tau}\\_p^2$:\n    > \n    > $$\n    > \\hat{\\tau}\\_p^2 = \\frac{\\alpha\\_{\\text{base}}^2 \\hat{C}^2}{\\hat{C}(1-\\hat{C}) + \\alpha\\_{\\text{base}}^2 \\hat{C}^2 \\sigma\\_p^2}.\n    > $$\n    > \n    > Assuming $\\hat{C} \\neq 0$ for now, this simplifies to:\n    > \n    > $$\n    > \\hat{\\tau}\\_p^2 = \\frac{\\alpha^2\\_{\\text{base}}}{\\frac{1-\\hat{C}}{\\hat{C}} + \\alpha^2\\_\\text{base}\\sigma\\_p^2} = \\frac{\\alpha^2\\_{\\text{base}}}{\\frac{1}{\\hat{C}} - 1 + \\alpha^2\\_\\text{base}\\sigma\\_p^2}.\n    > $$\n    > \n    > This relationship is highly non-linear. To illustrate, we can take the case where $\\alpha\\_\\text{base}=0.1$ and $\\sigma\\_p^2=1$ (a plausible assumption if the penalty term is also normalized, as we have done in Eq.7). The resulting function for $\\hat{\\tau}\\_p^2$ is plotted in the figure above [a functio plot is presented]. The curve is far from the linear relationship we desire; the penalty's influence is negligible for most of the difficulty range ($\\hat{C} < 0.8$) and then increases sharply only as the problem becomes extremely easy. This interaction obscures the difficulty-aware trade-off that we would like to intervene.\n    > \n    > The *Advantage Weighting* approach presented in the implementation section avoids all these issues by normalizing the outcome and penalty advantages separately *before* applying the weight, thus preserving the intended adaptive scaling and ensuring stable, effective training.\n    > \n    \n    Additionally, in our response to Reviewer zKPN, we have formally demonstrated that the sole source of bias in our Advantage Weighting technique originates from the same fundamental limitation present in the GRPO algorithm itself. By incorporating Dr.GRPO's modification to GRPO, we can have an **unbiased gradient estimator**. We kindly refer you to our response to Reviewer zKPN if you are interested. We are deeply grateful to both you and Reviewer zKPN for your insightful contributions, which have significantly strengthened the theoretical foundation of our work.\n    \n    We particularly appreciate your meticulous review in identifying the error in our proof. Your expert feedback has been invaluable in enhancing the rigor of our theoretical analysis through these corrections.\n    \n\n---\n\n### **2. Ablation on Normalizing the Penalty Term**\n\nThis is a good suggestion. To test the importance of normalizing the penalty term within our Advantage Weighting framework, we ran an ablation where we removed the GRPO-style normalization from the penalty advantage. The results are as follows:\n\n|Setting|MATH500P@1|MATH500Tok|AIME2024P@1|AIME2024Tok|AMC2023P@1|AMC2023Tok|OlympiadP@1|OlympiadTok|MinervaP@1|MinervaTok|AvgP@1|AvgTok|\n|---|---|---|---|---|---|---|---|---|---|---|---|---|\n|DIET|83.0%|3061|31.8%|10578|65.4%|6425|43.7%|6917|26.9%|3505|50.2%|6097.2|\n|Eq.5 w/o Norm.|82.5%|2856|29.0%|9994|66.5%|5724|43.8%|6384|26.9%|3070|49.7%|5605.6|\n|Eq.7 w/o Norm.|72.7%|780|17.4%|2646|48.3%|1957|34.6%|1444|24.1%|679|39.4%|1501.0|\n\nAs the results show, removing normalization leads to a substantial token reduction but also causes a significant degradation in performance. Some careful tuning of the weighting parameter might mitigate the issue, but normalization is an effective and simple approach. Therefore, we adopt the normalization. We hypothesize that the un-normalized penalty can have extreme values, which overly biases the policy gradient towards generating shorter sequences at the expense of correctness. We will add it to our ablation study.\n\n---\n\n### **3. Impact of Scaling N**\n\nThis is another important hyperparameter. We selected N=8 based on preliminary experiments. As requested, we have presented experiments for N=4 and 16.\n\n|N|MATH500P@1|MATH500Tok.|AIME24P@1|AIME24Tok.|AMC23P@1|AMC23Tok.|OlympiadP@1|OlympiadTok.|MinervaP@1|MinervaTok.|AvgP@1|AvgTok.|\n|---|---|---|---|---|---|---|---|---|---|---|---|---|\n|4|78.1%|2104|29.1%|8709|63.8%|4993|42.7%|5298|24.8%|2270|47.7%|4674.8|\n|8 (in our paper)|83.0%|3061|31.8%|10578|65.4%|6425|43.7%|6917|26.9%|3505|50.2%|6097.2|\n|16|82.2%|2613|28.3%|9899|66.6%|5328|44.0%|5970|26.3%|2834|49.5%|5328.8|\n\nThe results show a clear trend. Decreasing the sample count to N=4 results in a noticeable drop in performance. We attribute this to the less stable estimation of problem difficulty and reward baselines, which can introduce noise into the policy gradient updates. Increasing the sample count from N=8 to N=16 does not yield further performance improvements. However, it lowers the number of tokens, potentially because the difficulty estimation is more accurate, allowing for more token reduction for problems with suitable difficulty.\n\nConsidering the trade-off between performance and efficiency, we select N=8 in our main experiments. We will add it to the ablation study in our paper.\n\n---\n\n### **4. Functional Form of Difficulty vs. Token Budget**\n\nThis is an insightful observation. Indeed, as illustrated in Figure 2, the relationship between difficulty and token count follows a more logarithmic trend rather than a linear one. To explore this, we conducted a new experiment. We first fitted a logarithmic curve to the data in Figure 2, then we scaled and moved the curve to ranges from 0 to 8192 when x is in [0, 1], resulting in $y=3116.9*ln(21.33x+1.66)-1579.7$ We then used this function to sample target lengths in our Dynamic Target method, instead of sampling from a linear range. The experimental results are as follow:\n\n|Setting|MATH500P@1|MATH500Tok|AIME2024P@1|AIME2024Tok|AMC2023P@1|AMC2023Tok|OlympiadP@1|OlympiadTok|MinervaP@1|MinervaTok|AvgP@1|AvgTok|\n|---|---|---|---|---|---|---|---|---|---|---|---|---|\n|Linear (in our paper)|83.0%|3061|31.8%|10578|65.4%|6425|43.7%|6917|26.9%|3505|50.2%|6097.2|\n|Logarithmic|83.5%|2971|29.8%|10890|61.5%|6526|43.7%|6693|26.5%|3319|49.0%|6079.8|\n\nThe results show that the method is robust to the selection of the target function, the two choices do not differ significantly, and the simple linear function already works well. We will include it in our ablation study.\n\n---\n\n### **5. Clarification on Inference Scaling**\n\nThank you for the suggestion. We have revised Figure 3 to use a log scale, which provides a clearer comparison. Regarding the \"faster convergence\" claim, our intention was to compare methods that achieve **comparable final performance**. While TokenSkip's curve is initially steep, its accuracy quickly plateaus at a significantly lower level than the Base Model or our methods. We argue that rapid convergence to suboptimal performance does not constitute meaningful efficiency, as these token savings come at the expense of significantly degraded reasoning capability.\n\nOur key claim is that DIET achieves high accuracy more efficiently – that is, with a lower total token budget – compared to other high-performance methods. The log-scaled plot included in the revised manuscript provides clearer visualization of this advantage.\n\n---\n\nWe hope our rebuttal has adequately addressed your concerns. Incorporating your valuable suggestions has significantly strengthened and refined our paper. We would be deeply grateful if you could consider revising your assessment of our work accordingly. Thank you again for your efforts."}]

Author rebuttal:

Meta-review: This paper proposes DIET (DIfficulty-AwarE Training), an RL framework that reduces excessive token usage in reasoning models by adapting token compression to task difficulty. The method introduces dynamic trade-offs for length control and an Advantage Weighting scheme to stabilize optimization under GRPO. Experiments show DIET improves performance–efficiency trade-offs and enhances inference-time scaling, enabling more effective majority voting under fixed compute budgets.

**Strengths:**

- Tackles a timely and practically important problem (LLM overthinking).

- Difficulty-aware length control and Advantage Weighting are creative and well-motivated.

- Clear and well-written presentation.

- Strong empirical results, including new large-model (7B) experiments confirming scalability.

- Highlights inference scaling benefits overlooked by prior work.

**Weaknesses:**

- Missing direct comparisons to some concurrent baselines (e.g., L1).

- Theoretical justification of GRPO pitfalls relies on debatable assumptions, though partially clarified in rebuttal.

- Initial experiments limited to small models; later addressed.

- Some ablations were missing initially but added later.


**Review discussion and rebuttal:**

Reviewer zKPN: Concerned about missing baselines and theoretical bias; authors argued complementarity with L1, clarified Advantage Weighting, and added ablations. Reviewer acknowledged responses but stayed borderline (4).

Reviewer gh6m: Concerned about scalability; authors added 7B experiments showing improved efficiency and accuracy, leading to a score increase to 5.

Reviewer vKd6: Questioned GRPO analysis and scaling of hyperparameters; authors corrected proofs, added ablations (N, target functions), and clarified inference scaling. Reviewer found responses satisfactory and raised to 5.

Overall, the rebuttal strengthened the paper significantly, with two reviewers raising their scores and one remaining borderline.

--

Despite limitations in baseline coverage and theoretical depth, the paper makes a meaningful and timely contribution with a principled framework that improves both efficiency and reasoning quality, validated at scale. Overall, the strengths outweigh the weaknesses.
Decision: Accept (poster)