TokenSqueeze: Performance-Preserving Compression for Reasoning LLMs

AI summary

Overview

Research area: Efficient inference for large language models, specifically chain-of-thought compression for reasoning LLMs (Long2Short methods).

Technical level: Intermediate. Familiarity with preference optimization (DPO), chain-of-thought reasoning, and KL divergence will help, but the paper's core ideas are accessible.

Scope: The paper proposes TokenSqueeze, a three-stage offline preference-learning pipeline that shortens the reasoning traces of reasoning LLMs by roughly half while preserving or improving accuracy, using only the model's self-generated data.

What This Paper Is About

Reasoning LLMs like DeepSeek-R1 and OpenAI-o1 solve hard problems by writing very long chains of thought, which drives up latency and memory costs. Existing approaches to shorten these traces—either at inference time or via training-time length penalties—typically trade away accuracy, because compressing too aggressively removes essential reasoning steps. TokenSqueeze's goal is to condense reasoning traces while keeping accuracy intact, without relying on external teacher models or human-annotated short answers.

Key Contributions

1. Adaptive reasoning depth selection. A method for choosing which self-generated reasoning traces to keep as positive training examples, where the target length scales with problem difficulty via the quantile formula q = α(1−p), with p being the model's correctness rate on that problem.

2. Intra-step linguistic refinement. A per-step rewriting procedure that resamples candidate versions of each reasoning step and keeps the shortest one that stays within a KL divergence bound, preserving downstream semantics without needing an external LLM.

3. Length-aware preference objective (DPO-L). A DPO variant augmented with an adaptive margin term proportional to the log ratio of rejected-to-preferred response lengths, so pairs with larger compression gains produce stronger preference signals.

4. Composite training objective. A 50/50 combination of the DPO-L loss and supervised fine-tuning on the preferred (short, correct) examples, which stabilizes training and prevents reward collapse on positive samples.

Main Findings

• Accuracy preserved at roughly half the tokens on MATH500. The 7B distilled model reduced average correct-response length by 51.3% while scoring 92.4% accuracy versus the 92.8% baseline. Average length across all responses dropped 51.2%.

• Accuracy sometimes improves, not just holds. On AIME24, TokenSqueeze reached 57.5% accuracy versus the 55.5% baseline while cutting correct-response length by 31.6%. On LiveCodeBench, accuracy rose from 31.3% to 35.0% with 24.4% shorter average outputs.

• Higher AUC across all benchmarks. Since AUC measures accuracy under a 32K token budget, the consistent AUC gains (e.g., 48.5 vs 41.6 on AIME24 for the 7B model) indicate the method is more efficient at every budget level, not just on average.

• Large gains at tight token budgets. Under a 3K-token budget on AIME24, TokenSqueeze scored up to 15.5% higher accuracy than the base model; on MATH500 at 1K tokens, the gap reached 43.1%.

• Beats competing Long2Short methods. Compared with a reproduced Kimi-k1.5 (DPO) and DAST, TokenSqueeze achieved better accuracy at comparable or greater compression, and it outperformed GPT-4o-mini rewriting and TokenSkip pruning in the refinement ablation.

• Compression operates on two axes. Without refinement, the method mainly cuts the number of reasoning steps (267 to 198 on AIME24) while slightly increasing tokens per step. Full TokenSqueeze also shortens individual steps, dividing their lengths by roughly 25% while keeping step count similar.

• 7B model compresses more effectively than 1.5B. The larger model showed stronger compression on most metrics, suggesting the approach benefits from greater base capability.

• Adaptive depth beats fixed or shortest selection. In ablation, the dynamic quantile (Q-DYN) achieved the best accuracy/length trade-off. Selecting the shortest correct trace hurt accuracy, and including longer correct traces as extra negatives also degraded performance. α = 0.2 was the best setting.

• Both training components matter. Pure DPO cut length but dropped accuracy sharply (48.3% on AIME24). Pure SFT held accuracy but compressed weakly. The combined DPO-L + SFT objective gave the best accuracy and compression together.

Methodology in Plain English

TokenSqueeze works entirely offline, in three stages, and never touches data produced by another model.

Stage 1 — Pick the right reasoning depth. The base model is sampled many times per problem. Correct answers are sorted by token length. Rather than always taking the shortest correct trace, the method picks a slice of the shorter end whose size depends on how often the model got the problem right. If the model is usually correct (easy problem), it favors short traces. If it often fails (hard problem), it keeps longer traces to preserve the reasoning needed. Selected traces become positive examples and are paired against longer incorrect traces to form preference pairs, up to 64 pairs per problem.

Stage 2 — Rewrite each step more densely. Each reasoning step is resampled into 64 candidate rewrites, and the shortest candidate is kept—but only if the model's distribution over the following text barely changes. That constraint is measured by a KL divergence between the model's continuation distribution with the original step versus the rewritten one, approximated over a 512-token window. This prevents trimming that would silently change the logic.

Stage 3 — Train with two objectives. The model is fine-tuned with a mix of two losses: a length-aware DPO loss that prefers the short correct response over the long one, with a margin that grows when the length gap is larger; and a standard supervised loss on the short correct responses. The SFT half keeps training stable and prevents the model from degrading the preferred responses while the DPO half pushes toward brevity.

Why This Matters

Impact on research. The paper reframes Long2Short as a preference-learning problem about expression rather than reasoning depth—its central experimental claim is that beyond a certain length, more tokens do not correlate with better accuracy. It shows that a purely offline, self-generated-data pipeline can match or beat online RL with length penalties, avoiding the repeated sampling cost of PPO-style training. It also gives a principled alternative to crude token pruning or LLM-based rewriting that risk breaking logical coherence.

Real-world applications.

• Deploying reasoning models on edge devices or in latency-sensitive customer-facing products where token budgets are tight.
• LLM agent systems that make many sequential decisions, where verbose per-step reasoning compounds into unusable turnaround times.
• Code assistants that need to produce solutions quickly within IDE or CI constraints.
• Cost-sensitive API serving, where halving output tokens directly cuts inference bills at scale.

Industry relevance. Token cost is a dominant operating expense for anyone serving reasoning models. A method that halves output length without accuracy loss is directly commercially meaningful, and the fact that TokenSqueeze needs no external teacher or hand-curated short-answer data lowers the barrier to adopting it on proprietary or domain-specific models.

Future Directions

• Principled hyperparameter tuning. The KL threshold ε (set to 0.005) and other constants were chosen heuristically. The authors propose an adaptive mechanism that sets ε per context based on difficulty and local token divergence.

• Online extension. The current pipeline is fully offline. Moving to an online RL setting where preference generation, policy updates, and reward estimation are jointly optimized could enable continual self-improvement.

• Scaling and generality. The evaluation covers two distilled Qwen models up to 7B on math and code. Whether the approach holds for larger base models, non-distilled reasoning models, and other domains is open.

• Interaction between compression axes. The ablations isolate depth selection and linguistic refinement, but the interplay between step count and step length, and how to balance them optimally, is not fully mapped.

Target Audience

Researchers and engineers working on LLM inference efficiency, reasoning-model deployment, or preference-optimization training. Practitioners who serve reasoning models in production and need concrete token-reduction gains will find the results directly applicable; those studying Long2Short methods or DPO variants will find the methodological framing of interest. Readers should be comfortable with the basics of chain-of-thought reasoning and preference learning to get the most from the training sections.

Authors’ 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.