Compress the Easy, Explore the Hard: Difficulty-Aware Entropy Regularization for Efficient LLM Reasoning

Technical Summary: Compress the Easy, Explore the Hard (CEEH)

Problem Statement

Chain-of-Thought (CoT) prompting has significantly enhanced Large Language Models' (LLMs) reasoning capabilities by externalizing intermediate steps. However, this explicit reasoning often results in verbose outputs, leading to prohibitive inference latency and computational costs. Existing reasoning compression methods, which typically employ Reinforcement Learning (RL) with length constraints (e.g., explicit length penalties), face a critical failure mode: aggressively optimizing for brevity triggers rapid entropy collapse.

As the policy becomes overly deterministic to satisfy length constraints, the exploration space shrinks prematurely. This prevents the model from discovering valid reasoning paths for challenging problems that require extensive deduction, self-reflection, or error correction. Consequently, these methods often sacrifice reasoning accuracy for brevity. A central tension exists: uniform entropy regularization can sustain exploration but induces longer trajectories (counteracting compression), while uniform length penalties compress trajectories but collapse entropy (degrading accuracy).

Methodology: CEEH

The authors propose CEEH (Compress the Easy, Explore the Hard), a difficulty-aware framework that dynamically balances exploration and compression. The method integrates two synergistic components within a Group Relative Policy Optimization (GRPO) setting:

1. Difficulty-Aware Entropy Regularization
Instead of applying uniform regularization, CEEH selectively sustains exploration on "hard" questions while allowing aggressive compression on "easy" ones.

• Dynamic Difficulty Estimation: The system maintains a historical accuracy estimate for each question using an Asymmetric Exponential Moving Average (EMA). A question is classified as "hard" ($D_h$) if its historical accuracy falls below the global dataset average. This threshold evolves as the model improves.
• Selective Mechanism:
  • For Hard Questions ($D_h$): The method applies stronger entropy regularization. This is achieved by amplifying the regularization coefficient (e.g., a $5\times$ multiplier in the cosine annealing schedule) and incorporating an entropy-based advantage term ($\psi(H)$) that incentivizes the generation of high-entropy tokens. This preserves a diverse search space to prevent premature convergence.
  • For Easy Questions: Regularization is relaxed, allowing the policy to confidently exploit shorter, well-established reasoning paths.

2. Dynamic Optimal-Length Penalty
To counteract the potential length inflation caused by entropy regularization on hard questions, CEEH introduces a dynamic, question-level length penalty.

• Mechanism: The system tracks the historically shortest correct response length ($L_x$) for each specific question.
• Penalty Function: A penalty is applied only to correct responses that exceed this evolving baseline. The penalty is calculated as the relative deviation: $R_{len} = (L_{x, y_i} - L_x) / L_x$.
• Adaptive Signal: This design ensures the compression signal remains stable and discriminative even as the model's capabilities improve and response lengths naturally decrease. It prevents the model from exploiting the entropy regularization to generate unnecessarily long outputs on difficult problems.

Training Dynamics
The total reward combines the verifiable reward (accuracy) and the length penalty: $R_{total} = R - \beta R_{len}$. The method employs specific implementation strategies, such as increasing the upper clipping threshold (Clip-Higher) to encourage exploration and avoiding advantage standardization to prevent the attenuation of the length penalty signal.

Key Contributions

1. Identification of Entropy Collapse: The paper provides an entropy-based perspective, identifying rapid entropy decay as a primary factor causing the accuracy-length trade-off failure in existing reasoning compression methods.
2. CEEH Framework: A novel framework integrating difficulty-aware entropy regularization (to sustain exploration on hard tasks) and a dynamic optimal-length penalty (to stabilize training and enforce compression).
3. Empirical Validation: Demonstrates that CEEH achieves a superior accuracy-efficiency trade-off compared to strong baselines, effectively compressing reasoning without sacrificing correctness.

Experimental Results

The method was evaluated on six mathematical reasoning benchmarks (GSM8K, MATH500, AIME24/25, AMC, OlymBench) using R1-Distill-Qwen2.5 models.

• Accuracy vs. Length: CEEH consistently reduced response token counts by over 30% compared to base models while maintaining accuracy comparable to, or slightly better than, the base model.
• Comparison to Baselines:
  • Prompting/Offline Methods: Often showed fluctuating accuracy or failed to generalize across diverse difficulty levels.
  • Length-Only RL (e.g., Length-Penalty): While effective at reducing length, these methods frequently degraded Pass@k performance, indicating a loss in underlying reasoning capability due to entropy collapse.
  • CEEH: Maintained or improved Pass@k scores, suggesting that the method preserves the model's ability to discover correct solutions through controlled exploration.
• Robustness: The method generalized well across models of different scales (1.5B and 7B parameters) and diverse reasoning tasks.
• Training Efficiency: CEEH achieved higher training accuracy and better data efficiency compared to uniform penalty strategies, attributed to the high-entropy exploration on challenging questions.

Significance and Claims

The paper claims that CEEH resolves the tension between reasoning compression and capability preservation by explicitly separating where to compress and where to explore.

• Preservation of Reasoning Capability: By preventing entropy collapse on difficult instances, CEEH ensures that the model retains the ability to perform complex deductions and self-correction, rather than merely becoming confident in suboptimal, shortcut solutions.
• Stable Compression: The dynamic optimal-length penalty provides a robust signal that adapts to the model's improving capabilities, preventing the "length inflation" often seen when entropy regularization is applied uniformly.
• Generalizability: The approach demonstrates that efficiency gains in LLM reasoning do not necessitate a trade-off with accuracy, provided that exploration is managed adaptively based on instance difficulty.

The work concludes that principled entropy control during compression is essential for maintaining genuine reasoning capabilities in RL-based LLM optimization.