Research map / Large language models / Reasoning and chain of thought

Segment-Level Attribution for Selective Learning of Long Reasoning Traces

ICLR 2026 · Large language models › Reasoning and chain of thought · approach: Long CoT dynamics and efficiency · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes segment-level selective SFT that uses integrated-gradient attribution strength and direction consistency to pick important reasoning segments in long CoTs and masks loss on the rest, improving accuracy and shortening outputs.

Problem

Long chain-of-thought traces from large reasoning models contain much redundancy (repetition, truncation), and full-CoT SFT imitates this verbosity, limiting accuracy gains and efficiency.

Approach

Score each reasoning segment with integrated-gradient attribution using two metrics: length-normalized attribution strength and direction consistency (|sum IG|/sum|IG|). Segments with high strength (top 70-80% cumulative attribution) and moderate consistency (<=0.8) are treated as important; SFT loss is applied only to their tokens.

Key results

• Segment Selective SFT improves R1-Distill-Qwen-1.5B greedy accuracy averaged over six benchmarks from 44.8 (full CoT SFT) to 46.9 (+4.7%) while cutting output length by 18.2% (16520 to 13506 tokens).
• On R1-Distill-Qwen-7B, greedy accuracy rises from 62.1 to 64.5 (+3.9%) with length reduced 12.3%.
• Under temperature sampling, gains are smaller (e.g., +1.6% pass@1 and +1.1% pass@6 on 1.5B), and the method outperforms pruning-based baselines such as First-Correct Solution on accuracy/length trade-off.

Surprising findings

• Pruning unimportant segments (Pruned CoT SFT) reduces accuracy relative to full CoT SFT, whereas selectively masking loss on the same content maintains or improves accuracy.
• Using a finer segmentation keyword list raised greedy accuracy of the 1.5B model to 51.5 while reducing length to 12646.

Limitations

• Evaluated mainly on a small training set (817 LIMO questions) and math-focused benchmarks with a few model families.
• Hyperparameters (tau, beta) and segmentation keywords were tuned on the same analysis setup, and results depend on keyword-based segmentation of reasoning traces.
• Gains under temperature sampling are modest and within noise for some settings.

Experimental setup

Datasets: LIMO

Benchmarks: MATH500 AMC23 AIME24 GPQA-Diamond Minerva OlympiadBench

Metrics: accuracy pass@1 pass@6 output token length

Baselines: Full CoT SFT Pruned CoT SFT First-Correct Solution Confidence-Gain Segments Segment Perplexity Segment Entropy

Builds on: integrated gradients selective SFT

Closest papers by meaning

• Towards Long-Horizon Interpretability: Efficient and Faithful Multi-Token Attribution for Reasoning LLMs (ICML 2026 · similarity 0.89)
• On the Bias of Next-Token Predictors Toward Systematically Inefficient Reasoning: A Shortest-Path Case Study (NeurIPS 2025 · similarity 0.89)
• On Reasoning Strength Planning in Large Reasoning Models (NeurIPS 2025 · similarity 0.88)
• Probing the Trajectories of Reasoning Traces in Large Language Models (NeurIPS 2026 · similarity 0.88)
• Tracing the Traces: Latent Temporal Signals for Efficient and Accurate Reasoning (ICLR 2026 · similarity 0.88)
• Selection, Reflection and Self-Refinement: Revisit Reasoning Tasks via a Causal Lens (ICLR 2026 · similarity 0.88)