Segment-Level Attribution for Selective Learning of Long Reasoning Traces

TL;DR

This paper tackles redundancy in long reasoning traces by proposing an integrated gradients–based segment attribution framework. It defines two segment-level metrics—Attribution Strength and Attribution Direction Consistency—to identify important segments, then applies selective supervised fine-tuning (SFT) on those segments while masking unimportant ones. Across multiple models and benchmarks, the approach improves accuracy and reduces output length, outperforming full-CoT SFT and several baselines, with notable gains under greedy decoding. The method offers a principled way to compress reasoning traces without sacrificing performance, with potential for broader applications in reinforcement learning and efficient reasoning.

Abstract

Large Reasoning Models (LRMs) achieve strong reasoning performance by generating long chains of thought (CoTs), yet only a small fraction of these traces meaningfully contributes to answer prediction, while the majority contains repetitive or truncated content. Such output redundancy is further propagated after supervised finetuning (SFT), as models learn to imitate verbose but uninformative patterns, which can degrade performance. To this end, we incorporate integrated gradient attribution to quantify each token's influence on final answers and aggregate them into two segment-level metrics: (1) \textit{attribution strength} measures the overall attribution magnitude; and (2) \textit{direction consistency} captures whether tokens' attributions within a segment are uniformly positive or negative (high consistency), or a mixture of both (moderate consistency). Based on these two metrics, we propose a segment-level selective learning framework to identify important segments with high attribution strength but moderate consistency that indicate reflective rather than shallow reasoning. The framework then applies selective SFT on these important segments while masking loss for unimportant ones. Experiments across multiple models and datasets show that our approach improves accuracy and output efficiency, enabling more effective learning from long reasoning traces~\footnote{Code and data are available at https://github.com/SiyuanWangw/SegmentSelectiveSFT}.

Figures (5)

• Figure 1: Left: An illustrative CoT with important (green blocks) and redundant segments (gray blocks). Our metrics distinguish important from redundant segments (repetitions, truncations, superficial clarifications) with low strength and extremely high consistency. "Attribution Stre." denotes normalized strength across all segments. Right-top: Segment and token counts in correct vs. incorrect CoTs for the same queries. Right-bottom: Cumulative distribution function (CDF) of normalized segment strength in correct and incorrect CoTs, with segment ordered in descending strength.
• Figure 2: Change ($\Delta$) in correct answer confidence across different segment types.
• Figure 3: Log perplexity, entropy and BLEU similarity of important versus unimportant segments.
• Figure 4: Change in correct answer confidence ($\Delta$) across different segment types under varying $\tau$. The red star marks the selected threshold $\tau^*$.
• Figure 5: Model performance under different hyperparameters $\tau$ and $\beta$.