Towards Efficient Large Language Reasoning Models via Extreme-Ratio Chain-of-Thought Compression

AI summary

Overview

Research area: Efficient inference for large language reasoning models, specifically extractive chain-of-thought (CoT) compression on mathematical reasoning tasks.

Technical level: Advanced. The paper assumes familiarity with supervised fine-tuning, reinforcement learning from verifiable rewards, token-classification compressors, and reward shaping.

Scope: The paper introduces Extra-CoT, a three-stage pipeline (a question-aware CoT compressor, mixed-ratio SFT, and a hierarchical-reward RL stage called CHRPO) that aims to compress reasoning traces to extreme ratios while preserving or improving answer accuracy.

What This Paper Is About

Large reasoning models solve problems by writing out long step-by-step reasoning traces, but those traces consume a lot of computation at inference time, and models often "overthink" simple questions. Prior compression methods work reasonably at moderate budgets (roughly 50-60% of tokens) but collapse catastrophically at extreme ratios such as 20-30% of tokens, because they discard sparse but critical reasoning steps and lose logical fidelity. This paper builds a compression system that keeps mathematical formulas intact and teaches a model to choose very small reasoning budgets while still getting answers right.

Key Contributions

1. The Extra-CoT framework, a three-stage pipeline for extreme-ratio CoT compression that targets high-fidelity reasoning at ultra-low token budgets, addressing both semantic preservation and adherence to control commands.
2. A semantically-preserved, question-aware CoT compressor that marks question tokens with global attention and uses an index-based, formula-aware annotation method to produce high-fidelity supervision that keeps mathematical expressions atomic.
3. A unified SFT-RL training pipeline, in which a mixed-ratio SFT stage installs multi-ratio controllability and a trainable policy token, followed by the CHRPO algorithm with hierarchical rewards that explicitly incentivize accuracy at extreme compression ratios.
4. Controlled empirical validation across multiple mathematical reasoning benchmarks, backbones, and context lengths, including LLM-judge scoring and a human A/B preference study of compressor output quality.

Main Findings

• Large accuracy gains at extreme ratios. On MATH-500 at a target ratio of 0.2, Extra-CoT reaches 47.8 Acc@all, a +24.4 point gain over TokenSkip. On GSM8K at ratio 0.2 it reaches 80.2 accuracy (+21.1 points over TokenSkip), and on AMC2023 at ratio 0.2 it leads by +7.5 points. Even at ratio 0.8, Extra-CoT still gains +2.4 points on GSM8K over TokenSkip.
• Headline efficiency claim. On MATH-500 using Qwen3-1.7B, Extra-CoT achieves over 73% token reduction alongside a 0.6% accuracy improvement, and the paper reports it significantly outperforms state-of-the-art methods.
• Policy results beat the RL baseline on efficiency. The Extra-CoT (CHRPO) policy on MATH-500 reaches 64.8% accuracy at a 0.27 ActRatio, outperforming the Thinkless (DeGRPO) re-implementation by +1.2 points. On AMC2023 it reaches 50.0% accuracy at a 0.32 ActRatio, matching Thinkless while using less than half the tokens. On GSM8K it matches Thinkless's accuracy while using fewer tokens.
• Better ratio adherence, less "control collapse." TokenSkip shows nearly identical realized ratios for its 0.8 and 0.6 targets on MATH-500 (0.91 vs. 0.90), and its 0.6 GSM8K target overshoots to roughly 0.84. Extra-CoT reaches roughly 0.73 on GSM8K at the 0.6 target and clearly separates the 0.8 and 0.6 targets on MATH-500 (0.90 vs. 0.77).
• Compressor quality is the suspected root cause of control collapse. The paper argues the baseline model learns that low-ratio supervision is "untrustworthy" and disobeys commands, whereas Extra-CoT's high-fidelity data lets the model honor extreme-ratio commands.
• Both CHRPO reward components matter. Ablating both the mode reward and the control-head reward raises GSM8K average token count from 210 to 324 and gives 85.7 accuracy; ablating only the control-head reward gives 258 tokens and 82.1 accuracy; the full CHRPO gets 210 tokens with 85.8 accuracy. On MATH-500 the same ablation ladder is 604 tokens / 60.0 accuracy, 568 / 59.6, and 452 / 64.8 for full CHRPO.
• Robustness across backbones (SFT-level). On GSM8K with Qwen2.5-7B-Instruct, Extra-CoT at ratio 0.5 uses 189 tokens at 0.63 ActRatio with 89.4 accuracy, versus TokenSkip at 151 tokens, 0.51 ActRatio, 86.0 accuracy. With Llama3.2-3B-Instruct, Extra-CoT at ratio 0.5 uses 115 tokens at 0.50 ActRatio with 73.4 accuracy, versus TokenSkip at 123 tokens, 0.50 ActRatio, 72.5 accuracy.
• Compressor beats LLMLingua-2 on judged quality. Scored 1-5 by DeepSeek-R1 and Qwen3-32B on Math Fidelity, Reasoning Coherence, and Clarity & Readability across four ratios, the Extra-CoT compressor outperforms LLMLingua-2, with the largest gaps at the extreme 0.2 and 0.4 ratios. A blind A/B human study with five expert raters (250 pairwise votes per budget across four budgets, 1,000 votes total) shows the same trend, with strong preference for Extra-CoT under aggressive compression.
• End-to-end latency. The paper states that the practical efficiency gain is reflected in end-to-end latency results in Table 5; those specific latency numbers are not included in the provided paper content.
• Data scale. SFT data comes from MetaMathQA-395K: a fixed-ratio cohort of 60k samples (12k per ratio bucket) spanning ratios 0.2 to 1.0, plus a 10k-sample <COMP_POLICY> warm-up cohort. The compressor is trained on approximately 30K mathematical CoT examples generated by Qwen3-32B on the CAMEL dataset.

Methodology in Plain English

The pipeline has three linked stages.

Stage 1: Build a specialized compressor that understands math. The researchers take roughly 30K math CoT examples and break them into word-level spans, detecting LaTeX entities and inline math expressions and collapsing each into a single atomic unit so formulas cannot be fragmented. They then ask a teacher model (GPT-4o) to return only the indices of the spans that must be kept, given the original question. This index-only approach is designed to be robust to minor paraphrasing of formulas. Those span decisions are mapped back to token-level keep/drop labels. The compressor itself is a Longformer-large-4096 backbone with all question tokens set to global attention, so every CoT token can attend directly to the question — this is what makes it "question-aware." It is trained with a class-weighted Focal Loss over valid CoT tokens. At inference, tokens are sorted by keep-probability and selected greedily until the budget (floor of γ times the original think-only length) is met.

Stage 2: Teach the model to follow ratio commands. Six special tokens are added to the tokenizer: <COMP_20>, <COMP_40>, <COMP_60>, <COMP_80>, <COMP_100>, and <COMP_POLICY>. Under a Prefix-Mirror Protocol, the model receives the question plus a control token, and outputs a mirrored ratio token followed by the compressed rationale. If the input is a fixed-ratio token, the model copies it verbatim; if the input is <COMP_POLICY>, the model must first choose a ratio on its own. Training on a spectrum of budgets (not just the target ratio) is described as pre-conditioning that prevents collapse and makes command-following reliable.

Stage 3: Reinforce the choice of budget. The RL data uses a monotonically correct teacher rule: the teacher budget is the smallest ratio at which the SFT model answers correctly provided all larger ratios also answer correctly, which filters out unstable instances. CHRPO then splits rewards into two locations. A main reward applied at the end of the sequence combines an accuracy reward of +1/-1, a mode reward that pays only for successful shortening (with a small cap κ, e.g. 2) and penalizes shortening that fails, a Huber-like budget-calibration reward with tolerance ε, and a penalty for failing to produce a <think> block; the sum is clipped to [-1, 1]. A separate control-head reward is applied only at the first token to give the ratio-selection policy a direct, immediate signal. The asymmetric penalties are designed to give three guarantees: shortening is only rewarded when correct, failing while shortening is penalized more than failing while being too verbose (so the policy "fails safe" and does not get stuck at short budgets), and the reward for successful compression is capped so progress toward extreme ratios is gradual.

Evaluation setup. Experiments use a matched-ratio protocol, comparing methods at identical target ratios and reporting the resulting ActRatio. All token counts use think-only accounting: only tokens strictly inside <think>...</think> blocks count. Accuracy is measured over the full test set. A single evaluation script handles answer extraction and correctness checking across all methods, and a fixed random seed of 42 is used for sampling, initialization, and decoding.

Why This Matters

Impact on research. The paper reframes extreme-ratio CoT compression as a supervision-quality and reward-design problem rather than a token-scoring problem. It shows that a domain-aware, question-conditioned compressor produces training data that a downstream model actually trusts, and that hierarchical reward placement can overcome the instability that makes aggressive compression risky. It also introduces a matched-ratio evaluation discipline (reporting ActRatio alongside accuracy) that exposes "control collapse" — a failure mode where models quietly ignore their compression commands, which accuracy numbers alone would hide.

Real-world applications.

• Serving reasoning models under tight inference budgets, such as on-device or edge assistants where every generated token costs latency and battery.
• High-volume math and STEM tutoring systems that must answer many student queries per second without sacrificing step-by-step correctness.
• Cost control in API-based or cloud-hosted reasoning services, where token count maps directly to billing.
• Long-context reasoning deployments, where the paper validates a 16K-context setting on Pangu-Embedded-7B-V1.1 and evaluates out-of-domain STEM performance on the MMLU STEM subset.

Industry relevance. The author list includes Huawei Foundation Model Dept. researchers, and the paper discloses that an auxiliary long-context evaluation uses Pangu-Embedded-7B-V1.1. The work targets a direct production concern — reducing reasoning token cost without losing answer quality — and is evaluated across Qwen3-1.7B, Qwen2.5-7B-Instruct, Llama3.2-3B-Instruct, and Pangu-Embedded-7B-V1.1, making it applicable to both small on-device models and larger served models. Code is released at https://github.com/Mwie1024/Extra-CoT.

Future Directions

• Beyond mathematics. The compressor supervision relies on math-specific formula-aware annotation and LaTeX detection. Extending the same index-based annotation approach to code, formal logic, or multi-step scientific reasoning is an open question the paper does not address.
• Latency and deployment measurements. The paper points to Table 5 for end-to-end latency, but the specific numbers are not present in the provided content, so the actual wall-clock benefit of the token savings remains to be quantified in the material available here.
• Compressor scope and cost. The compressor is a Longformer-large-4096 backbone trained on roughly 30K examples, and the supervision data is produced by GPT-4o with Qwen3-32B-generated CoTs. Whether cheaper supervision or a smaller compressor preserves the fidelity gains is left open.
• Reward design transfer. CHRPO's asymmetric, two-location reward scheme is tuned for ratio selection under a discrete grid of {0.2, 0.4, 0.6, 0.8, 1.0}. Whether the same risk-sensitive shaping generalizes to other controllability axes, such as tool-use budgets or adaptive test-time compute, is not explored.
• OOD behavior. Full results for SVAMP, MultiArith, MetaMath-1k, and the OOD benchmarks are placed in the appendix and are not included in the provided content, so the extent of out-of-domain generalization is not visible here beyond the stated evaluation design.

Target Audience

Researchers and engineers working on inference efficiency for large reasoning models, especially those interested in controllable CoT compression, reward shaping for RL post-training, and mathematical reasoning benchmarks. It is also relevant to practitioners deploying reasoning models under strict token, latency, or cost constraints, and to readers studying how supervision data quality propagates into downstream model behavior. The paper is not introductory: it assumes comfort with extractive token classification, SFT data construction, and policy-gradient-style optimization with shaped rewards.

Authors’ abstract

Chain-of-Thought (CoT) reasoning successfully enhances the reasoning capabilities of Large Language Models (LLMs), yet it incurs substantial computational overhead for inference. Existing CoT compression methods often suffer from a critical loss of logical fidelity at high compression ratios, resulting in significant performance degradation. To achieve high-fidelity, fast reasoning, we propose a novel EXTreme-RAtio Chain-of-Thought Compression framework, termed Extra-CoT, which aggressively reduces the token budget while preserving answer accuracy. To generate reliable, high-fidelity supervision, we first train a dedicated semantically-preserved compressor on mathematical CoT data with fine-grained annotations. An LLM is then fine-tuned on these compressed pairs via a mixed-ratio supervised fine-tuning (SFT), teaching it to follow a spectrum of compression budgets and providing a stable initialization for reinforcement learning (RL). We further propose Constrained and Hierarchical Ratio Policy Optimization (CHRPO) to explicitly incentivize question-solving ability under lower budgets by a hierarchical reward. Experiments on three mathematical reasoning benchmarks show the superiority of Extra-CoT. For example, on MATH-500 using Qwen3-1.7B, Extra-CoT achieves over 73\% token reduction with an accuracy improvement of 0.6\%, significantly outperforming state-of-the-art (SOTA) methods. Our source codes have been released at https://github.com/Mwie1024/Extra-CoT.