TokenSqueeze: Performance-Preserving Compression for Reasoning LLMs

TL;DR

TokenSqueeze tackles the inefficiency of long chain-of-thought reasoning by learning succinct yet faithful reasoning traces from self-generated data. It introduces adaptive reasoning-depth selection, intra-step linguistic refinement under distributional alignment with KL constraints, and a length-aware Direct Preference Optimization objective to promote concise but correct reasoning. The method achieves substantial token reduction (e.g., up to 50% on MATH500) while maintaining or improving accuracy across math and programming benchmarks, and yields higher AUC under limited token budgets. By avoiding external datasets or teacher models and providing an offline training pipeline, TokenSqueeze offers a practical, scalable approach to efficient reasoning in large language models.

Abstract

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.

Figures (7)

• Figure 1: Overview of the TokenSqueeze. It first selects self-generated responses based on adaptive reasoning depth, then rewrites them under KL constraints to improve information density, and finally trains the model with a composite objective that promotes accuracy and brevity.
• Figure 2: Example of reasoning trace before and after intra-step linguistic refinement. The refined version achieves higher information density while preserving logical integrity.
• Figure 3: TokenSqueeze outperforms the base model across token budgets on AIME24 and MATH500, with up to 15.5% higher accuracy on AIME24 (3K tokens) and 43.1% on MATH500 (1K tokens).
• Figure 4: Effect of the $\alpha$ parameter in adaptive quantile selection. Moderate values (e.g., $\alpha=0.2$) provide the best balance between accuracy and token efficiency, while extreme values degrade performance.
• Figure 5: Reasoning trace comparison for a quadratic equation problem before and after refinement.