Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Shorter Thoughts, Same Answers: Difficulty-Scaled Segment-Wise RL for CoT Compression

Ye Tian, Aijun Liu

TL;DR

Difficulty-Scaled Segment-Wise GRPO (DSS-GRPO), which decomposes returns into think and answer components, computes group-relative advantages per segment, and routes them with hard token masks so compression updates act only on think while answer alignment acts only on answer.

Abstract

Chain-of-thought (CoT) improves reasoning reliability but increases token cost, motivating post-training compression of explicit reasoning traces. However, the shortest sufficient reasoning is not universal: it depends on difficulty, model capacity, and training state, making fixed length targets brittle. In practice, naive RL-based compression can also undesirably shorten the user-facing answer, because a single completion-level learning signal leaks across the think/answer boundary. We propose Difficulty-Scaled Segment-Wise GRPO (DSS-GRPO), which decomposes returns into think and answer components, computes group-relative advantages per segment, and routes them with hard token masks so compression updates act only on think while answer alignment acts only on answer. DSS-GRPO uses prompt-wise within-group shaping and difficulty-aware scaling to encourage concise reasoning without collapsing answer behavior.

Shorter Thoughts, Same Answers: Difficulty-Scaled Segment-Wise RL for CoT Compression

TL;DR

Difficulty-Scaled Segment-Wise GRPO (DSS-GRPO), which decomposes returns into think and answer components, computes group-relative advantages per segment, and routes them with hard token masks so compression updates act only on think while answer alignment acts only on answer.

Abstract

Chain-of-thought (CoT) improves reasoning reliability but increases token cost, motivating post-training compression of explicit reasoning traces. However, the shortest sufficient reasoning is not universal: it depends on difficulty, model capacity, and training state, making fixed length targets brittle. In practice, naive RL-based compression can also undesirably shorten the user-facing answer, because a single completion-level learning signal leaks across the think/answer boundary. We propose Difficulty-Scaled Segment-Wise GRPO (DSS-GRPO), which decomposes returns into think and answer components, computes group-relative advantages per segment, and routes them with hard token masks so compression updates act only on think while answer alignment acts only on answer. DSS-GRPO uses prompt-wise within-group shaping and difficulty-aware scaling to encourage concise reasoning without collapsing answer behavior.
Paper Structure (37 sections, 14 equations, 2 figures, 3 tables, 1 algorithm)

This paper contains 37 sections, 14 equations, 2 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Related Work
  4. Understanding when CoT helps (and when it does not).
  5. Explicit CoT compression and controllable length.
  6. Implicit and latent CoT for token-efficient reasoning.
  7. CoT distillation, structure, and data quality.
  8. GRPO, group-relative RL, and credit assignment.
  9. Methodology
  10. Segmentation and segment masks
  11. Quality gate: restricting structural rewards to reliable samples
  12. Difficulty-aware scaling
  13. Think compression reward
  14. Answer length alignment reward
  15. Overall objective of Difficulty-Scaled Segment-Wise GRPO
  16. ...and 22 more sections

Figures (2)

  • Figure 1: GSM8K-test length distributions at evaluation $T=0.7$ (Base vs. DDS-GRPO LoRA variants): think (left) and answer (right).
  • Figure 2: GSM8K Example: Base vs DSS-GRPO.