TokenSqueeze: Performance-Preserving Compression for Reasoning LLMs

We sincerely thank you for your careful review and positive feedback. We appreciate your recognition of the originality of our approach, the practicality of our objective, and the overall clarity of the paper. Your encouraging comments are very helpful to us, and we address the remaining points in detail below.

W1: While the proposed DPO-L formulation integrates response length into the preference loss, the core idea of using length-based regularization in preference optimization has been explored in prior works.[1,2]

Response: In contrast to approaches like O1-Pruner and Kimi-k1.5, which primarily rely on adding a length regularization term within a policy gradient framework, our core idea is not simply to impose a regularizer. Instead, we propose a systematic approach to preference optimization that improves inference efficiency without sacrificing correctness. Our method enables the model to implicitly prefer responses that are both accurate and concise, offering a fundamentally different paradigm from methods that focus solely on length penalties.

Moreover, the way training data is constructed also differs fundamentally. We automatically derive $(\mathrm{chosen\ sample}, \mathrm{rejected\ sample})$ pairs from model outputs: for each prompt, the “chosen” response is correct and shorter, while longer, incorrect ones are labeled as “rejected.” In contrast, O1-Pruner and Kimi-k1.5 construct their training data as $(\mathrm{sample}, \mathrm{reward})$ pairs, where each response is assigned a scalar reward.

W2: The formulation in Equation (2) lacks clarity. Specifically, it is unclear whether the ( j )-th token refers to the position following the ( i )-th reasoning step, as implied in the main text. A more precise definition would improve interpretability.

Response: We appreciate your comment and will clarify this in the camera-ready version if the paper is accepted. Specifically, in Equation $2$:

We clarify that $t_j$ refers to the $j$-th token in the consecutive reasoning steps following the $i$-th step $s_i$. Specifically, $t_1$ corresponds to the first token of $s_{i+1}$. If the tokens in $s_{i+1}$ are insufficient to fill the window, we continue accumulating tokens from subsequent steps until the window is filled or the response ends.

To provide more insight into Equation 2, we provide the derivation from Equation 1 to Equation 2 below.

Proof (from Equation (1) to Equation (2))

Notation and Setup

We begin by stating the two key KL divergence expressions from the paper:

• Equation (1):

  $$D_{\mathrm{KL}}\left( P_{\theta}(\cdot \mid p, s_{\leq i}) \| P_{\theta}(\cdot \mid p, s_{<i}, s_i') \right)$$

• Equation (2):

  $$D_{\mathrm{KL}}^{\text{full}} \approx \sum_{j=1}^{\min(T,L)} D_{\mathrm{KL}}\left( Q_{\theta}(\cdot \mid p, s_{\leq i}, t_{1:j-1}) \| Q_{\theta}(\cdot \mid p, s_{<i}, s_i', t_{1:j-1}) \right)$$

These expressions measure the change in predicted continuation distributions before and after rewriting step $s_i$.

We now define the key variables and distributions involved in the KL computation:

• $p$: the input prompt to the model.
• $s_i$: the $i$‑th reasoning step generated by the model.
• $s_{\le i}$: the reasoning trace up to and including step $i$ (i.e., $s_1, s_2, \dots, s_i$).
• $s'_i$: a rewritten version of step $s_i$ used in intra-step refinement.
• $t_j$: the $j$‑th token after $s_i$ (or $s'_i$), starting the continuation sequence.
• $t_{1:T}$: a full continuation sequence of $T$ tokens (including EOS).

We define the following probability distributions:

• $Q_\theta(\cdot \mid \cdot)$: the model’s next-token distribution under the given context (token-level), parameterized by $\theta$.
• $P_\theta(\cdot \mid \cdot)$: the model’s sequence-level distribution over continuations $t_{1:T}$, given the input context.

To simplify notation, we define:

• $\mathsf{A}(t_{1:T}) := P_\theta(t_{1:T} \mid p, s_{\le i})$
• $\mathsf{B}(t_{1:T}) := P_\theta(t_{1:T} \mid p, s_{< i}, s'_i)$

These represent the two sequence distributions compared in the KL divergence.

Step 1: Sequence‑Level KL Definition

The KL term in Equation (1) is:

Expanding both $\mathsf{A}$ and $\mathsf{B}$ autoregressively using $Q_\theta$:

Substituting into the KL definition:

This shows that the sequence-level KL equals the expected sum of per-token conditional KL divergences.

Step 2: Monte Carlo Approximation

Computing the full expectation over all possible token histories $t_{1:j-1}$ is intractable in practice. To address this, we use a Monte Carlo approximation: instead of averaging over all continuations, we compute the KL terms along an existing sampled trajectory $t_{1:T}$ drawn from $\mathsf{A}$:

Step 3: Token-Window Truncation

To reduce cost further, we truncate the sum at a fixed window size $L$ (e.g., 512 tokens):

This is exactly the formulation given in Equation (2) of the main paper.

The approximation introduces a residual term:

which is small when the effect of rewriting $s_i$ fades over time.

W3: The paper would benefit from including experiments with alternative backbone models to verify the generalizability of the proposed method across different architectures.

Response: To further demonstrate the effectiveness of our model across different backbones and LLM scales, we conducted experiments using Phi4-Reasoning-14B. The results are presented below:

It is important to note that, due to time constraints, we reducing K in intra-step linguistic refinement from 64 to 16. We expect that further hyperparameter tuning would yield even better performance.

We sincerely thank you for your careful review and positive assessment of our work. Your feedback is greatly valued, and we respond to your points in detail below.

Response to W1: We conducted an experiment using 1,000 randomly sampled initial question prompts on a single node with 8×A100 GPUs. The runtime (in hours) for each major preprocessing stage is shown below:

This runtime is comparable to standard preprocessing pipelines in industrial practice, such as reward-based rejection sampling, and remains within acceptable bounds for large-scale training. Moreover, since our method constructs the dataset offline, it incurs this cost only once and can be reused across multiple training runs, unlike online reinforcement learning methods that require repeated sampling during training.

To further reduce computational cost, we have implemented several optimization strategies:

• KV-cache reuse: When processing step $s_{i+1}$, we reuse the key-value cache computed during the sampling of step $s_i$ in both resampling and KL divergence calculation, significantly reducing the cost of autoregressive decoding.

• Selective rewriting: We first filter generated samples to identify preference-worthy candidates (e.g., concise and correct answers) before rewriting. This ensures rewriting is applied only to a targeted subset.

• KL window approximation: The KL divergence is approximated within a fixed-length window of 512 tokens:

  $$D_{\mathrm{KL}}^{\text{full}} \approx \sum_{j=1}^{\min(T, L)} D_{\mathrm{KL}}\left( Q_{\theta}(\cdot \mid p, s_{\leq i}, t_{1:j-1}) \,\|\, Q_{\theta}(\cdot \mid p, s_{<i}, s_i', t_{1:j-1}) \right),$$

  This avoids full-sequence computation and offers significant resource savings.

• Choice of $K$:
  We would like to clarify that the resampling parameter $K = 64$ was chosen as a strong setting to push performance to its limit for benchmark comparison. However, our empirical observations indicate that smaller values such as $K = 32$ still achieve highly competitive performance while significantly reducing computational cost. This suggests that such a large value is not strictly necessary in practice. We believe that a more comprehensive investigation of the impact of different $K$ values is an important direction for future work, and we plan to explore this systematically in subsequent versions of the paper or follow-up research. As shown in our response to Reviewer 1 (W1), the experiments on Phi-4-Reasoning-14B with $K = 16$ still demonstrate a significant improvement in the model’s efficiency.

Response to W2: As shown in Figure 4 of the paper, we analyzed the performance impact of different $\alpha$ values and found that extreme values (too large or too small) negatively affect both accuracy and token efficiency. We selected $\alpha = 0.2$ as it achieves the best balance between brevity and correctness.

For the KL threshold $\varepsilon$, we selected its value based on empirical observations. To further investigate the effect of $\varepsilon$, we additionally conducted a quantitative experiment with $\varepsilon = 0.02$. The results are shown below:

These results confirm that excessively high values of $\varepsilon$ can negatively affect accuracy. We plan to explore the effect of varying $\varepsilon$ more systematically in future work.

In addition to the quantitative results, we also provide a qualitative example that illustrates how different values of $\varepsilon$ influence the form and length of the generated rewrites:

Response to W3: Thank you for your question. Due to space limitations, we could not include the full derivation in the main paper. If the paper is accepted, we will provide the complete proof in the appendix of the camera-ready version.

Since both data generation and step rewriting are performed on the same model, we define the original step $s_i$ and its refined version $s_i'$ as semantically equivalent if they satisfy the following condition:

as formulated in Equation (1). This constraint ensures that rewriting does not compromise the semantic integrity of the reasoning trace.

Below we provide the derivation from Equation (1) to Equation (2):

Proof (from Equation (1) to Equation (2))

Notation and Setup

We begin by stating the two key KL divergence expressions from the paper:

• Equation (1):

  $$D_{\mathrm{KL}}\left( P_{\theta}(\cdot \mid p, s_{\leq i}) \| P_{\theta}(\cdot \mid p, s_{<i}, s_i') \right)$$

• Equation (2):

  $$D_{\mathrm{KL}}^{\text{full}} \approx \sum_{j=1}^{\min(T,L)} D_{\mathrm{KL}}\left( Q_{\theta}(\cdot \mid p, s_{\leq i}, t_{1:j-1}) \| Q_{\theta}(\cdot \mid p, s_{<i}, s_i', t_{1:j-1}) \right)$$

These expressions measure the change in predicted continuation distributions before and after rewriting step $s_i$.

We now define the key variables and distributions involved in the KL computation:

• $p$: the input prompt to the model.
• $s_i$: the $i$‑th reasoning step generated by the model.
• $s_{\le i}$: the reasoning trace up to and including step $i$ (i.e., $s_1, s_2, \dots, s_i$).
• $s'_i$: a rewritten version of step $s_i$ used in intra-step refinement.
• $t_j$: the $j$‑th token after $s_i$ (or $s'_i$), starting the continuation sequence.
• $t_{1:T}$: a full continuation sequence of $T$ tokens (including EOS).

We define the following probability distributions:

• $Q_\theta(\cdot \mid \cdot)$: the model’s next-token distribution under the given context (token-level), parameterized by $\theta$.
• $P_\theta(\cdot \mid \cdot)$: the model’s sequence-level distribution over continuations $t_{1:T}$, given the input context.

To simplify notation, we define:

• $\mathsf{A}(t_{1:T}) := P_\theta(t_{1:T} \mid p, s_{\le i})$
• $\mathsf{B}(t_{1:T}) := P_\theta(t_{1:T} \mid p, s_{< i}, s'_i)$

These represent the two sequence distributions compared in the KL divergence.

Step 1: Sequence‑Level KL Definition

The KL term in Equation (1) is:

Expanding both $\mathsf{A}$ and $\mathsf{B}$ autoregressively using $Q_\theta$:

Substituting into the KL definition:

This shows that the sequence-level KL equals the expected sum of per-token conditional KL divergences.

Step 2: Monte Carlo Approximation

Computing the full expectation over all possible token histories $t_{1:j-1}$ is intractable in practice. To address this, we use a Monte Carlo approximation: instead of averaging over all continuations, we compute the KL terms along an existing sampled trajectory $t_{1:T}$ drawn from $\mathsf{A}$:

Step 3: Token-Window Truncation

To reduce cost further, we truncate the sum at a fixed window size $L$ (e.g., 512 tokens):

This is exactly the formulation given in Equation (2) of the main paper.

The approximation introduces a residual term:

which is small when the effect of rewriting $s_i$ fades over time.

We sincerely thank you for your careful reading and constructive feedback. We appreciate your recognition of the overall clarity of the paper, the novelty of the intra-step refinement component, and the effectiveness of our method in reducing token usage while preserving accuracy. Your comments have helped us identify several areas for clarification and improvement, and we address each of your points in detail below.

W1 W2 W3 W4:

• In line 13, the authors mention "we propose to selectively identify self-generated samples." While the general concept becomes clearer in subsequent sections, this point is initially vague and could benefit from clarification upfront.

• In line 48, the paper states that when the reasoning length exceeds a certain token threshold, the correlation between token count and model performance significantly weakens. This observation would be more appropriately framed as part of the motivation in the Methods section rather than being introduced in the Experiments section.

• The notation $s$ is used inconsistently: in line 135 it refers to reasoning traces, while in line 151 it denotes a reasoning step. This dual usage creates confusion and should be clarified.

• The symbol $\alpha$appears twice with different meanings (e.g., line 132 and Equation 4). These should be clearly distinguished or redefined to avoid ambiguity.

Response: Thank you for pointing out these important issues. If the paper is accepted, we will make the following revisions in the camera-ready version:

1. In the abstract, we will revise the sentence “we propose to selectively identify self-generated samples that preserve the model’s competency while minimizing token usage” to “we propose to select self-generated samples whose reasoning depth is adaptively matched to the complexity of the problem” to improve clarity.

2. In lines 132 and 135, we will revise the overloaded symbols used for quantile computation and reasoning trace collections to eliminate ambiguity and improve readability.

3. Regarding the placement of the token-threshold observation (line 48), we acknowledge your suggestion and will revise it accordingly in the camera-ready version if the paper is accepted. Due to space limitations, we omit the details here.

W3: The paper lacks an ablation study on the choice of $M$, as referenced in line 139. This analysis is important for understanding the impact of the selection strategy.

Response: To further investigate the impact of the number of samples per instance $M$ during the data collection phase, we additionally trained DeepSeek-R1-Distill-Qwen-7B with $M = 32$. The results are presented below, where Baseline refers to the original unmodified DeepSeek-R1-Distill-Qwen-7B model:

These additional experiments allow us to assess how varying the sampling number $M$ influences the trade-off between reasoning performance and inference efficiency.

W6: In line 137, it is unclear how the "longer incorrect response" is obtained. More detail on this process would be helpful.

Response: The “long and incorrect” samples are defined relative to the chosen ones. Specifically, during data construction, for each chosen sample, we select rejected samples that satisfy both of the following conditions:

1. The response is longer than the chosen sample in token length.
2. The response is classified as incorrect based on final answer comparison.

This design ensures that the preference objective simultaneously penalizes both verbosity and incorrectness, aligning the optimization with our intended trade-off between accuracy and efficiency.

W7: In line 157, the paper states that the level of information preservation is "tunable." Is this tunability controlled manually, or is it learned during training? This should be specified.

Response: The tunability of information preservation is implemented at training time during data preprocessing.

Our method allows adjusting the strictness of information preservation in the rewriting phase via a KL divergence threshold $\varepsilon$. This hyperparameter governs how much semantic deviation is tolerated when rewriting each reasoning step. Specifically, given a set of candidate rewrites $\{s_i^{(k)}\}$, we select the shortest one that satisfies the following constraint:

Here, $\ell(s_i')$ denotes the token length of the rewritten step, and $P_\theta$ is the model's sequence-level distribution. This KL-based constraint ensures that the semantic integrity of the original reasoning trace is preserved during the offline data refinement stage.

Below, we provide a concrete example illustrating how different threshold values affect the resulting rewrites and their lengths:

W8: In Section 3.2, the method for accurately segmenting a reasoning step is unclear, especially since reasoning steps may contain many tokens and ambiguous boundaries. How is the stopping criterion determined?

Response: We segment reasoning steps based on double newline characters \n\n, which is a common convention adopted by most mainstream open-source LLMs. This formatting is naturally used by model family such as DeepSeeks and Qwens when generating step-by-step reasoning outputs. As a result, our segmentation aligns well with model behavior and does not require additional annotation or post-processing.

W9: The generated reasoning lengths appear unstable. Were multiple random seeds used in the experiments to evaluate robustness?

Response: Thank you for raising this point. For smaller datasets such as AIME24 and AIME25, we adopted an evaluation strategy designed to ensure robustness against potential variability in response lengths. Specifically, we averaged the results over 16 independent inference runs to obtain stable performance estimates. All generations were conducted using the default random seed provided by the vLLM framework.

W10: It would strengthen the paper to compare results with state-of-the-art methods such as O1-Pruner, L1, and others.

Response: Since O1-Pruner and L1 are built upon different base models than ours, their results are not directly comparable to our setting. Therefore, in the paper, we compared our method with DAST and TrainEffi, which utilize the same base model.

Considering that L1 focuses more on controllable reasoning length, while O1-Pruner emphasizes reasoning efficiency, we additionally reproduced O1-Pruner using its open-sourced codebase. Following the official configuration of O1-Pruner, we reproduced their method on a single node with 8 × A100 80GB GPUs, using DeepSeek-R1-Distill-Qwen-7B as the base model. We strictly followed the hyperparameter settings described in the original paper, without any modifications.

The reproduction results are summarized below:

As shown in the results, our method achieves significantly higher accuracy on challenging tasks like AIME24, while also delivering superior AUC and substantially reduced reasoning lengths compared to O1-Pruner.

The authors addressed my concerns, and I would like to keep the positive rating.

We sincerely thank you for your valuable feedback and positive evaluation of our work. We have carefully addressed each of your suggestions and questions in the responses below.

Response to Weakness 1: To further demonstrate the effectiveness of our model, we conducted experiments using Phi4-Reasoning-14B. The results are presented below:

It is important to note that, due to time constraints, we reducing K in intra-step linguistic refinement from 64 to 16. We expect that further hyperparameter tuning would yield even better performance.

Response to Weakness 2: We will include a dedicated section on limitations and future work in the camera-ready version, if the paper is accepted. Some of the limitations and future works we plan to include are listed below:

1. Heuristic Hyperparameter Selection: Certain hyperparameters (e.g., the KL threshold $\varepsilon$ used in intra-step rewriting) were selected heuristically based on preliminary experiments. We did not conduct a systematic analysis. In future work, we plan to investigate how varying $\varepsilon$ affects the trade-off between semantic fidelity and brevity.
2. Offline-only Setting: The current pipeline is restricted to offline preference optimisation. We plan to extend it to online reinforcement learning to enable tighter integration between data generation and model updates, potentially improving the trade-off between reasoning efficiency and accuracy.

Response to Question 1: At the time of submission, DAST and TrainEffi had not provided results on the AIME25 and LiveCodeBench datasets in their papers, and neither released training code or model checkpoints. However, in June, DAST made its model publicly available. We evaluated their model on AIME25 and LiveCodeBench, and report the results below:

Our model outperforms DAST in both accuracy and AUC on AIME25 and LiveCodeBench, with especially notable gains on LiveCodeBench, demonstrating stronger generalization.

Response to Question 2: The two KL terms serve different purposes. In the intra-step rewriting, the KL divergence aims to ensure that the semantic information of the data is preserved before and after rewriting. In contrast, the KL divergence in the training objective is used to constrain the model from deviating too far from the original model (to prevent training collapse). Given the multi-stage training process (data collection → data rewriting → model training), it is currently challenging to combine these two KL terms into a single joint objective. However, for online reinforcement learning methods such as PPO and GRPO, data collection and model training occur simultaneously during the RL process, and we believe there is potential to use a unified optimization objective. We will continue to explore this alignment in future work.

Response to Question 3: In Equation $2$:

We clarify that $t_j$ refers to the $j$-th token in the consecutive reasoning steps following the $i$-th step $s_i$. Specifically, $t_1$ corresponds to the first token of $s_{i+1}$. If the tokens in $s_{i+1}$ are insufficient to fill the window, we continue accumulating tokens from subsequent steps until the window is filled or the response ends.

In fact, KL divergence at the token level can be used as an unbiased approximation for sequence-level KL divergence. In Equation 1, we use sequence-level KL divergence, while in Equation 2, we approximate it using token-level KL divergence. Below is a proof that shows how we transition from sequence-level to token-level KL divergence.

Proof (from Equation (1) to Equation (2))

Notation and Setup

We begin by stating the two key KL divergence expressions from the paper:

• Equation (1):

  $$D_{\mathrm{KL}}\left( P_{\theta}(\cdot \mid p, s_{\leq i}) \| P_{\theta}(\cdot \mid p, s_{<i}, s_i') \right)$$

• Equation (2):

  $$D_{\mathrm{KL}}^{\text{full}} \approx \sum_{j=1}^{\min(T,L)} D_{\mathrm{KL}}\left( Q_{\theta}(\cdot \mid p, s_{\leq i}, t_{1:j-1}) \| Q_{\theta}(\cdot \mid p, s_{<i}, s_i', t_{1:j-1}) \right)$$

These expressions measure the change in predicted continuation distributions before and after rewriting step $s_i$.

We now define the key variables and distributions involved in the KL computation:

• $p$: the input prompt to the model.
• $s_i$: the $i$‑th reasoning step generated by the model.
• $s_{\le i}$: the reasoning trace up to and including step $i$ (i.e., $s_1, s_2, \dots, s_i$).
• $s'_i$: a rewritten version of step $s_i$ used in intra-step refinement.
• $t_j$: the $j$‑th token after $s_i$ (or $s'_i$), starting the continuation sequence.
• $t_{1:T}$: a full continuation sequence of $T$ tokens (including EOS).

We define the following probability distributions:

• $Q_\theta(\cdot \mid \cdot)$: the model’s next-token distribution under the given context (token-level), parameterized by $\theta$.
• $P_\theta(\cdot \mid \cdot)$: the model’s sequence-level distribution over continuations $t_{1:T}$, given the input context.

To simplify notation, we define:

• $\mathsf{A}(t_{1:T}) := P_\theta(t_{1:T} \mid p, s_{\le i})$
• $\mathsf{B}(t_{1:T}) := P_\theta(t_{1:T} \mid p, s_{< i}, s'_i)$

These represent the two sequence distributions compared in the KL divergence.

Step 1: Sequence‑Level KL Definition

The KL term in Equation (1) is:

Expanding both $\mathsf{A}$ and $\mathsf{B}$ autoregressively using $Q_\theta$:

Substituting into the KL definition:

This shows that the sequence-level KL equals the expected sum of per-token conditional KL divergences.

Step 2: Monte Carlo Approximation

Computing the full expectation over all possible token histories $t_{1:j-1}$ is intractable in practice. To address this, we use a Monte Carlo approximation: instead of averaging over all continuations, we compute the KL terms along an existing sampled trajectory $t_{1:T}$ drawn from $\mathsf{A}$:

Step 3: Token-Window Truncation

To reduce cost further, we truncate the sum at a fixed window size $L$ (e.g., 512 tokens):

This is exactly the formulation given in Equation (2) of the main paper.

The approximation introduces a residual term:

which is small when the effect of rewriting $s_i$ fades over time.

If the paper is accepted, we will include the full proof in the appendix of the camera-ready version.