Table of Contents
Fetching ...

Reinforcement Learning for Chain of Thought Compression with One-Domain-to-All Generalization

Hanyu Li, Jiangshan Duo, Bofei Gao, Hailin Zhang, Sujian Li, Xiaotie Deng, Liang Zhao

TL;DR

This paper tackles the inefficiency and latency of long chain-of-thought reasoning in large language models by introducing a sample-level, gated soft compression mechanism that activates only after the model has mastered a problem. By using a mastery gate and per-sample safe-length targets, it penalizes only unnecessary verbosity without truncating unmastered solutions, and it combines this with a dynamic mixture-based training regime within a GRPO-style reinforcement learning framework. The approach yields 20–40% reductions in average response length while maintaining or improving accuracy, and exhibits strong cross-domain generalization from math to tasks like code, instruction following, and general knowledge QA. A stable accuracy→compression→accuracy curriculum with early stopping (pre-collapse optimum) provides a practical, repeatable post-training phase to make reasoning more efficient and generalizable.

Abstract

Chain-of-thought reasoning in large language models often creates an "overthinking trap," leading to excessive computational cost and latency for unreliable accuracy gains. Prior work has typically relied on global, static controls that risk penalizing necessary reasoning. We introduce a sample-level, soft reinforcement learning compression method that penalizes inefficiently long rollouts, but only on problems where the model has already mastered and already produced a more concise rollout. Our experiments show that this method reduces average response length by 20-40% with comparable or higher accuracy. Crucially, the compression exhibits strong cross-domain generalization; a model trained on math spontaneously shortens responses on unseen tasks like code, instruction following, and general knowledge QA, with stable or improved accuracy. We demonstrate a stable post-training curriculum (accuracy-compression-accuracy) that can ultimately produce models that are more accurate and reason more concisely, arguing that such compression method should be a standard phase in developing efficient reasoning models.

Reinforcement Learning for Chain of Thought Compression with One-Domain-to-All Generalization

TL;DR

This paper tackles the inefficiency and latency of long chain-of-thought reasoning in large language models by introducing a sample-level, gated soft compression mechanism that activates only after the model has mastered a problem. By using a mastery gate and per-sample safe-length targets, it penalizes only unnecessary verbosity without truncating unmastered solutions, and it combines this with a dynamic mixture-based training regime within a GRPO-style reinforcement learning framework. The approach yields 20–40% reductions in average response length while maintaining or improving accuracy, and exhibits strong cross-domain generalization from math to tasks like code, instruction following, and general knowledge QA. A stable accuracy→compression→accuracy curriculum with early stopping (pre-collapse optimum) provides a practical, repeatable post-training phase to make reasoning more efficient and generalizable.

Abstract

Chain-of-thought reasoning in large language models often creates an "overthinking trap," leading to excessive computational cost and latency for unreliable accuracy gains. Prior work has typically relied on global, static controls that risk penalizing necessary reasoning. We introduce a sample-level, soft reinforcement learning compression method that penalizes inefficiently long rollouts, but only on problems where the model has already mastered and already produced a more concise rollout. Our experiments show that this method reduces average response length by 20-40% with comparable or higher accuracy. Crucially, the compression exhibits strong cross-domain generalization; a model trained on math spontaneously shortens responses on unseen tasks like code, instruction following, and general knowledge QA, with stable or improved accuracy. We demonstrate a stable post-training curriculum (accuracy-compression-accuracy) that can ultimately produce models that are more accurate and reason more concisely, arguing that such compression method should be a standard phase in developing efficient reasoning models.
Paper Structure (14 sections, 3 equations, 8 figures, 2 tables)

This paper contains 14 sections, 3 equations, 8 figures, 2 tables.

Figures (8)

  • Figure 1: Overview of our compression method across diverse tasks. Each panel shows one benchmark: the orange curve is response length (tokens, normalized to 100% at that benchmark's starting point) and the blue curve is the core metric (accuracy or score), illustrating that compression training only on math problems broadly shortens chains of thought while maintaining or improving downstream performance.
  • Figure 2: Training dynamics of response length (measured in tokens) and mathematical accuracy for our method (ours) versus global soft penalties (DAPO-lite and DAPO-heavy). Left: validation response length in tokens, with absolute values shown in scientific notation on the left axis and percentages relative to the initial length of ours on the right axis. Middle and right: validation accuracies on AIME24 and AIME25 as training progresses, where ours maintains performance while global penalties compress more slowly and offer weaker accuracy--efficiency trade-offs.
  • Figure 3: Training dynamics of the sample-level hard truncation baseline. Top: actor gradient norm and entropy loss, where truncation leads to gradient explosion and collapse of entropy, indicating highly unstable optimization. Bottom: AIME accuracy (average of AIME24 and AIME25) and validation response length in tokens, showing that hard truncation quickly destroys performance while failing to provide a controlled length–accuracy trade-off.
  • Figure 4: Length--accuracy dynamics over two accuracy $\rightarrow$ compression $\rightarrow$ accuracy loops on Qwen3-30B-A3B. The orange curve and left axis show validation response length in tokens (displayed in scientific-notation form with relative percentages with respect to the initial length), while the blue curve and right axis show average AIME accuracy (averaging AIME24 and AIME25). The first loop (150 accuracy steps, 15 compression steps, and 145 accuracy steps) achieves a substantial length reduction with almost no loss in accuracy, and the second loop (10 compression steps and 280 accuracy steps) further improves accuracy while maintaining a net length reduction.
  • Figure 5: Length--accuracy dynamics over a single accuracy $\rightarrow$ compression $\rightarrow$ accuracy loop on MiMo-7B-RL. The orange curve and left axis again show validation response length in tokens (in scientific-notation units and relative percentages), while the blue curve and right axis show average AIME accuracy (averaging AIME24 and AIME25). After roughly 20 compression steps that yield a sizable length reduction with stable accuracy, about 90 subsequent accuracy steps recover and slightly improve accuracy while the length stabilizes at a modest reduction relative to the original.
  • ...and 3 more figures