TokenSqueeze: Performance-Preserving Compression for Reasoning LLMs

Yuxiang Zhang; Zhengxu Yu; Weihang Pan; Zhongming Jin; Qiang Fu; Deng Cai; Binbin Lin; Jieping Ye

Source-linked AI summary

TokenSqueeze: Performance-Preserving Compression for Reasoning LLMs

Yuxiang Zhang, Zhengxu Yu, Weihang Pan, Zhongming Jin, Qiang Fu, Deng Cai, Binbin Lin, Jieping Ye

arXiv:2511.13223v1Nov 17, 2025cs.LGcs.AI

TL;DR

Long reasoning traces improve difficult-task performance but increase token use, latency, and memory, while existing Long2Short methods can sacrifice accuracy. TokenSqueeze uses self-generated data, adaptive reasoning-depth selection, linguistic refinement, and length-aware preference optimization to compress reasoning. Experiments report reduced token usage while preserving or improving accuracy across math and code reasoning benchmarks.

Problem
Long chain-of-thought reasoning increases inference costs, while existing compression methods can remove essential reasoning steps and reduce accuracy.
Method
TokenSqueeze constructs self-generated long-short reasoning pairs using adaptive depth selection and intra-step linguistic refinement, then trains with length-aware preference optimization.
Results
TokenSqueeze significantly reduces token usage while preserving or improving accuracy on reasoning benchmarks including math and code tasks.
Takeaways & Limitations
Self-generated training data can support efficient, high-fidelity reasoning without manual annotations or teacher models.
Takeaways & Limitations
The offline framework cannot continuously adapt its reasoning strategy to new feedback or shifts in task distribution.

Abstract

from arXiv · show

Emerging reasoning LLMs such as OpenAI-o1 and DeepSeek-R1 have achieved strong performance on complex reasoning tasks by generating long chain-of-thought (CoT) traces. However, these long CoTs result in increased token usage, leading to higher inference latency and memory consumption. As a result, balancing accuracy and reasoning efficiency has become essential for deploying reasoning LLMs in practical applications. Existing long-to-short (Long2Short) methods aim to reduce inference length but often sacrifice accuracy, revealing a need for an approach that maintains performance while lowering token costs. To address this efficiency-accuracy tradeoff, we propose TokenSqueeze, a novel Long2Short method that condenses reasoning paths while preserving performance and relying exclusively on self-generated data. First, to prevent performance degradation caused by excessive compression of reasoning depth, we propose to select self-generated samples whose reasoning depth is adaptively matched to the complexity of the problem. To further optimize the linguistic expression without altering the underlying reasoning paths, we introduce a distribution-aligned linguistic refinement method that enhances the clarity and conciseness of the reasoning path while preserving its logical integrity. Comprehensive experimental results demonstrate the effectiveness of TokenSqueeze in reducing token usage while maintaining accuracy. Notably, DeepSeek-R1-Distill-Qwen-7B fine-tuned using our proposed method achieved a 50\% average token reduction while preserving accuracy on the MATH500 benchmark. TokenSqueeze exclusively utilizes the model's self-generated data, enabling efficient and high-fidelity reasoning without relying on manually curated short-answer datasets across diverse applications. Our code is available at https://github.com/zhangyx1122/TokenSqueeze.

1 Introduction

Long chain-of-thought improves complex reasoning but raises latency, memory, and overthinking costs, while existing compression methods can remove essential reasoning steps. TokenSqueeze addresses this tradeoff with self-generated, complexity-aware reasoning compression and length-aware preference learning.

Long CoT improves complex reasoning but increases inference latency and memory usage, while redundant steps can impede practical applications.The problem is especially relevant for time-sensitive, resource-constrained, and multi-turn agent settings.
Existing inference- and train-time compression methods often reduce output length at the cost of essential reasoning steps and accuracy.The paper describes this tradeoff as the reasoning oversimplification dilemma.
TokenSqueeze treats Long2Short compression as a preference-learning problem that favors concise expression while preserving reasoning depth matched to problem complexity.The approach relies on self-generated reasoning data rather than external teacher models or additional annotations.
The method constructs long-short reasoning pairs through adaptive depth selection and intra-step linguistic refinement that preserves essential steps and meaning.Adaptive selection addresses problem complexity, while refinement increases information density without changing the underlying reasoning path.
A length-aware preference objective reinforces concise reasoning while experiments report improved reasoning efficiency without sacrificing model performance across reasoning tasks.The introduction characterizes the method as broadly applicable and balancing efficiency with accuracy.

2 Related Work

Related work compresses reasoning through online reinforcement learning or offline response construction, but these approaches involve computational, data-quality, or external-supervision tradeoffs. Offline methods avoid repeated sampling, whereas dataset construction commonly selects or rewrites shorter responses.

2.1 Online Reinforcement Learning with Length Penalty: Online reinforcement learning methods reward concise and correct responses through modified reward functions, but response sampling at every training step makes them computationally expensive.The added sampling burden limits scalability for large-model training and becomes heavier when many traces are needed.
2.2 Offline Optimization: Offline optimization avoids repeated sampling during training by constructing concise reasoning datasets and then applying supervised fine-tuning or direct preference optimization.This design offers substantial computational savings relative to online approaches.
2.2 Offline Optimization: Several offline methods select the shortest correct responses from model-generated pools or use prompting to guide shorter-answer generation before selection.Examples include Kimi-k1.5 (DPO) [29], Sky-T1-Flash, Self-Training [23], and Token-Budget [12].
2.2 Offline Optimization: Other methods improve training data by rephrasing model responses with external LLMs or pruning content through rule-based token elimination.C3oT [17] uses GPT-4 [1] for refinement, while TokenSkip [33] uses importance-based token elimination.

3 Methodology

TokenSqueeze builds concise reasoning data from self-generated traces, then trains with objectives that preserve logical fidelity while promoting brevity. Its pipeline combines adaptive depth selection, KL-constrained intra-step refinement, and length-aware preference optimization.

3.1 Adaptive Reasoning Depth Selection: Adaptive depth selection retains shorter chains for easier problems and longer chains for harder ones, preserving critical reasoning steps across difficulty levels.The method computes q = α · (1 − p), selects the shortest correct traces up to k = ⌈q · c⌉, and pairs them with longer incorrect responses.
3 Methodology: The overall workflow selects self-generated traces by adaptive depth, rewrites them under KL constraints, and trains for both accuracy and brevity.Figure 1 summarizes the three-stage pipeline, while Figure 2 illustrates increased information density with preserved logical integrity after refinement.
3.2 Intra-Step Linguistic Refinement via Distributional Alignment: Intra-step refinement resamples candidate rewrites and selects the shortest one whose KL divergence remains within an information-preservation threshold.The procedure uses local token-window approximation with a fixed window size L = 512 to estimate divergence over subsequent reasoning tokens.
3.3 Composite Optimization Objective: The composite objective combines length-aware DPO with supervised fine-tuning to encourage compact, correct reasoning without destabilizing preferred responses.Its logarithmic length-ratio term strengthens preference signals for pairs with larger compression gains, while preserving standard behavior for comparable-length pairs.

4 Experiments

Across four benchmarks and two model scales, TokenSqueeze improves token efficiency while maintaining or improving accuracy. Ablations show that adaptive reasoning depth, linguistic refinement, and composite optimization jointly support this balance.

4.2 Evaluation on General Reasoning Benchmarks: TokenSqueeze consistently improves accuracy while significantly reducing token consumption across four benchmarks and two model scales, yielding higher AUC than baselines.Table 1 compares training methods across model sizes; AUC measures performance under token-budget constraints.
4.2 Evaluation on General Reasoning Benchmarks: On AIME24, MATH500, and AIME25, TokenSqueeze compresses reasoning substantially without sacrificing accuracy, while LiveCodeBench results show improved accuracy and shorter outputs.The 7B model generally shows stronger compression than the 1.5B model across most metrics.
4.2 Evaluation on General Reasoning Benchmarks: 15.5% higher accuracy on AIME24 at 3K tokens and 43.1% on MATH500 at 1K tokens are achieved over the base model under matched budgets.These gains are reported alongside consistent AUC improvements across datasets.
4.3.1 Impact of Adaptive Reasoning Depth Selection: Q-DYN achieves the highest accuracy on both ablation datasets while maintaining competitive response length by adapting reasoning depth to problem complexity.Q-FIX improves accuracy over selecting only the shortest trace, whereas adding correct traces as negatives reduces accuracy through excessive compression.
4.3.1 Impact of Adaptive Reasoning Depth Selection: α = 0.2 provides the best balance between accuracy and token efficiency, while extreme values degrade performance through excessive compression or overly long reasoning chains.The adaptive threshold is q = α × (1 − p), where p is the correctness rate.
4.3.2 Impact of Intra-Step Linguistic Refinement: TokenSqueeze’s full refinement further shortens individual reasoning steps after depth compression, preserving essential information while maintaining high accuracy.Compared with no refinement, GPT-4o-mini rewriting, and TokenSkip, the proposed refinement achieves a better balance between conciseness and correctness.

5 Limitation

TokenSqueeze retains efficiency and accuracy gains but remains limited by heuristic hyperparameter selection and its offline-only training setting.

Heuristic Hyperparameter Selection: The KL threshold ε was chosen heuristically rather than through exhaustive optimization, leaving its sensitivity and parameter interactions insufficiently explored.A lower ε can restrict compression, whereas a higher ε may cause semantic drift.
Offline-only Setting: The framework operates entirely offline with pre-generated preference pairs, limiting adaptation to new feedback and task-distribution shifts.The authors propose extending preference generation, policy updates, and reward estimation into online reinforcement learning.

6 Conclusion

The conclusion presents TokenSqueeze as a self-generated-data method for compressing reasoning traces while preserving task performance. Across math and code benchmarks, it reduces token usage while preserving or improving accuracy.

6 Conclusion: TokenSqueeze combines adaptive reasoning-depth selection, intra-step linguistic refinement, and length-aware preference optimization to compress reasoning traces without manual annotations or teacher models.The method uses only self-generated training data.
6 Conclusion: TokenSqueeze significantly reduces token usage while preserving or even improving accuracy across reasoning benchmarks, including math and code tasks.The conclusion frames this as a practical path toward efficient, high-fidelity reasoning for deployment.

A.1 Data Construction and Compression Settings

The data-construction pipeline samples diverse reasoning traces, refines individual steps under a KL constraint, and selects preference pairs using problem-adaptive reasoning depth.

A.1 Data Construction and Compression Settings: The refinement stage samples 64 candidate rewrites per reasoning step and selects the shortest one when its KL divergence is at most 0.005.Self-sampled traces use temperature 0.9, while rewrites use temperature 1.0.
A.1 Data Construction and Compression Settings: The prompting format separates reasoning traces from final answers using <think> and <answer> tags throughout data collection and training.This format explicitly distinguishes internal reasoning from the user-facing answer.
A.1 Data Construction and Compression Settings: Preference-pair construction uses the adaptive threshold q = α(1 − p), with α = 0.2 and p representing correctness rate.The threshold adapts reasoning depth to problem difficulty by balancing informativeness and brevity.

A.2 Training Configurations

Training fine-tunes both DeepSeek-R1-Distill-Qwen models with a combined supervised-preference objective and an explicit length-aware penalty for concise outputs.

A.2 Training Configurations: The training objective combines supervised fine-tuning and DPO losses with equal weights of 0.5 each.The implementation is based on the LLaMAFactory DPO pipeline.
A.2 Training Configurations: Both DeepSeek-R1-Distill-Qwen-7B and 1.5B are fine-tuned with full-parameter training.Training uses a learning rate of 5×10−6, batch size 128, and maximum context length 9000 tokens on 8 NVIDIA A100 GPUs.
A.2 Training Configurations: The length-aware DPO objective uses penalty coefficient λ = 1 to explicitly promote concise outputs.This objective is integrated into the modified DPO training pipeline.

A.3 Evaluation Protocols

The evaluation covers four benchmarks with repeated sampling and fixed decoding limits, while LiveCodeBench uses a leakage-controlled date range and lower temperature.

Evaluation spans AIME24, MATH500, AIME25, and LiveCodeBench under a maximum generation limit of 32,768 tokens.
AIME24, MATH500, and AIME25 use temperature 0.6 with results averaged over 16 independent runs.
LiveCodeBench follows DeepSeek-R1's setup, restricts problems to August 1, 2024–January 31, 2025, and uses temperature 0.2 over 16 runs.

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

The proof expands a sequence-level KL divergence into per-token conditional divergences, approximates the expectation with one sampled trajectory, and truncates computation to a fixed token window. The appendix also illustrates refined traces intended to compress reasoning while preserving core logic and improving clarity.

Proof from Equation (1) to Equation (2): Sequence-level KL divergence becomes an expected sum of per-token conditional KL divergences after autoregressive expansion.
Proof from Equation (1) to Equation (2): A single trajectory sampled from A approximates the otherwise intractable expectation over all continuation prefixes.
Proof from Equation (1) to Equation (2): The computation is truncated to a fixed token window L, such as 512 tokens, producing the main-text approximation with a residual term.
Proof from Equation (1) to Equation (2): The residual is considered small in practice because rewriting a reasoning step has diminishing influence over long continuations.
Illustrative refined traces: Illustrative appendix examples cover quadratic equations, constrained letter counting, and circle geometry, showing refined traces intended to preserve core logic and improve clarity.