PACE: Prefix-Protected and Difficulty-Aware Compression for Efficient Reasoning

TL;DR

PACE tackles inefficiency and overthinking in Language Reasoning Models by introducing a dual-level compression framework that preserves essential early reasoning through prefix-protected optimization and adapts length penalties to task difficulty. The prefix rollout anchors valid solution paths and decays over time, while the difficulty-aware penalty modulates compression based on empirical task difficulty, yielding strong token efficiency without sacrificing accuracy. Empirical results across in-domain math benchmarks and out-of-domain domains show substantial token reductions (up to ~55%) and accuracy gains, with robust ablations demonstrating the necessity of both levels. The approach generalizes across model sizes and domains, offering a practical method to improve the efficiency–accuracy trade-off in reasoning copilots and CoT systems.

Abstract

Language Reasoning Models (LRMs) achieve strong performance by scaling test-time computation but often suffer from ``overthinking'', producing excessively long reasoning traces that increase latency and memory usage. Existing LRMs typically enforce conciseness with uniform length penalties, which over-compress crucial early deduction steps at the sequence level and indiscriminately penalize all queries at the group level. To solve these limitations, we propose \textbf{\model}, a dual-level framework for prefix-protected and difficulty-aware compression under hierarchical supervision. At the sequence level, prefix-protected optimization employs decaying mixed rollouts to maintain valid reasoning paths while promoting conciseness. At the group level, difficulty-aware penalty dynamically scales length constraints based on query complexity, maintaining exploration for harder questions while curbing redundancy on easier ones. Extensive experiments on DeepSeek-R1-Distill-Qwen (1.5B/7B) demonstrate that \model achieves a substantial reduction in token usage (up to \textbf{55.7\%}) while simultaneously improving accuracy (up to \textbf{4.1\%}) on math benchmarks, with generalization ability to code, science, and general domains.

Figures (12)

• Figure 1: Two limitations of uniform length penalties. (a) Over compression: Using a uniform length penalty will skip steps and lead to a "shortcut" answer. (b) Indiscriminative compression: We report avg pass@1 (32 samples) and split queries into Simple ($>0.75$) and Hard ($\le 0.75$) on MATH500; Hard accuracy drops sharply during training.
• Figure 2: Overview of PACE. PACE consists of two stages: (1) Prefix-Protected Optimization, which anchors early reasoning by generating a short protected prefix with a frozen prefix policy and gradually decays the prefix length; and (2) Difficulty-Aware Penalty, which scales the length penalty by estimated task difficulty.
• Figure 3: Training dynamics of the PACE-7B and its ablations over RL steps, reporting training accuracy, average response length, gradient normalization, and policy entropy.
• Figure 4: Impact of different prefix length on MATH500 dataset.
• Figure 5: Left: Accuracy and average length change across different difficulty levels on MATH500 dataset. Right: The distributional shift of four reasoning behaviors (Advance, Reflect, Refine, Verify) across training steps, with a sample AIME 2024 solution color-coded by behavior. Detailed prompt and explanation can be seen on Appendix \ref{['app:prompt']}.