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.

We have addressed the identified weaknesses below and have conducted new experiments and revised our manuscript to incorporate this valuable feedback.

1. On the Pitfalls of Naive Reward Weighting in GRPO

We thank the reviewer for raising these two important points regarding the analysis in Appendix B.

• 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.

• 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.

  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:

  ...[content from Line 523 to Line 552, with $\alpha_{ada}$ replaced with $\alpha_{base}$]

  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}$.

  To visualize this, we can analyze its squared form, $\hat{\tau}\_p^2$:

  $$\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}.$$

  Assuming $\hat{C} \neq 0$ for now, this simplifies to:

  $$\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}.$$

  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.

  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.

  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.

  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.

2. Ablation on Normalizing the Penalty Term

This 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:

As 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.

3. Impact of Scaling N

This is another important hyperparameter. We selected N=8 based on preliminary experiments. As requested, we have presented experiments for N=4 and 16.

The 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.

Considering 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.

4. Functional Form of Difficulty vs. Token Budget

This 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:

The 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.

5. Clarification on Inference Scaling

Thank 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.

Our 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.

We 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.

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.

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.

From these new results, we draw two key conclusions:

1. 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.
2. 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.

These 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.

Thank you for running additional experiments with a 7B model. This further strengthens the potential significance of the results.

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.

Below, we respond to each of the questions and weaknesses identified.

1. On the Comparison with L1

We view our work and L1 as addressing different and complementary aspects of the token compression problem.

• 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.
• 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.

As 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.

2. Difference from GRPO and Unbiased Gradient Estimation

Thank you for these critical questions on the nature of our advantage calculation.

• 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.
• 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.

Here's the major part of the added proof in our manuscript (refresh the page if it is not correctly rendered):

[...]

First, we define the optimization objective $J(\theta)$, which combines the task performance reward with the difficulty-aware length penalty:

$$J(\theta) = \mathbb{E}\_{x \sim \mathcal{D}, y \sim \pi\_{\theta}(\cdot|x)} \left[ r\_{\text{total}}(x, y) \right], \quad \text{where} \quad r\_{\text{total}}(x, y) = r\_{\text{outcome}}(x, y) - \alpha'(x, \pi\_{\theta}) \cdot p(y).$$

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$.

According to the Policy Gradient Theorem, the true gradient of this objective is:

$$\nabla\_{\theta} J(\theta) = \mathbb{E}\_{x, y} \left[ \nabla\_{\theta} \log \pi\_{\theta}(y|x) \cdot r\_{\text{total}}(x, y) \right].$$

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:

$$\mathbb{E}\_{y \sim \pi\_{\theta}(\cdot|x)} \left[ \nabla\_{\theta} \log \pi\_{\theta}(y|x) \cdot B(x) \right] = B(x) \cdot \int\_{y} \pi\_{\theta}(y|x) \nabla\_{\theta} \log \pi\_{\theta}(y|x) dy = B(x) \cdot \int\_{y} \nabla\_{\theta} \pi\_{\theta}(y|x) dy = B(x) \cdot \nabla\_{\theta} \int\_{y} \pi\_{\theta}(y|x) dy = B(x) \cdot \nabla\_{\theta}(1) = 0$$

Thus, the true gradient can be equivalently written as:

$$\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].$$

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:

$$B(x) = \mu_{\text{outcome}}(x) - \alpha'(x, \pi_{\theta}) \cdot \mu_p(x).$$

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:

$$\begin{aligned} A(x, y_i) = \left( r_{\text{outcome}, i} - \alpha' p_i \right) - \left( \mu_{\text{outcome}} - \alpha' \mu_p \right) = (r_{\text{outcome},i} - \mu_{\text{outcome}}) - \alpha'(x, \pi_{\theta}) \cdot (p_i - \mu_{p}). \end{aligned}$$

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.

We 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):

While 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!

3. Ablation Studies on Dynamic Variables

We appreciate the request for more insight into the ablation study, as you have proposed in your questions 2 and 5.

Ablation on Dynamic Weight ($\alpha_{base}$): We tested several values for the base trade-off parameter under the Adaptive Weighting setting.

We chose $\alpha_{base}=0.5$ for our main experiments as it provided the best balance between significant token reduction and high performance.

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.

To 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:

The 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.

Abaltion on $\delta$ and $L_{max}$ in $t(x,\pi_{theta})$

We added ablation study for $\delta$ and $L_{max}$, the results are as follows:

Reducing $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.

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.

We'll include these results in the final version. Your feedback has been invaluable in strengthening our paper's completeness.