Segment-Level Attribution for Selective Learning of Long Reasoning Traces

AI summary

Overview

Research area: Natural language processing, specifically efficient training of large reasoning models (LRMs) on long chains of thought (CoTs).

Technical level: Advanced. The work combines integrated gradients attribution with a modified supervised fine-tuning loss, so it assumes familiarity with transformer training, attribution methods, and SFT objectives.

Scope: A single paper-length study proposing two segment-level attribution metrics, an importance-based selection rule, and a selective SFT procedure evaluated on three base models over six mathematical and scientific reasoning benchmarks.

What This Paper Is About

Large reasoning models produce very long chains of thought, but the paper states that only a small fraction of those tokens meaningfully contributes to the final answer, while the rest is repetitive, truncated, or low-value clarification. When models are fine-tuned on these full traces, they imitate the verbose and uninformative patterns, which can reduce both accuracy and output efficiency. The paper's goal is to measure which segments of a long reasoning trace actually influence the correct answer, and then train only on those segments.

Key Contributions

1. Two segment-level attribution metrics derived from integrated gradients. "Attribution strength" sums absolute token-level IG values inside a segment with square-root length normalization (then normalized across segments); "direction consistency" is the ratio of the absolute signed sum to the absolute-value sum, capturing whether a segment pushes uniformly in one direction or mixes supportive and corrective contributions.

2. A criterion for identifying important segments. Important segments are the top-ranked segments by normalized strength whose cumulative strength reaches a threshold τ (they use τ = 0.7), restricted to those with direction consistency at or below β (they use β = 0.8), on the reasoning that high strength combined with moderate consistency indicates reflective rather than shallow reasoning.

3. A selective SFT framework. Rather than pruning the trace, the method keeps the full trajectory intact but masks the loss for tokens in unimportant segments, so parameter updates come only from important segments.

4. Empirical validation and analysis. The paper documents redundancy in long CoTs (30%∼40% of segments carry over 80% of total attribution), characterizes important versus unimportant segments, and compares the method against full-CoT SFT, pruning-based SFT, and four alternative importance measures.

Main Findings

• Heavy attribution concentration: Only 30%∼40% of segments accumulate over 80% of total attribution in both correct and incorrect CoTs, confirming substantial redundancy.

• Important segments raise answer confidence more: Sequentially appending segments shows that segments with high attribution strength and moderate direction consistency produce significantly larger gains in correct-answer confidence than other segment types.

• Important segments have lower perplexity and entropy: Aggregated at the segment level, important segments consistently show significantly lower log perplexity and entropy than unimportant ones, suggesting more constrained deduction-like content.

• Unimportant segments are repetitive and truncated: Unimportant segments have higher BLEU similarity against preceding segments (including substantially more content with BLEU > 0.8), and a Qwen3-8B judge classified 49% of unimportant segments as truncated versus only 26% of important segments.

• Selected thresholds: A greedy search selected τ = 0.7, under which about 33% of segments (45% of tokens) are labeled important, and β = 0.8 was identified as the optimal consistency threshold. Incorrect CoTs are typically correlated with more segments and tokens than correct CoTs for the same queries.

• Accuracy and efficiency gains over full-CoT SFT (greedy decoding, "Overall" column): R1-Distill-Qwen-1.5B improved from 44.8 to 46.9 (↑4.7%) with output length dropping from 16520 to 13506 tokens (↓18.2%); R1-Distill-Qwen-7B from 62.1 to 64.5 (↑3.9%) with length 9693→8499 (↓12.3%); Qwen2.5-7B-Instruct from 44.2 to 45.6 (↑3.2%) with length 10317→9852 (↓4.5%).

• Gains under temperature sampling are smaller but present: For R1-Distill-Qwen-1.5B, pass@1 rose from 50.9 to 51.7 (↑1.6%) and pass@6 from 72.4 to 73.2 (↑1.1%) with length 10043→9388 (↓6.5%); R1-Distill-Qwen-7B reached 65.8 pass@1 (↑0.5%) and 80.0 pass@6 (↑1.0%) with length ↓2.8%; Qwen2.5-7B-Instruct reached 45.0 pass@1 (↑3.7%) and 66.5 pass@6 (↑0.8%) with length ↓7.1%.

• Pruning hurts, selective masking helps: On R1-Distill-Qwen-1.5B, Pruned CoT SFT scored 43.9 greedy / 49.2 pass@1 versus 44.8 / 50.9 for full CoT SFT, while selective SFT with the paper's important segments scored 46.9 / 51.7.

• Segment-level beats token-level and random selection: Selective SFT with random segments (45.1 / 50.8), top 45% tokens by absolute IG (46.1 / 50.8), top 45% tokens by original IG values (45.2 / 50.3), and high-strength segments alone (46.9 / 51.4) all trailed or matched the full criterion (46.9 / 51.7), with the full criterion achieving the greatest token reduction (13506 greedy, 9388 sampling).

• Better than alternative importance measures averaged across datasets: The IG-based segments (46.9 / 51.7) outperformed First-Correct Solution (46.2 / 51.0), Confidence-Gain Segments (44.7 / 50.3), Segment Perplexity (44.7 / 50.2), and Segment Entropy (44.5 / 51.2). Confidence-Gain Segments identified only 24% of segments on average.

• Positional analysis (appendix): 40% of important segments appear after the decision segment, while 57% of unimportant segments occur before it; most low-strength unimportant segments (64%) occur before the decision segment, and most high-consistency unimportant segments (72%) appear after it.

Methodology in Plain English

The researchers first split each long chain of thought into segments using natural transition keywords such as "Wait" and "Alternatively." They then ask how much each token inside those segments changes the model's probability of producing the correct answer, using integrated gradients — a technique that traces the model's output sensitivity from a neutral baseline embedding to the actual input embedding along a straight interpolation path, and averages gradients over a number of interpolation steps.

Each token's attribution is aggregated to the segment level in two ways. Strength is the sum of absolute attributions divided by the square root of the segment length, which keeps long segments from being favored automatically; these strength scores are then normalized across all segments in the trace. Consistency is the magnitude of the signed attribution sum divided by the sum of absolute attributions, so a value near 1 means the segment pushes entirely one way and a moderate value means it mixes positive and negative contributions. Absolute values are used because negatively attributed tokens may still represent necessary exploration.

Selection is greedy on strength: sort segments by normalized strength, take the smallest prefix whose cumulative strength reaches τ, then keep only those in the prefix whose consistency is at most β. Training then uses a standard cross-entropy loss, but with an indicator that zeroes out the contribution of tokens in unselected segments. The paper compares this against full-CoT SFT, against pruning the unimportant segments away, and against four other importance measures, using R1-Distill-Qwen2.5-1.5B, R1-Distill-Qwen2.5-7B, and Qwen2.5-7B-Instruct trained on 817 LIMO math questions with full-parameter fine-tuning at a maximum sequence length of 16384. Evaluation covers MATH500, AMC23, AIME24 in-domain and GPQA-Diamond, Minerva, and OlympiadBench out-of-domain, under greedy decoding and temperature sampling (temperature=0.6, top-p=1.0) with zero-shot CoT and a maximum response length of 32,768 tokens.

Why This Matters

The work reframes "efficient reasoning" as a data-selection problem: instead of throwing away parts of a long trace, it changes which tokens the loss sees. That distinction matters because the paper shows pruning degrades accuracy while masking does not, and because the attribution signal can in principle be reused wherever per-token credit assignment is needed.

Real-world applications:

• Training smaller reasoning models on distilled long traces, where teacher outputs are verbose and expensive to imitate faithfully.
• Reducing inference cost in deployed reasoning assistants, since the paper reports output length reductions alongside accuracy gains.
• Curriculum or reward shaping for reinforcement learning, which the authors explicitly flag as a possible extension for emphasizing policy updates on important content.
• Quality auditing of synthetic reasoning data, where the strength/consistency metrics can flag repetitive or truncated segments before a dataset is used.

Industry relevance: any organization fine-tuning reasoning models on long chain-of-thought data faces the tradeoff between supervision quality and token cost. This paper offers a measurable criterion for choosing which parts of a trace to learn from, and reports both accuracy and token-length deltas, which are the two numbers that most directly affect serving cost.

Future Directions

• Applying the identification to reinforcement learning: the conclusion explicitly raises emphasizing policy gradient updates on important content as an extension.
• Reducing reliance on integrated gradients: IG requires many forward passes along an interpolation path; cheaper or approximated attribution signals with comparable selection quality are an open question.
• Testing beyond math and science reasoning: the experiments use LIMO math questions and six math/science benchmarks; generalization to other domains is not reported.
• Understanding the moderate-consistency criterion better: the paper ties moderate consistency to "reflective" reasoning, and a finer-grained account of which reasoning behaviors fall into that band is left open, along with the appendix's observation that high-consistency segments after the decision point mostly reinforce an already-found answer.

Target Audience

Researchers and engineers working on reasoning model training, chain-of-thought distillation, and supervised fine-tuning efficiency. It will be most useful to readers comfortable with attribution methods and loss-level training interventions, and to practitioners deciding how to curate long reasoning traces for fine-tuning.

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