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

TL;DR

This paper introduces a reinforcement learning-based method to compress chain-of-thought reasoning in large language models, reducing response length significantly while maintaining or improving accuracy and enabling cross-domain generalization.

Contribution

It presents a novel mastery-gated, sample-level reinforcement learning approach for chain-of-thought compression that adapts dynamically and generalizes across multiple domains and tasks.

Findings

Reduces response length by 20-40% with similar or better accuracy.

Enables models trained on math to generalize to code, instruction, and QA tasks.

Significantly decreases reasoning steps in tool-use agents, improving efficiency.

Abstract

Chain-of-thought reasoning in large language models can trigger an "overthinking trap": longer rollouts raise cost and latency yet often yield unreliable accuracy gains. Existing methods use global, static controls that may suppress needed reasoning. We propose mastery-gated, sample-level, soft reinforcement learning compression that penalizes long rollouts only when the model already solves the problem and has produced a shorter rollout. Across benchmarks, it cuts response length by 20-40% with comparable or higher accuracy and generalizes across domains: a model trained on math spontaneously shortens unseen tasks (code, instruction following, general-knowledge QA) without hurting accuracy. We further show two-way transfer between non-agent CoT and tool-use agents: non-agent training reduces SWE-Bench Verified rounds by 13%, while compressing a thinking agent cuts SWE trajectories by 67% tokens and 52% rounds and shortens non-agent outputs by up to 44%. Compression is thus not cosmetic brevity, but an inherent computation policy -- what to keep, and what to forget.