Our new X account is live! Follow @wizwand_team for updates
WorkDL logo mark

Segment-Level Attribution for Selective Learning of Long Reasoning Traces

About

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}.

Siyuan Wang, Yanchen Liu, Xiang Ren• 2026

Related benchmarks

TaskDatasetResultRank
Mathematical ReasoningMinerva--
138
Scientific ReasoningGPQA Diamond
Accuracy28.8
45
Mathematical ReasoningOlympiadBench
Pass@141.9
12
Mathematical ReasoningMATH 500
Accuracy95.2
9
Mathematical ReasoningAMC 23
Accuracy90
9
Mathematical ReasoningOverall
Length (tokens)7.71e+3
6
Mathematical ReasoningLIMO
Length1.60e+4
6
Mathematical ReasoningAIME24
Accuracy56.7
2
Showing 8 of 8 rows

Other info

Follow for update