Technical Summary: Extra-CoT

Problem Statement

Large Language Reasoning Models (LRMs) have demonstrated remarkable capabilities in complex logical inference by generating step-by-step Chain-of-Thought (CoT) reasoning. However, this performance comes at a substantial computational cost, often leading to "overthinking" where models generate redundant reasoning paths even for simple queries. While existing CoT compression methods (e.g., TokenSkip, CTS) can effectively prune tokens at moderate compression ratios (50–60%), they suffer from a critical "fidelity catastrophe" at extreme ratios (e.g., 20–30%). At these levels, general-purpose compressors fail to preserve the sparse, critical reasoning steps and mathematical integrity, leading to a catastrophic loss of logical fidelity and significant performance degradation. The core challenge is achieving high-fidelity reasoning under ultra-low token budgets without sacrificing answer accuracy.

Methodology: The Extra-CoT Framework

The authors propose Extra-CoT, a novel three-stage framework designed to push the boundaries of CoT efficiency while maintaining high-accuracy reasoning. The pipeline integrates a semantic compressor, mixed-ratio supervised fine-tuning (SFT), and a specialized reinforcement learning (RL) algorithm.

1. Semantically-Preserved, Question-Aware Compressor

To generate reliable supervision data for extreme compression, the authors train a dedicated compressor on mathematical CoT data.

• Architecture: Based on Longformer-large-4096, the model utilizes global attention for all tokens in the input question, making the compressor "question-aware." This allows every reasoning token to directly attend to the problem context.
• Annotation Strategy: The training data undergoes a unique index-based, formula-aware annotation process. Mathematical expressions and LaTeX entities are collapsed into atomic units to prevent fragmentation. A teacher model (GPT-4o) is prompted to return only the set of indices to be preserved, ensuring robustness against rendering variations.
• Objective: The compressor is trained with a class-weighted Focal Loss to predict whether to keep or drop each token, prioritizing the preservation of mathematical integrity.

2. Mixed-Ratio Supervised Fine-Tuning (SFT)

A reasoning LLM is fine-tuned on the high-fidelity compressed data to learn robust controllability.

• Control Vocabulary: The tokenizer is augmented with special control tokens (e.g., <COMP 20>, <COMP 40>, ..., <COMP POLICY>).
• Two-Stage Data Cohort:
  • Fixed-Ratio Cohort: Trains the model to strictly adhere to specific compression commands (e.g., <COMP 40>), establishing robust controllability across a spectrum of budgets.
  • Policy Warm-up Cohort: Trains the model to autonomously predict a ratio token when given the <COMP POLICY> trigger, establishing a trainable mechanism for the subsequent RL stage.
• Goal: This stage provides a stable initialization for RL, teaching the model to follow a spectrum of compression budgets rather than just a single target.

3. Constrained and Hierarchical Ratio Policy Optimization (CHRPO)

To explicitly incentivize the selection of lower budgets while maintaining accuracy, the authors introduce CHRPO, a hierarchical RL algorithm.

• Teacher Budget Selection: For each query, a "teacher budget" ($r^*$) is determined as the smallest ratio at which the SFT model still produces a correct answer, ensuring the RL target is achievable.
• Hierarchical Reward Structure:
  • Main Reward ($R_{main}$): Applied to the entire sequence, integrating four components: Accuracy, Ratio-Optimized Mode (rewarding successful shortening, penalizing failure), Budget Calibration (ensuring the realized ratio matches the selected ratio), and Rationale Integrity (penalizing missing reasoning blocks).
  • Control-Head Reward ($R_{ctrl}$): Applied exclusively to the first token (the ratio selection token). This provides an immediate gradient signal to shape the policy's decision-making, addressing the credit assignment problem where the final accuracy signal is too distant from the initial ratio choice.
• Risk-Sensitive Shaping: The reward functions employ asymmetric penalties to ensure the policy only shortens when correct ("Safe-Shortening") and prefers to increase the budget rather than fail on hard instances ("Fail-Fast Recovery").

Key Contributions

1. Extra-CoT Framework: A novel three-stage pipeline (Compressor $\to$ Mixed-Ratio SFT $\to$ CHRPO) that achieves high-fidelity reasoning at ultra-low token budgets, effectively alleviating the challenges of semantic preservation and control adherence.
2. Semantically-Preserved Compressor: A question-aware architecture with formula-aware annotation that generates high-fidelity supervision data, preserving mathematical integrity and preventing the fragmentation of symbolic reasoning components.
3. Unified SFT-RL Pipeline: A training strategy that uses mixed-ratio SFT to establish robust controllability, followed by CHRPO with hierarchical rewards to explicitly incentivize accuracy under extreme compression ratios.

Experimental Results

The framework was evaluated on three mathematical reasoning benchmarks: GSM8K, MATH-500, and AMC2023, primarily using the Qwen3-1.7B backbone.

• Performance at Extreme Ratios: On MATH-500, Extra-CoT achieved a 73% token reduction (0.27 ActRatio) with an accuracy improvement of 0.6% over the base model, significantly outperforming state-of-the-art methods like TokenSkip. At a 0.2 target ratio, Extra-CoT maintained robust performance (47.8% accuracy on MATH-500), whereas TokenSkip suffered a catastrophic drop to 23.4%.
• Control Adherence: Unlike baselines that exhibit "control collapse" (where realized ratios deviate significantly from targets), Extra-CoT demonstrated precise adherence to compression commands across all ratios.
• Efficiency: In end-to-end inference latency tests, Extra-CoT reduced latency by 3.24× on GSM8K and 2.35× on MATH-500 compared to the base model. In contrast, TokenSkip showed marginal gains or even increased latency on some benchmarks.
• Generalization: The method demonstrated robustness across different backbone models (Qwen2.5-7B, Llama3.2-3B) and long-context settings (Pangu-Embedded-7B-V1.1 with 16K context), confirming that the improvements stem from the quality of compressed supervision rather than specific model architectures.

Significance and Claims

The paper claims that Extra-CoT demonstrates that high-fidelity extreme compression is a viable path to efficient reasoning. By addressing the "fidelity catastrophe" through semantically-preserved supervision and hierarchical RL optimization, the framework enables Large Language Reasoning Models to operate effectively under ultra-low token budgets without the performance collapse seen in existing methods. The work highlights that preserving the logical integrity of reasoning chains is critical for maintaining problem-solving capabilities, even when the token count is drastically reduced. The authors provide their code to facilitate further research in this direction.