Technical Summary: Segment-Level Attribution for Selective Learning of Long Reasoning Traces

Problem Statement

Large Reasoning Models (LRMs) achieve high performance by generating extended chains of thought (CoTs). However, these traces often suffer from significant redundancy, where only a small fraction of tokens meaningfully contribute to the final answer, while the majority consists of repetitive content, truncated thoughts, or verbose but uninformative elaborations. This redundancy is problematic because Supervised Fine-Tuning (SFT) on such verbose traces causes models to imitate inefficient patterns, wasting learning capacity on trivial continuations and potentially degrading reasoning performance.

Existing methods for identifying important reasoning parts face limitations:

• Token-level analysis often neglects semantic integrity and fails to interpret redundancy in terms of meaningful reasoning units.
• Indirect segment-level metrics (e.g., perplexity or entropy) are prone to false positives (overemphasizing superficial scaffolding text) and false negatives (filtering out independent verification steps that are low-entropy but crucial for accuracy).
• Pruning-based approaches that compress CoT supervision often compromise reasoning accuracy.

Methodology

The paper proposes a segment-level selective learning framework that leverages Integrated Gradient (IG) attribution to quantify the influence of reasoning segments on final answer prediction.

1. Segment Importance Metrics

The authors partition long CoTs into segments using transition keywords (e.g., "Wait," "Alternatively"). They then aggregate token-level IG attributions into two segment-level metrics:

• Attribution Strength: Measures the overall magnitude of a segment's influence on the model's prediction. It is calculated as the sum of absolute IG values within a segment, normalized by the square root of the segment length to prevent bias toward longer segments.
• Attribution Direction Consistency: Captures whether a segment's attributions are uniformly positive/negative or a mixture of both. It is defined as the ratio of the absolute sum of signed IG attributions to the sum of absolute IG attributions.
  • High Consistency: Indicates shallow or skewed reasoning (e.g., superficial clarifications or severe distortions).
  • Moderate Consistency: Indicates reflective reasoning where the model explores possibilities, refines understanding, and eliminates incorrect content.

2. Important Segment Identification

The framework identifies "important" segments based on a dual-criteria threshold:

1. High Strength: Segments are ranked by normalized attribution strength. The top $k$ segments are selected such that their cumulative strength exceeds a threshold $\tau$ (e.g., 80% of total attribution).
2. Moderate Consistency: Among the top-ranked segments, those with direction consistency below a threshold $\beta$ (e.g., 0.8) are retained. This filters out segments that are high-strength but exhibit extreme uniformity (indicating shallow reasoning).

3. Selective SFT

Instead of compressing the CoT (which disrupts coherence) or training on the full trace (which includes noise), the framework applies Selective SFT.

• The loss function is computed only on tokens belonging to the identified important segments.
• Loss for tokens in unimportant segments is masked.
• This acts as implicit regularization, preventing the model from overfitting to uninformative content while preserving the coherence of the full trajectory during inference.

Key Results

Experiments were conducted across multiple models (R1-Distill-Qwen-1.5B, 7B, and Qwen2.5-7B-Instruct) and datasets (MATH500, AMC23, AIME24, GPQA, Minerva, OlympiadBench).

• Accuracy Improvements: The proposed method consistently outperformed standard Full-CoT SFT. Under greedy decoding, it achieved accuracy gains of up to 4.7% (e.g., on R1-Distill-Qwen-1.5B).
• Efficiency Gains: The approach significantly reduced output length, with token reductions of up to 18.2% compared to full-CoT SFT, while maintaining or improving performance.
• Ablation Studies:
  • Segment vs. Token Level: Segment-level selection outperformed token-level selection (based on IG or random sampling), demonstrating that preserving reasoning units maintains coherence better.
  • Metric Comparison: The IG-based strength-consistency criterion outperformed alternative importance measures such as "First-Correct Solution," "Confidence-Gain," "Segment Perplexity," and "Segment Entropy."
  • Pruning vs. Selective Learning: When applied to pruning (removing unimportant segments entirely), the IG-based method maintained better performance than other pruning strategies, though the authors note that Selective SFT (masking loss) is superior to pruning for preserving accuracy.

Significance and Claims

The paper claims to provide a principled method for distinguishing truly important reasoning segments from various forms of redundancy (repetition, truncation, superficial clarification). By utilizing Integrated Gradients, the method captures both direct and indirect influences of reasoning steps, addressing the limitations of indirect metrics like perplexity.

The primary contribution is a segment-level selective learning framework that enables models to learn more effectively from long reasoning traces. The authors assert that this approach:

1. Improves reasoning efficacy and output efficiency simultaneously.
2. Avoids the accuracy degradation typically associated with CoT compression or pruning.
3. Offers a broader potential for application in other contexts, such as emphasizing policy gradient updates on targeted content in reinforcement learning.

The work suggests that not all reasoning steps are equally valuable; by focusing learning signals on segments with high attribution strength and moderate direction consistency (indicative of reflective reasoning), models can achieve better performance with fewer tokens.