Segment-Level Attribution for Selective Learning of Long Reasoning Traces
Authors: Siyuan Wang
Yanchen Liu
Xiang Ren
Venue: ICLR 2026 Poster (2026)
Keywords: long CoTs
selective learning
integrated gradient
segment attributions
Primary area: foundation or frontier models, including LLMs

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.

Published as a conference paper at ICLR 2026
SEGMENT-LEVEL ATTRIBUTION FOR SELECTIVE LEARNING OF LONG REASONING TRACES
Siyuan Wang, Yanchen Liu, Xiang Ren
University of Southern Californa {sw_ 641, 1241,1iuyee 1iuyanch, xichh,,iian xiangren] @usc. edu
ABSTRACT
Large Reasoning Models (LRMs) achieve strong reasoning performance by gen- erating long chains of thought (CoTss, yet only small fraction of these traces meaningfully contributes to answer prediction, while the majority contains repet- itive 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, incorporate integrated gradient attribution to quantify each token's 's influence on final answers and ag- gregate them into segment-level metrics: (I) attribution strength measures the overall attribution magnitude; and (2) direction consistency captures whether tokens' attributions within segment are uniformly positive negative (high con- sistency), or a mixture of both (moderate consistency). Based on these two metrics, we propose a segment-level selective learning framework identify important segments with high attribution strength but moderate consistency that indicate reflective rather than shallow reasoning. The framework then applies selective SFT 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]
INTRODUCTION
Recent Large Reasoning Models (LRMs) (Jaech al. 2024 Guo al., 2025 Yang et al. 2025
have demonstrated strong capabilities in solving complex problems. Their effectiveness is largely attributed test-time scaling by increasing computation during inference to produce extended chains of thought (CoT) (Wei al. 2022) that include detailed problem understanding, step-by-step solution processes, and comprehensive verification. These long-CoT trajectories have also become valuable supervisory resources for cold-start supervised finetuning (SFT) (Muennighoff et al. 2025).
However, current reasoning CoTs often span thousands of tokens, of which only small fraction meaningfully contributes to reaching the correct answer or improving confidence (Sui et al. 2025).
substantial portion consists of redundant repetition or incomplete truncated thoughts (Wang
et al. 2025d), illustrated in Fig. T (left). More critically, verbosity without substance actively degrade reasoning performance as CoT length increases (Wu et al.) 2025) Huang et al., 2025). This phenomenon also evident by the right-top panel of Fig. [ where incorrect CoT's are typically correlated with more segments and tokens than correct CoTS for the same queries. Training LLMs on verbose CoT supervision with sparse positive contribution further exacerbates these issues. Models learn to imitate redundant behaviors, waste learning capacity on trivial continuations, and fail to prioritize the crucial high-impact parts of reasoning sequences Lin et al. (2024). result, finetuned models tend to achieve limited accuracy gains and generate inefficient outputs.
Prior studies have explored various strategies identify important parts in long reasoning chains to construct compressed CoT supervision for efficiency purposes. However, they either focus on fine- grained token-level analysis (Xia et al. 2025b) that neglects semantic integrity and fails interpret redundancy in terms of meaningful reasoning units, or rely segment-level perplexity (Cui et al.) 2025a) or entropy (L et al. 2025b) calculations. These indirect metries provide not entirely consistent
"oode and data available at htw.lh/w....pppp./......ooo////n

---

Published as a conference paper at ICLR 2026
digit product postitin ropts locertio ( (oog (1995)a)s
Awe Awerage Sugnan Number
Correct Biegmentt
(Segren()
the last dighs of product polliive
the equalion -U1995)
Deg (o9ss] of Anlnda 1995 (2000-. 5/2 2000" 22 socohs 5
4,000000 - 20,00 00,000 25 belaaaalutting 1ewp.iiinn.000000 030 20,000
S p00coo0.oo lete Then 5880000 25 pes Vntesa 1995" is 3.93cee helun herefor.... productof pouriie roats 59900255
dight,
frres diits are cet, which
Perrage
doub ke-check saloviorians sure I didn't make any misiones Wireiest the or orginsl squatton.ubb99) x =0og__ (1905) 10 af took log bass 1995 of both dess. Aftributton Stre. 14%
8980.05, So yes, the last three dights are 025.
Segmeen Number
Taten Number
Truncated Verification
1982025 Weeaaiuuriing wytemmmaaan 2000 -100000 Anlunee (-5) 1-5)
..
2000
Corrulative Aftrisution Strengh
14
Segnart Strongth
Attributiin Stre.
0.
Repested
... ..llernallelyy kil me compule 1995 1995 us (2000 572,, which & 2000 2"2000"s
55, which is 4,000,000 20,000
indees 1990,025. Se that part in
0.4
dight
6%
Superficial clarification
0.4
digis.
Attribution Sire. 13%
Cornect CeT Encernect GOT
thrse 025 So the should Diraction Coneistency
0.2
All-ibuiion
in compate onother 997s.lnlate 1990,000-
VelnCompptt
1e9c000..nnrssbbee resutt toliihrarreeeeee Wassedjo25)
20
60 Segrrent Percentage Sortell)
20
100
Attribution Stre. 37%
Figure I: Left: An illustrative CoT with important (green blocks) and redundant segments (gray blocks). Our metries distinguish important from redundant segments (repetitions, truncations, super- ficial 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 Cumulative distribution function (CDF) of normalized
CoT's for the same queries. segment strength in correct and incorrect CoT', with segment ordered in descending strength.
signals of importance and are prone to both false positives and false negatives. False positives occur when methods overemphasize superficial scaffolding text (c.g.. "So, let's calculate step by step") that contributes little to actual reasoning yet serves as linguistic bridges whose removal disrupts subsequent textual coherence. False negatives arise when methods filter out independent verification or intermediate conclusions that, while exhibiting low-entropy and their removal not affecting linguistic fluency, significantly enhance the probability of reaching correct final answers. Consequently, existing methods cannot accurately and comprehensively distinguish truly important segments from various forms of redundant content that meaninglessly contribute to reasoning accuracy.
In this work, we systematically identify important segments that directly contribute to correct answer prediction within long CoT's, and show that unimportant segments cluster into distinctive redundant patterns. Specifically, we utilize integrated gradient (IG) attribution (Sundararajan et al.) (017)) to calculate each token's direct influence on improving correct answer prediction and aggregate token-level attributions at the segment level to obtain two metrics: (1) Attribution strength quantifies the overall magnitude of a segment's influence on the model's prediction by summing absolute IG values within each segment with length normalization.
(2) Attribution direction consistency is defined as the ratio between the absolute sum of signed IG attributions and the sum of absolute IG attributions, captures how uniformly a segment contributes in one direction (either positively or negatively toward the correct answer). Extremely high consistency often reflects shallow or skewed reasoning, such as segments with uniformly positive token IGs that merely provide superficial clarification (see the penultimate segments in Fig. Ieft), or uniformly negative token IGs that severely distort the reasoning process toward incorrect conclusions. In contrast, moderate consistency indicates more reflective reasoning that mixes supportive and corrective attribution within segment, which is more critical for problem solving.
We first verify significant redundancy in long CoTs using attribution strength. showing that 30~40% of segments accumulate over 80% of total attribution both correct and incorrect CoT's (Fig. T right-bottom). Building on this insight, we propose segment-level selective learning framework that learns the most critical parts of long CoTs leveraging attribution strength and direction consistency. It identifies segments with high strength but moderate consistency as important, filtering out redun- dancies like repeated content, truncated thoughts, and dispensable clarifications with minimal gains on correct answer probability. Unlike pruning- based SFT methods that compress CoT supervision but compromise accuracy, our framework applies selective SET (Lin et al.) 2024) that trains only on important segments while masking loss for unimportant ones. This acts as implicit regularization by
2

---

Published as a conference paper at ICLR 2026
preventing overfitting to uninformative content. Experiments across multiple models and datasets, using both self- gelf-genertte and reference long CoTs, show that our method outperforms full-CoT SFT by improving reasoning efficacy (up to 4.7%) while reducing output length (up to 18%). Notably, our important segment identification can be broadly applied to other contexts, such as emphasizing policy gradient updates on important content in reinforcement learning.
METHODOLOGY
2.1 PRELIMINARY
long-form CoT typically comprises multiple segments, each focusing on distinct aspects such as detailed problem understanding, intermediate reasoning exploring different solutions, or multiple verification process. Moreover, it inevitably contains redundant content, including repeated or truncated thought, and excessive clarification of self-evident facts, which diminishes the overall quality and efficiency of the COT. To facilitate segment-level importance analysis of long CoTs, we first partition each long CoT T into multiple segments T {S1. S2, .... Sn} using common transition keywords (e.g.. Wnit",,t",, "ttenntieraatvveel""" ) that naturally occur within reasoning traces, following (Lu al. 2025). The complete list of segmentation keywords is in the Appendix
Importance Definition We define segment within the full reasoning trace as important if it contributes to the final answer prediction, either by guiding transition from incorrect to correct answers by enhancing confidence of correct answers. An intuitive approach is to sequentially append segments and force answer generation when adding each new segment to compute the change in correct answer probability as its attribution. However, this method underestimates segments that contribute indirectly, such as problem understanding or exploration of incorrect alternatives, which may not immediately improve answer prediction but establish crucial foundations for subsequent reasoning steps. Leave-one-out methods that mask individual segments and measure their effect on the final prediction, suffer from the same limitation, as omitting these segments typically does not significantly alter final answer prediction when their subsequent content remains intact. Therefore, we adopt a more principled approach using integrated gradients (IG) (Sundararajan al., 2017) to measure each segment's s attribution to the final answer prediction.
2.2 INTEGRATED GRADIENTS BASED SEGMENT IMPORTANCE
IG measures input token attribution by computing partial derivatives of the model output with respect to each input feature along a straight- line interpolation path from uninformative baseline to the actual input embedding. By accumulating gradients across all interpolated points, IG estimates the total attribution of each input token to the final prediction. This approach captures both direct and indirect influences, making it particularly suitable for identifying the importance of reasoning segments that may contribute implicitly rather than immediately affecting the final answer prediction.
Formally, given a model F, an input token embedding x and its baseline value x' (typically the padding token embedding), the integrated gradient for the i-th input dimension is formulated as
(F + ax (x (e==x))
x²)) do
(G) (x: x!)x
Ox:
1
OF(x' + j/J (x x'))
( x
Ox.
j=1
where z; and 2 refer the i-th dimension of x and x'. respectively. The integral is approximated using J interpolation steps where j denotes the j-th step.
Segment-level Aggregation After computing IG attribution of each input token as IG(z) Σ, (z), we aggregate them at the segment level to obtain segment-wist importance scores. While IG values can be either positive or negative, where positive values indicate increased likelihood of predicting the correct answer and negative values indicate decrease, tokens with negative IG values should not be dismissed as unimportant. Negatively attributed tokens may represent incorrect but necessary exploratory reasoning that ultimately guides the model toward the correct solution. For
3

---

Published as a conference paper at ICLR 2026
example, initial segments often exhibit overall negative IG sums, but these segments are typically important for establishing problem understanding and initial exploration. Therefore, we utilize the absolute IG value of each token to capturing the magnitude of influence regardless of direction.
We aggregate the absolute IG values for all tokens within each segment and define two segment-level measures: (1) Attribution Strength: computed as the sum of absolute IG values within the segment with square root length normalization applied to prevent bias toward longer segments; (2) Attribution Direction Consistency: measures the extent which segment exhibits consistent positive or negative contributions versus internally conflicting mixed attributions.
Formally, given segment S ; ON) with N tokens, where each token On is associated with an IG value IG,),,, we define the segment-level attribution strength and direction consistency as:
Strength(S) Eones [[G().)] Consistency(S)
|Eones (G))
VN
Eo.es [IG(o,)]
To ensure strength comparability across all segments within a CoT, further normalize all strength scores across the set of segments (S1, , SM) Specifically, for each segment Smy the normal- ized attribution strength is defined
(trength'(S... Strength
(4)
Strength (S²)
2.3 SEGMENT-LEVEL SELECTIVE LEARNING
Important Segment Identification We utilize our segment-level metrics to distinguish critical segments from unimportant ones for more targeted learning of long CoTs Important segments should exhibit higher attribution strengths that occupy a substantial portion of total attribution across all segments, and moderate rather than extremely high attribution direction consistency. Extremely consistent attribution directions (most positive or most negative) may indicate shallow or skewed reasoning (e.g., dispensable explanation of final answers, and uninformative or severe wrong explorations, as analyzed in Appendix A) while moderate consistency often reflects reflective and critical reasoning patterns where the model explores different possibilities, refines its understanding, and eliminates incorrect content, which are more valuable for learning effective reasoning strategies. Specifically, given a CoT with segments {S1, S2. SM }. we first rank segments in descending order of their normalized attribution strength. Let ı be permutation such that:
Strength' Soaggth'(S,aa (S=(1)) 2 Strength' (5x)2) 2 ... 2 Strength' (S,())))
(5)
We then identify the minimal number of top-ranked segments k- whose cumulative attribution exceeds a predefined threshold r (e.g., 80%);
arg
min KEL.2....M
Strength (Sx(s)) 2 2
(6)
Finally, we define the important segment set as those among the top-k* segments whose attribution direction consistency is below a threshold /3 (e.g., 0.8), while the remaining as unimportant ones.
Simportant = { < k*, Consistency(Saoo) s 3}
Selective SFT Long-CoT trajectories serve as valuable supervisory resources for cold-start of reasoning models via SFT, which optimizes parameters 0 using the cross-entropy loss over all tokens.
1 Σ log P(o.loce:9:0))
LSFT(() =
where T denotes the length of the reasoning trajectory, o, is the predicted token position t, Oct = {01, 02, 02...., 0 represents the preceding token sequence, and is the input query.
To mitigate learning from meaningless parts, prior efficiency works mainly prune redundant portions of each complete CoT, yielding more concise supervision, but pruning typically degrades performance. Instead, we follow (Lin et al. 2024) and propose selectively train only tokens within important
4

---

Published as a conference paper at ICLR 2026
segments. This strategy ensures that parameter updates are driven solely by the most critical reasoning parts, while preserving coherence of the full trajectory. The selective SFT loss is formulated as:
Tieeeectvee SFT SFT(0) Σ, ((r)
(o.) log Pioeloce, R:0)
where ((rr) indicating whether token or belongs to segment Sm with S. E Simportant.
SEGMENT IMPORTANCE ANALYSIS
Before applying segment-level selective learning strategy, first validate that our importance metric effectively distinguishes meaningful reasoning parts from redundant behaviors and determine the optimal hyperparameters 7 and
Analysis Setup
We conduct an investigation on self-generated long reasoning trajectories from (Guo et al., 2025). From 32
the L.MO datasets (Ye al. 2025) using R R1-Distill-Qwen2.5-77 candidate samples per problem, we select the shortest correct and incorrect outputsm, split them into segments and compute their attribution strength and consistency. Segments
ranked as Eq. 5
with cumulative attribution averaged across percentage intervals, as illustrated in the top-right part of Fig. The analysis reveals that only 30%~~40% of segments contribute significantly (80%) to the final prediction, regardless of correctness, while most segments exhibit low attribution. This verifies substantial redundancy long CoTs and motivates our approach for identifying important segments.
3.1 IMPORTANT VERSUS UNIMPORTANT SEGMENTS
We aggregate important and unimportant segments, respectively from all correct reasoning outputs" and compare the distinct patterns exhibited by these two subsets. We first intuitively set the threshold 80%, 3 0.95 as Eq. for this comparative analysis.
1. Important segments yield greater improve- ment in correct answer predictions. We se- quentially append segments to the input and en- force answer generation after adding each new seg- ment to compute the change A in correct answer confidence relative to without that segment. The model's 's confidence on the correct answer is defined as the probability that it generates the correct an- swer. We empirically estimate this correct answer confidence by calculating the fraction of multiple cemperature-samlles generations (i.e., 32 samples) that produce the correct answer. As shown Fig. 2 segments with higher attribution strength (top 80% of total attribution), and moderate direction consis- tency achieve significantly larger gains in correct answer confidence compared to their counterparts. This suggests that segments with high strength but moderate consistency reflect the critical exploratory reasoning that contribute to accurate predictions.
Correct Answer Confidence
8-
6
4-
Top streghh Strongth
Law Consisteryy High Cornnneny
Important Unimp Unimportonn
Figure 2: Change (A) in correct answer confi-
dence across different segment types.
Important segments exhibit relatively lower perplexity and entropy. We calculate the log perplexity and entropy of all tokens in each long reasoning trace and aggregate them at the segment level to compare the linguistic properties between important and unimportant segments. As shown in Fig. 3 important segments consistently exhibit significantly lower log perplexity and entropy compared to unimportant segments. This indicates that our identified critical segments tend be more predictable and carry lower uncertainty, likely reflecting processes such as problem understanding, logical deduction or essential computations that follow more constrained patterns. In contrast, unimportant segments, such incomplete truncations associated with higher uncertainty (Wang et al. 2025d), as well as verbose explanations or dispensable elaborations that exhibit higher variability and greater linguistic freedom, resulting in higher perplexity and entropy.
We focus reasoning outputs, objective is leverage reliable reasoning for SFT.

---

Published as a conference paper at ICLR 2026
3.0
1.0
BLEU wur.t Preceding Segments
2.5
2.0
Log Perplexity
2.0
1.5
Entropy
0.6
1.5
1.0
0.4
0.0
0.0
Important
Unimportant
Important
Unimportant
Important
Unimportant
Figure Log perplexity, entropy and BLEU similarity of important versus unimportant segments.
Unimportant segments are more characterized by repetition or incomplete truncation. We further examine whether unimportant segments closely associated with redundant reasoning behaviors such as repetition and incomplete truncation. To this end, we calculate BLUE similarity of each segment against all preceding segments within the same CoT (rightmost panel, Fig 3) Unimportant segments exhibit overall higher BLEU similarity scores compared to important segments. In particular, unimportant segments contain substantially more highly repetitive content (e.g.. BLEU 0.8) that reiterate previously established reasoning. In addition, we prompt Qwen3-8B (Yang al. 2025) assess whether each segment incompletely truncated, meaning that subsequent segments failing to logically follow it (as implemented in Appendix [B. The results show that 49% of unimportant segments are classified truncated compared with only 26% of important segments. These results demonstrate that unimportant segments identified by our method are more strongly associated with repetition and truncation, contributing minimal additional value to the reasoning chain and can even introduce noise that obscures the logical flow. We additionally analyze the positional distribution of important versus unimportant segments within each CoT Appendix
3.2 HYPERPARAMETER SEARCH
To determine the optimal threshold r for identify- ing high-impact segments for targeted training, we adopt a greedy search strategy by selecting r that maximizes the difference in correct answer confi- dence change A between resulted important and unimportant segments while minimizing false neg- atives (i.e., ensuring that unimportant segments exhibit the lowest confidence gain). shown Fig. we set the optimal threshold as 0.7.
Unimportant
Impor ton--uummmprrton
Correct Answer Confidence (96)
10
8-
On average, this setting yields about 33% of seg- ments classified important, accounting for 45% of the tokens within each self-gemerated COT. This imbalance is because important segments tend be longer, whereas unimportant ones often either repeat part of the previous content are prema- turely truncated in the middle. We further ablate different values of under 0.7 and provide experimental comparison of varying 7 overall performance in Appendix D and identify 3 0.8 as the optimal consistency threshold.
2
0.55
0.65
0.75
0.85
0.95
Threshold
Figure Change in correct answer confidence (A) across different segment types under varying
т. The red star marks the selected threshold r".
4 EXPERIMENTS
4.1 SETUP
Training Details We compare our selective SFT on IG-based important segments against standard SFT full long CoT supervision, experimenting both cnerruction-oollovviii and reasoning baseline

---

Published as a conference paper at ICLR 2026
Table 1: Comparison results on different baseline models. "Overall" reports the average accuracy under greedy decoding, or pass @ 1 and pass @6 6 under temperature sampling, together with the average output token count over all benchmarks. The percentages parentheses indicate the relative accuracy improvement and token reduction of our method compared to full CoT SFT during generation. The best results are highlighted in bold.
In Domain MATH500 AMC23 AIME24 GPQA Minerva Olympiad Acc/Pass @1
Out of Domain
Overall Pass@6
Models
Length
Greedy Decoding
RI-Disiill Owen-1.5B
70.6
52.5
16.7
27.3
21.7
31.4
36.7
20244
Full CoT SFT
80.8
70.0
26.7
25.3
26.8
39.4
44,8
/
16520
Segment Selective SFT
82.4
60.0
40.0
28.8
27.6
42.8
46.9 (*4.7%)
/
13506 (118.2%%
Qwen-nQ
85.8
85.0
46.7
42,9
38.2
47.9
57.7
12518
Full CoT SFT
91.2
90.0
50.0
42.4
41.2
57.9
62.1
9693
Segment Selective SFT
95.2
90.0
56.7
40.4
46.3
58.7
64.5 ("3.9%)
8499 (2.3233%)
Qwen2.5-77 Instruct
77.0
50.0
10.0
38.9
34.9
37.5
41.4
1405
CoT SFT
77.2
55.0
23.3
29.8
40,9
44.2
10317
Segment Selective SFT
76.6
57.5
23,3
44,4
30.5
41.5
45.6 ("3.2%%
9852 (14.5%)
Temperature
Sampling
RI-Distill- Qwen-1.5BB
84.0
73.3
29.8
32.8
32.0
44.8
49.5
71.1
9791
Full CoT SFT 85.1
74.7
35.1 1
31.0
32.0
47.3
50,9
72.4
10043
Segment Selective SFT
85.1
75,9
36.1
32.8
31.9
48.1
51.7 (*1.6%)
73.2
(11))%
9388 (16.5%)
RI-Distill- Qwen-7h
92.8
90.5
45.5
64.5
79.2
Full CoT SFT
93.6
91.5
60.2
41.8
45.6
60.3
65.5
79.2
7934
Segment Selective SFT
93.3
91.9
61.0
42.3
45.8
60.3
65.8 ("0.5%) 80.0 (11.0%)
7709 (12.8%)
Qwen2.5-7B-Inss Inn--tsuttttt
75.7
51.9
12.0
36.9
35,4
37.2
41.5
60.2
910
Full CoT SFT
77.0
55.4
16.4
38.9
31.6
40.8
43.4
66.0
9902
Segment Seleetive SFT
77.1
58.5
16.4
43.7
32.4
41.9
45.0
('3.7%)
66.5
(10.8%)
9195 (17.1%)
RI-Distill-Qwen2.5-. (Guo et al. 2025 and Qwen2.5-7B --Bttunttt (Team] 2024). We use 817 high-quality questions from the LIMO mathematical dataset Ye et al. 2025) and evaluate both provided and self-generated CoT supervision. For self- generated supervision, we prompt R 1 1.-Distill-Qen2255-... Qwen2.5-7B to generate 32 candidate responses per question and select the shortest correct response for training (as described in Sec. 33 We utilize provided higher-quality CoT supervision for RI-Distill-Qwenn25--13. and self-generated CoT's for Qwen2.5-7B-Instruee We apply the same hyperparameter configurations of r = 0.7 and / = 0.8 to both sources of CoT supervision (as their hyperparameter search in Sec. 3.2 shows similar results). All models are fine-tuned with full parameters with maximum sequence length of 16384. Additional training details are provided in Appendix C.3 model performance under different combination of and / in Appendix D
models of different scales: RI
and EL--iisiil-Qwen2.5-1
We also provide
Evaluation Details We conduct evaluation on both in-domain and out-of-domain benchmarks follow-
ing the similar setup in (Ye et al. 2025). The in-domain benchmarks include Math500 (Hendryeks)
et al. 2021). AMC23 and AIME24 while
(Rein et al. 2024). Minerva (Lekkowyzz 2024) out-of-domain. We report accuracy using
al. 2022) and OlympiadBench (He al.
greedy decoding, and pass@1 1 and pass @ 6 metrics using temperature sampling (temperature=0.6, top-p=1.0) as no universally optimal decoding strategy exists across all models and benchmarks.
We employ a zero-shot CoT setting with a maximum response length of 32,768 tokens. For larger benchmarks (MATH500, GPQA-Diamond, Minerva, and OlympiadBench), we sample responses per instance, while for smaller benchmarks (AMC23 and AIME24), we sample 32 responses per instance. We compute the average token count across all sampled outputs as the response length.
4.2
MAIN RESULTS
Table repoott the overall comparison across six benchmarks. Relative to the baseline performance without SFT training, our segment-level selective SFT consistently outperforms full CoT SFT on both in-domain and out-of-domain datasets, across different baseline models and under both greedy decoding and temperature sampling. In addition, compared to full CoT SFT, our approach achieves a substantial reduction in token usage, leading to better optimization while simultaneously improving output efficiency. Notably, the gains in both accuracy and token reduction are more pronounced under greedy decoding. This is because the randomness in multiple temperature sampling tends to
7

---

Published as a conference paper at ICLR 2026
smooth out and partially offset the advantages of training. Nevertheless, even under such stochastic conditions, our approach still yields measurable improvements, underscoring its robust effectiveness.
4.3 ABLATION STUDY
To investigate the contributions of our important segment identification and selective SFT, we conduct with results averaged across datasets. Taking full CoT SFT as the baseline, we compare two categories Pruned CoT SFT according to our identified important segments and Selective SFT using various criteria: (1) randomly selecting roughly 33% of segments from each CoT; (2) selecting the top 45% of tokens ranked by absolute IG values; (3) Selecting the top 45% of tokens ranked by original IG values; (4) Selecting only high-strength segments; (5) selecting segment with high strength and moderate direction consistency (our method).
an ablation study using RI RI-Distill-Qwen-Ia.
These ratios are chosen to ensure the amount of selected content is comparable across methods.
As shown in Table 2 pruning unimportant segments degrades accuracy, whereas our selective learning can improve SFT performance while reducing token usage. Among selective SFT, our IG-based important segments outperforms random seg- ments and token-level selection based on ab- solute original IG values, with absolute IG values proving more effective than original IGs for identifying importance. Furthermore, the improvement over only high-strength segments shows that incorporating moderate direction con- sistency can further filter out low-impacted seg- ments to achieve greater token reduction without accuracy loss. Overall, these consistent improve- ments demonstrate that segment-leve granular- ity better preserves reasoning coherence than token-level selection, and our IG-based strength- consistency criterion effectively identifies the most contributive reasoning components.
Table Ablation study on RI--Ditill-Qwen--l.. Our important segments have high attribution strength and moderate direction consistency.
Greedy Acc.
Sampling
Methods
Length
Pass
Length
Full COT SET Pruned CoT SFT
44.8
16520
50.9
10043
43.9
15097
49.2
7766
Sclective SFT Random Segments
45.1
15138
50.8
10117
High-Abs IG Tokens
46.1
14612
50.8
9495
Hligh-Orig-ICC -IG Tokens
45.2
14747
50.3
9876
High Strength Segments
46.9
14715
51.4
9466
Our Important Segments
46.9
13506
51.7
9388
4.4 COMPARISON TO EXISTING SEGMENT IMPORTANCE MEASURES
To demonstrate the superiority of our IG-based importance segment identification, we apply Selective SFT using our identified segments and compare against alternative importance measures: (1) First- Correct Solution (Chen et al. 2024): retaining the first correct answer and all preceding segments as important: (2) Confidence-Gaii Segments (Xu al., 2025): identifying segments that directly improve the model's confidence in the correct answer compared to when the segment is removed (as described in Sec. B.Ip. (3) Segment Perplexity (Cui et al. 2025b): identifying critical segments whose removal leads to substantial increase in CoT perplexity. (4) Segment Entropy Li al. (2025)): prioritizing segments with higher aggregated token-level entropy as they indicate greater informational contribution. We set thresholds for methods (3) and (4) ensure the selected important content occupies approximately the same 45-550% token ratio as our method for a fair comparison.
Results averaged across all datasets R1-
Distill-Qwen-1.5p are shown in Table 3 Our method consistently outperforms others both accuracy and token reduction, demonstrating its ability to more precisely identify segments that are truly contributive to reasoning, while per- plexity and entropy are not entirely consistent signals of importance. The weak performance of Confidence-Gaii Segments stems from its ne-
Table Selective SFT performance on important segments using different importance measures on RI-Distill-Qwen-I 1.5B, Accuracy and token length averaged across all datasets.
Greedy Acc.
Sampling
Methods
Longth
Pass@) I Length
Full CoT SFT
44.8
16520
50.9
10043
glect of indirectly contributive segments, such as those supporting early-stage question under- standing eliminating incorrect solution paths, as evident in its identification of only 24% of segments average. Notably, our approach
First- Cortect Solution
46.2
15238
51.0
9492
Confidence- Segments
44.7
15479
50.3
9856
Segment Perplexity
44.7
14973
511.2
9816
Segment Entropy
44.5
16288
51.2
10148
IG-bused Important Segments
46.9
13506
51.7
9388
8

---

Published as a conference paper at ICLR 2026
Solution, suggesting that segments beyond the first correct answer, such as answer verification or alternative solutions, are also meaningful by improving correct answer confidence. Moreover, it achieves greater token reduction than Porst-Corrett Solution, indicating can more selectively mask unnecessary segments before the first correct answer.
surpasses
4.5 PRUNING-BASED COMPARISON
We further apply our IG-based importance segment identification to segment-level CoT pruning for SFT, and compare against other strategies introduced in Sec. 4.4 We do not compare with token-level pruning methods as they would significantly disrupt text coherence and fluency in long CoTs and severely degrades SFT performance. We adjust thresholds for our method and compared methods (3) and (4) to prune nearly 30% of tokens from unimportant segments, excessive pruning ratios substantially impact SFT performance. As shown in Table our method better maintains performance in both decoding settings, compared to other CoT pruning strategies. Among them, First-Correct Solution performs best because retains consecutive and complete reasoning processes, while other methods inevitably affect the information coherence of supervi- sions, leading to more severe performance drops and limited token reduction in long- CoT scenarios. However, performance still declines after First-Correct Solution prun- ing. Compared Table 3 these results fur- ther demonstrate the effectiveness of our proposed segment-level selective learning across various importance metrics, which can maintain or even improve performance while reducing token consumption.
Table 4: Performance of different segment-level pruning methods for CoT SFT on R R1-Distill-Qwan 1.5B.. Accu- racy and token length are averaged across all datasets.
Greedy
Sampling
Methods
Acc.
Length
Pass@
Length
Full CoT SFT
43.5
17140
50.7
10163
First-Correct Solution
45.2
11987
47.2
6846
Confidence-Gain Segments
40.3
11344
44.5
4180
Segment Perplexity
38.6
15892
47.9
8847
Segment Entropy
43.4
14321
45.9
7307
IG-based Important Segments
43.9
15097
49.2
7766
RELATED WORK
LLMs Reasoning Efficiency LRMs exhibit strong reasoning capabilities but typically generate verbose and redundant traces, leading to high computational costs and error accumulation (Sui et al.
2025) Wang al., 2025a d). Recent research explores various approaches improve reasoning efficiency, which can be grouped into four categories. First, prompt-based methods design input prompts to control task difficulty and token budgets (Han et al. 2025 Renze & Guven 2024). Second, decoding strategies dynamically reduce redundancy and shorten outputs during inference (Wang et al.)
2025c Liao et al. 2025). Third, latent reasoning approaches represent reasoning trajectories in latent spaces (Hao et al. 2024 Deng et al. 2024 Dong et al. 2025). Finally, training-based methods use
SFT on compressed CoT supervision (Xia et al. 2025a Lu et al. 2025) and reinforcement learning (RL) with conciseness rewards Yuaan et al. 2025 Lou et al. 2025) encouraged exploration (Xing
ct al. 2025). We focus on SFT-based methods, as they better balance efficiency and performance while serving as essential cold starts for RL-based training.
Importance Measurement To construct compressed CoT supervision for SFT-based methods, a key direction is measure importance of different parts within full CoTs and prune less important content. Token-level methods typically utilize token logits and entropy to prioritize salient tokens but
ignore semantic coherence (Xia et al., 2025a Cheng al. 2025 Wang al., 2025b) Segment-level
methods assess importance at segment granularity, better aligning with human reasoning units (Li
al. 2025a| Cui et al. 2025b). Moreover, Xu al. (2023) propose leave-one-out and greedy forward selection to estimate segment attribution. Recent Wang al. (2025e) suggests that I. LRMs rely both reasoning and memory retrieval mechanisms, and introduces perturbation-bant
attribution to promote cassoing--oominnn behavior. Unlike these indirect measures, we integrated mandients-basee attribution method to quantify the direct contribution of each segments to final answer predictions.
Selective Training Recent research challenges traditional full sequence learning paradigm that uniformly optimizes loss on all tokens (Lai et al. 2024 Lin et al. (022)). Lin et al. (2024) selectively apply loss only useful tokens and improve pretraining performance. Hans al. (2024) propose training on randomly selected token prevents memorization without performance degradation. Kim et al. (2025) groups tokens in each sample by importance and optimizes weighted loss that adap-

---

Published as a conference paper at ICLR 2026
tively emphasizes challenging groups. In this work, propose segment-lever selective learning framework that masks unimportant segments labeled by our importance measurement during SFT.
CONCLUSION
In this work, we proposed segment-level selective learning framework for more effective learning from long reasoning traces. By incorporating integrated gradient attribution, we introduced two segment-level metrics: attribution strength and attribution direction consistency to identify important segments with high strength but moderate consistency from full CoTs. Selective SFT on these segments improves both reasoning accuracy and reduces output length, outperforming full-CoT SFT. Our approach provides a principled way to identify critical reasoning segments and offers broader potential for emphasizing policy updates on targeted content in reinforcement learning.
REPRODUCIBILITY STATEMENT
To support reproducibility, provide detailed descriptions of our metrics and framework in Sec. and Appendix C.4 Additionally, we include comprehensive implementation details to reproduce our results in Sec. 4.1 and Appendix. C.3 covering hyperparameters selection, training and evaluation details, training frameworks. We will release our data and code after the anonymous review process.
ACKNOWLEDDGENNT
This research is supported in part by the Office of the Director of National Intelligence (ODNI), Intelligence Advanced Research Projects Activity (LARPA), via the HIATUS Program contract #2022- 22072200006, the Defense Advanced Research Projects Agency with award HR001 I 12220046, and NSF IIS 2048211. We would like to thank all the collaborators in USC INK research lab for their constructive feedback on the work.
REFERENCES
Xingyu Chen, Jiahao Xu, Tian Liang, Zhiwei He, Jianhui Pang, Dian Yu, Linfeng Song, Qiuzhi Liu, Mengfei Zhou, Zhuosheng Zhang, et al. Do not think that much for 2+ 3=? on the overthinking of ol-like 1lms. arXiv preprint arXiv:2412.21187. 2024.
Daixuan Cheng, Shaohan Huang, Xuckai Zhu, Bo Dai, Wayne Xin Zhao, Zhenliang Zhang, and Furu Wei. Reasoning with exploration: An entropy perspective. arXiv preprint arXiv:2506.14777 2025.
Yingqian Cui, Pengfei He, Jingying Zeng, Hui Liu, Xianfeng Tang, Zhenwei Dai, Yan Han, Chen Luo, Jing Huang, Zhen Li, Suhang Wang, Yue Xing, Jiliang Tang, and Qi He. Stepwise perplexity- guided refinement for efficient chain ofain-of- reasoning in large language models, 2025a. URL
inttps://arxian wrg. hatgg//bs//a 2502 13260
Yinggian Cui, Pengfei Hc, Jingying Zeng, Hui Liu, Xianfeng Tang, Zhenwei Dai, Yan Han, Chen Luo, Jing Huang. Zhen Li, et al. Stepwise perplexity-guide refinement for efficient Slain-of-thoughh reasoning in large language models. arXiv preprint arXiv:2502. 13260, 2025b.
Michael Han Daniel Han and Unsloth team. Unsloth, 2023. URL http::/ com www.ooohhhinnsas...
Yuntian Deng, Yejin Choi, and Stuart Shieber. From explicit to implicit cot: Learning to internalize cot step by step, 2024. [ttts:::: http://xxxxixrgaarra/a/a
14838
Shuai Dong, Siyuan Wang, Xingyu Liu, Chenglin Li. Haowen Hou, and Zhongyu Wei. Interleaved latent visual reasoning with selective perceptual modeling. arXiv preprint arXiv:2512.05665, 2025,
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseck-rl: Incentivizing reasoning capability in 1lms via reinforcement learning. arXiv preprint arXiv:2501 /225500...22..... 2025.
10

---

Published as a conference paper at ICLR 2026
Tingxu Han, Zhenting Wang. Chunrong Fang, Shiyu Zhao, Shiqing Ma, and Zhenyu Chen. Token- budget-aware IIm reasoning, 2025. [RLttttss://axiiii org/ abs/ 221224122.. 18547
Abhimanyu Hans, Yuxin Wen, Neel Jain, John Kirchenbauer, Hamid Kazemi, Prajwal Singhania, Siddharth Singh, Gowthami Somepalli, Jonas Geiping, Abhinav Bhatele, and Tom Goldstein. Be like a goldfish, don't memorize! mitigating memorization in generative 1lms, 2024. URL
bttpss://axxie http///xxxxvv...aa///////a 2406. 10209
Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu, Jason Weston, and Yuandong Tian. Training large language models to reason in continuous latent space, 2024. |RLhttps
//arxiv. aww////ba//aa 2tt.b///222..7....
Chaoqun He, Renjie Luo, Yuzhuo Bai, Shengding Hu, Zhen Leng Thai, Junhao Shen, Jinyi Hu, Xu Han, Yujie Huang, Yuxiang Zhang, et al. Olympiadbench: A challenging benchmark for promoting agi with olympiad-level bilingual multimodal scientific problems. arXiv preprint arXiv:2402.14008, 2024.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arrXiv:2105 038844, 2021.
Kerui Huang, Shuhan Liu, Xing Hu, Tongtong Xu, Lingfeng Bao, and Xin Xia. Reasoning efficiently through adaptive cfain-of-thought compression: A self-optimizing framework, 2025.
7/arxiv. org. abs http://bbb//50// 2209. 14093
Aaron Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El- El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Carney, al. Openai ol system card. arXiv preprint 1.Niv:22/2.1667000 2024.
Gyuhak Kim, Sumiran Singh Thakur, Su Min Park, Wei Wei, and Yujia Bao. Sit-go: Supervised fine-tuning with group optimization for large language models. arXiv preprint arXiv:2506. 15021, 2025.
Xin Lai, Zhuotao Tian, Yukang Chen, Senqiao Yang, Xiangru Peng, and Jiaya Jia. Step-dpo: Step- wise preference optimization for long-chain reasoning of llms, 2024. URL https://axxiv org/abs/2405. 18629
Aitor Lewkowyez, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Ra- masesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, et al. Solving quantitative reasoning problems with language models. Advances in neural information processing systems, 35:3843-3857, 2022.
Zeju Li, Jianyuan Zhong, Ziyang Zheng, Xiangyu Wen, Zhijian Xu, Yingying Cheng, Fan Zhang, and Qiang Xu. Compressing chain-of-thoughh in Ilms via step entropy. arXiv preprint NXXv:25088.0346 2025a.
Zeju Li, Jianyuan Zhong, Ziyang Zheng, Xiangyu Wen, Zhijian Xu, Yingying Cheng, Fan Zhang, and Qiang Xu. Compressing chain-of -hain-of-thought in Ilms via step entropy, 2025b. URL https
//axxiv.. httg//bb//500.. 03346
Baohao Liao, Yuhui Xu, Hanze Dong, Junnan Li, Christof Monz, Silvio Savarese, Doyen Sahoo, and Caiming Xiong. Reward-guided speculative decoding for efficient Ilm reasoning, 2025, URL
htttps://arxinn org/ httgg/bs////oo// /os.//5111... 19324
Zhenghao Lin, Zhibin Gou, Yeyun Gong, Xiao Liu, Yelong Shen, Ruochen Xu, Chen Lin, Yujiu Yang, Jian Jiao, Nan Duan, et al. Rho-1: Not all tokens are what you need. arXiv preprint 2xXi::2404.07765, 2024.
Zicheng Lin, Tian Liang, Jiahao Xu, Qiuzhi Lin, Xing Wang, Ruilin Luo, Chufan Shi, Siheng Li, Yujiu Yang, and Zhaopeng Tu. Critical tokens matter: Token-level contrastive estimation enhances lim's reasoning capability, 2025. URL https: /ttp:////vvvvvv... http:///wwwii..o///////111.1111.
19943
11

---

Published as a conference paper at ICLR 2026
Chenwei Lou, Zewei Sun, Xinnian Liang, Meng Qu, Wei Shen, Wenqi Wang, Yuntao Li, Qingping Yang. and Shuangzhi Wu. Adacot: Pareto-optimal adaptive chain-of-thouugh triggering via reinforcement learning, 2025. DRLttt:s:://ax:
httpa/awwww...aaa///aaa///////// 11896
Ximing Lu, Seungju Han, David Acuna, Hyunwoo Kim, Jaehun Jung. Shrimai Prabhumoye, Niklas Muennighoff, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro, et al. Retro-search: Exploring untaken paths for deeper and efficient reasoning. arXiv preprint arXiv:250.00338, 2025.
Niklas Muennighoff, Zitong Yang, Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hannanch Hajishirzi, Luke Zettlemoyer, Percy Liang, Emmanuel Candès, and Tatsunori Hashimoto. s1: Simple test-time scaling, arXiv preprint arXiv:2501. 193933, 2025.
David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel Bowman. Gpqa: graduate-level google-prost proof q&a benchmark. In First Conference Language Modeling, 2024.
Matthew Renze and Erhan Guven. The benefits of concise chain of thought problem-solving in large language models. In 2024 2nd International Conference on Foundation and Large Language (ttp:///mmmm222222....20 URL Utt.tttp///ax......c//00.11///// 11w.001111000111112222 TTT.LLMM622............ 2024. 2
Models (FLLM), pp. 476 483. IEEE, November 2024. doi: 10. 10.1109/Allm612//
Yang Sui, Yu-Neng Chuang, Guanchu Wang, Jiamu Zhang. Tianyi Zhang, Jiayi Yuan, Hongyi Liu, Andrew Wen, Shaochen Zhong, Hanjie Chen, et al. Stop overthinking: A survey on efficient reasoning for large language models. arXiv preprint arXiv:2503.16419, 2025.
Mukund Sundararajan, Ankur Taly, and Qiqi Yan. Axiomatic attribution for deep networks. In International conference on machine learning, pp. 3319-3328, PMLR, 2017.
Qwen Team. Qwen2 technical report. arXiv preprint arXiv:2407.10671, 2024.
Rui Wang, Hongru Wang, Boyang Xue, Jianhui Pang. Shudong Liu, Yi Chen, Jiahao Qiu, Derek Fai Wong, Heng Ji, and Kam-Fai Wong. Harnessing the reasoning economy: A survey of efficient rea- soning for large language models, 2025a. UtLpptss://axxx org/ 2ttp//550..337777
Shenzhi Wang, Yu, Chang Gao, Chujie Zheng, Shixuan Liu, Rui Lu, Kai Dang, Xionghui Chen, Jianxin Yang, Zhenru Zhang, et al. Beyond the 80/20 rule: High-entropy minority tokens drive effective reinforcement learning for I1m reasoning. arXiv preprint vxXiv:5566.09933 2025b.
Siyuan Wang. Enda Zhao, and Xiang Ren. Stepwise informativeness search for improving Ilm reasoning. In Proceedings of the 2025 Conference Empirical Methods in Natural Language Processing, pp. 25291-25300, 2025c.
Yue Wang, Qiuzhi Liu, Jiahao Xu, Tian Liang, Xingyu Chen, Zhiwei He, Linfeng Song, Dian Yu, Juntao Li, Zhuosheng Zhang, al. Thoughts are all over the place: On the underthinking of ol-like lms. arXiv preprint arXiv:2501.18 18585, 2025d.
Yuhui Wang, Changjiang Li, Guangke Chen, Jiacheng Liang, and Ting Wang. Reasoning or retrieval? study of answer attribution on large reasoning models. arXiv preprint 2rXiv:2509.24155 2025e.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of- thougtt prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824 24837, 2022.
Yuyang Wu, Yifei Wang, Ziyu Ye, Tianqi Du, Stefanie Jegelka, and Yisen Wang. When more is less: Understanding chain-of-though length in Ilms, 2025,
Heming Xia, Chak Tou Leong Wenjie Wang, Yonggi Li, and Wenjie Li. Tokenskip: Controllable of-innof-thought compression in llms. arXiv preprint 25X2 12067, 2025a.
Heming Xia, Chak Tou Leong, Wenjie Wang, Yongqi Li, and Wenjie Li. Tokenskip: Control- lable chain-of- tf-tough compression in Ilms, 2025b. URL https: ///rriiv..... wwt//bo/////.eeeoeeo...oeee abb/2502 12067
12

---

Published as a conference paper at ICLR 2026
Shangyu Xing, Siyuan Wang. Chenyuan Yang, Xinyu Dai, and Xiang Ren. Lookahead tree-base baseb rollouts for enhanced trajectory-level exploration reinforcement learning with verifiable rewards. arXiv preprint arXhv:2510.2430,, 2025.
Zhentao Xu, Fengyi Li, Albert Chen, and Xiaofeng Wang. Procut: Llm prompt compression via attribution estimation, 2025. UtLhsttt::/// wttpv.aaivvvv..///.. org/ abs abs/2508. 02053
An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint 00982250..09388. 2025.
Yixin Ye, Zhen Huang, Yang Xiao, Ethan Chern, Shijie Xia, and Pengfei Liu. Limo: Less is more for reasoning. arXiv preprint arXiv:2502.03387, 2025.
Danlong Yuan, Tian Xie. Shaohan Huang, Zhuocheng Gong, Huishuai Zhang, Chong Luo, Furu Wei. and Dongyan Zhao. Efficient training for reasoning models via length-aware optimization, 2025. httpstts::/ //tr//axivvvv/ org/ abs/ /205..255.... 12284
13

---

Published as a conference paper at ICLR 2026
A POSITIONAL DISTRIBUTIONAL OF SEGMENTS IMPORTANCE
We analyze the positional distribution of segment importance within long CoT's by identifying the decision segment, defined the first segment that derives the correct answer. We observe 40% of important segments appear after the decision segment, indicating that additional verification or exploration beyond the initial correct answer can still play crucial role. In contrast, 57% of unimportant segments occur before the decision segment, suggesting that the search for the correct answer often involves redundant behaviors such as repetition and truncation. This also highlights that simply splitting by decision segments may introduce a considerable number of false positives and false negatives (Wang al., 2025d). finer-grained analysis further reveals that most low- strength unimportant segments (64%) occur before the decision segment. In contrast, the majority of high-consistency unimportant segments (72%) appear after it. Combined with Figure 3.1 shows that high-consistency segments typically induce minimal improvement in correct-answer confidence, this suggests that once the correct answer is found, such segments are prone to shallow or redundant elaboration formalized verification or reinforcement of the already-four answer rather than introducing deeper critical examination.
B INCOMPLETELY TRUNCATED SEGMENT IDENTIFICATION
To determine whether a given segment is complete or truncated by alternative content, we employ Qwen3-8B with carefully designed prompt template. The template, shown in Table [ describes the task and includes few-shot examples. instructs the model to first reason step by step in thinking mode, and then output final judgment whether Segment 2, given the surrounding context of Segment and Segment 3, is a truncated segment. We frame this as a 0/1 classification task rather than using Yes/No answers, since we find that restricting the output to binary digits yields more reliable performance.
For each segment, we conduct up two inference rounds. In round the model runs in thinking mode with a token budget of 2000, temperature 0.2, top 0.7, and repetition Beaallteeeeaaaaeeeee 1.1. If the model fails reach to a conclusion in this round, we will inject an Early-Stopping Prompt (as shown in Ta. and run the model in round 2. with temperature 0.1, top_p 0.7, max tokens 1000, and repetition__pentt penalty 1.1. The input for round consists of the original input and the Round output, followed by the injected Early-Stopping Prompt.
C SUPPLEMENTARY IMPLEMENTATION DETAILS
C.I LLM USAGE
Our use of LLMs in this paper comprises two main components: writing assistance and tools for data analysis and construction. During writing, we utilize ChatGPT 5 and Claude Sonnet 4 to help polish the exposition. For data analysis and construction, as described in Section 4. we employ RI-Distill-Qwen225-72 to generate its CoT trajectories leading to correct answers as supervision for SFT, and to compute integrated gradient attributions on each token.
C.2 SEGMENTATION KEYWORDS
To partition each long COT into individual segments, select transition keywords following LLu et al. 2025) which includes "But", "Wait", "Alternatively", "However", "Hmm", "Hmmm"", "Not
sure". "Going back". "Backtrack", "Trace back". and "Another". Specifically, most of our keywords are directly adopted from (Lu et al.) (025). including "nn", "IuunAlteraatively"" "unnHowwverr" "'n'nNot sure". "nn'nGooig back". "annnBacktrack.. "unlnTrace back", "nn'nnnothhr".. We prepend "Anin" to each keyword because many LRMs naturally structure their reasoning output with paragraph breaks, making this prefix reliable delimiter for identifying reasoning segments. Our keywords differ from Lu al. (2025) in two ways. First, instead of the broad keyword "But", we use more specific variants including Bntnnut wait", "'n'nBut alternatively" and "'nlnBut just to". This change is motivated by empirical observations that LRMs, especially smaller distilled models like RI-Distill- Qwen2.5-7B 7B tend to use "But" as conversational fillers rather than a meaningful reasoning transition. Second, we exclude fillers such "Hmm" and "Hmmm", which also find be frequently
14

---

Published as a conference paper at ICLR 2026
Prompt for Judging Truncated Segments
I will provide you with segments of text which are generated continuously by reasoning model. Each segment itself represents complete thinking step, and the segments together represent continuous thinking steps during reasoning process.
Your task:
Taking Segment I and Segment 3 as context, determine whether Segment 2 is "truncated" segment, ie., unfinished thinking step.
Example a non-truncated segment:
> But let me double-check my calculations to be sure didn't make any mistakes, AnknFirst, the original cquation:Innn, (1995) x^[log_ 1998__99950000 x) xmll took log hase 1995 of both sides, correctly applied the logarithm rules, and ended up with quadratie in y =log__19000 1995 Solving that quadratic gave me two solutions, which translated back to gave me two roots. Multiplying them gave me 1995?, which 3,980,025. So. yes, the last three digits are 025.
The above segment shows complete thought process, because: 1. Internally, this segment has clear logic, where calculations and reasoning are derived progressively. 2. It reaches a conclusion, which the final answer to this reasoning step. 3. Externally, if given context, this segment may continue or refute it's previous segment, and next segment may also continue reasoning based on this segment, which forms a coherent reasoning chain.
Example of truncated segment:
Wait, hold on second. Let me verify if 1995² is indeed 3,980.025 3,980.025.unnnaaluuaaae WwhnCalcullainng 1995 1995/n1n1 compute 2000 * 2000 4.000.000.mmnnTenn -10,000..nnnAmd (-5) =
2000 (-5) -1,0000000...iimmme.. (-5) 2000
The above segment is a truncated segment, because: 1. Internally, this segment lacks clear logic, and it only shows partial calculations without reaching the final conclusion of 1995?, 2. Externally, given more context, this segment may abandon it's previous segment, and the next segment may also abandon this segment, which breaks the reasoming chain.
Now, based on the definitions and examples above, please judge whether Segment is a truncated segment. Here are the given 3 segments:
Segment I:
(SEGMENT 1
Segment 2:
> (SEGMENT
Segment 3:
> (SEGMENT 3}
Please reason step step, and answer "1" or "0" in the following format:
My final answer is:lnnnSS 0)
NOTE:
'The final answer must he either "1" "O", which means yes no that Segment 2 is a truncated segment.
TTee final answer must be \boxed format.
You must not explain anything more after giving the final answer.
Table 5: The prompt for judging truncated segments using Qwen3-8B.
Early-Stopping Prompt
Inin Considering the limited time by the user, have to immediately stop reasoning and give the answer (1 or 0) directly now www.in</thinkkAnnn My final answer is:Ann
Table 6: Early-stopping prompt for terminating the model's thinking process.
15

---

Published as a conference paper at ICLR 2026
used verbal habits instead of genuine segment boundaries. These design choices aim to balance generalizability across LRMs while ensuring that segmented reasoning units are meaningful rather than overly fragmented.
To illustrate the effectiveness of our method using different keywords, we also experiment with using the exact keyword list from (Lu et al. 2025). We apply these keywords to segment the original .5B. This results in much finer-grained segmentation, increasing the average number of segments from 22.5 to 33.5 per reasoning trace. The selective SFT results in Table show that this finer segmentation leads to higher accuracy and more concise outputs. We attribute this to that LIMO traces are generated by nigher-quality models (e.g. DeepSeek-F - QwQ-32B), whose reasoning contains fewer meaningless fillers. In such cases, finer segmentation yields more precise identification of important reasoning units.
LIMO CoT's and retrain R R1-Distill-Qwen22.
Table Results on R1 R1-Distill-Qwen 1.5B using new transition keywords from (Lu al. 2025
Greedy Acc.
Sampling
Methods
Length
Pass@1
Length
R1-Distill- H.-Diitill-Qeen-I. Qwen-1.5B
36.7
20244
49.5
9791
Full CoT SFT
44.8
16520
50.9
10043
Segment Selective SFT (Original Keywords)
46.9
13506
51.7
9388
Segment Selective SFT (New Keywords)
51.5
12646
51.5
9167
C.3 TRAINING & EVALUATION DETAILS
All models are trained using the Unsloth training framework (Daniel Han team 2023). Specifically, we train R1-Dit 21-Distill-Qeen2.5- 1.555 and R -.-Distill-wwen2--77 for 10 epochs respectively with learning rate of 3c-5 and 1.5e-5. For Been2..5-B-Instruue which does not original support long-chain reasoning, adopt two-stage SFT strategy that first conducts full CoT training for 7 epochs with a learning rate of 1 1.5e-5, and continued training for 4 additional epochs using only the identified important segments, with learning rate of 8e-6. For fair comparison, the full CoT SFT baseline is trained with the same two-stage schedule. All experiments employ a cosine learning rate schedule. The overall performance (including greedy accuracy, sampling pass @ 1, pass 6 and output length) reported in this paper is calculated using macro averaging across all datasets, treating each dataset as equally important.
C.4 METHODOLOGY DETAILS
For calculating token-level integrated gradients, we use J = 50 interpolation steps to estimate the integral in the IG computation for the cost-precision trade-off as indicated by (Sundararajan et al. 2017) that 20~3300 integration steps typically approximate the path integral within about 5% approximation error. We specifically utilize R1-Diisill-Qwen22.5-- to calculate token IG values on both reference long CoT's provided by the LIMO dataset and self-generated long CoT supervisions by Distill- Seeill-Qwee2.5-7BB Each costs roughly 7 for full SFT run of the same 7B model. In deploying our method and all baselines, we additionally include the first and last segments which are responsible for establishing the problem understanding and for explicitly formatting the final answer, together with the identified important segments as grusing-based based settings. For segment-level aggregation, we have explored several strategies: (1) nength--ormalizze summation, (2) direct summation, and (3) averaging the top 20% of token IG values. We find that the first two achieve comparable performance, while length- length-normaln normalized summation better mitigates bias toward longer segments.
compared to the 8 (GPUxhours) required
learning objectives in both selective SFT and pruning-
D DIFFERENT THRESHOLD HYPERPARAMETERS
We further experiment with different selection of the hyperparameters T and beta on RI-Distill- Qwen2.5-1.5B. We first explore different values of 3 E (0.7. 0.8, 0.9) under fixed 7 0.7 and report the temperature sampling results in the left panel of Figure 5 We observe that B = 0.8 achieves the best performance, which is also adopted the main experimental setting in our paper. Based on 3 = 0.8, we further validate whether T 0.7 indeed achieves the best experimental performance as expected in our hyperparameter search. We vary E (0.6, 0.7, 0.8, 0.9) as shown in the right panel
16

---

Published as a conference paper at ICLR 2026
of Figure 5 Results show that increases, segments selected important among all candidates. Consequently, the token reduction after training becomes smaller. However, performance decreases because higher 7 introduces more false positives, and denser loss mask negatively impacts training. Overall, r = 0.7 and 3 = 0.8 constitute our final experimental setting.
52.0
10500
52.0
10500
Psss
Pass@1
Output Length
Output Length
51.5
51.5
-10000
10000
51.0
51.0
9500
9500
50.5
50.5
50.0
9000
50.0
9000
0.7
0.9
0.6
0.7
0.8
0.9
Threshold B
Threshold
Figure 5: Model performance under different hyperparameters and 8.
E
ADDITIONAL RESULTS ON LLAMA3. 1---MA3.1-BB-ISSTRUI
To further demonstrate the effectiveness of our method beyond the Qwen model family, addition- ally train 1-t.aa..-----Intttte using the LIMO dataset. As shown in Table g Segment Selective SFT still achieves performance improvement, demonstrating that our method generalizes well to different model family.
Table 8: Comparison results on LLaMA3.1-SB-Instru
Greedy Acc. Length
Sampling
Methods
Pass 1 Length
LLaMA3. 1-8B-Inntruce
24.0 9691
23.5
4161
Full CoT SFT
30.4
17474
33.0
15067
Segment Seleetive SFT
33.8
16005
33.5
14669
17


Reviews:
[{"review_id": "aG1liinOvB", "reviewer": "Reviewer_Jvus", "summary": "The paper proposes a segment-level selection method for CoT reasoning that identifies and trains only those segments that meaningfully contribute to correct answer generation (for SFT). Using Integrated Gradients (IG), the authors measure each segment’s attribution to the final answer and summarise this using two metrics: Attribution Strength and Attribution Direction Consistency. An interesting design choice is that Attribution Strength is defined using the absolute IG values (rather than the signed values), and that the method selects segments with moderate, rather than high, Direction Consistency. The stated motivation is that (i) even reasoning that initially explores incorrect directions may still be useful later, and therefore should not be discarded purely because it has negative IG, and (ii) segments with very high consistency may reflect shallow or overly biased reasoning rather than careful reflection. Across multiple models (R1-Distill-Qwen2.5-1.5B/7B, Qwen2.5-7B-Instruct) and datasets (LIMO, MATH500, AIME, etc.), the proposed approach improves accuracy while reducing output length.", "questions": "- How were the segmentation keywords in Appendix C.2 selected? The transition keywords used here appear to differ from the list in Lu et al. (2025). Have the authors evaluated the effect of varying this keyword list on the final results?\n\n- In Table 2, could the authors provide an additional comparison for the “High Strength Segments” condition where segments with very large negative IG attribution are excluded? This would help clarify whether strongly negative-attribution segments are genuinely useful or whether they can be safely filtered out.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- It is interesting to see the use of IG at the segment level to score the importance of different parts of the CoT. The definitions of Attribution Strength and Attribution Direction Consistency are thoughtful in that they are not purely driven by raw gradient magnitude, but instead attempt to capture properties specific to reasoning structure.\n\n- Through extensive experiments, the authors present hyperparameter search, ablations, and comparisons with alternative importance measures.\n\n- The writing is clear and the figures effectively illustrate the ideas and motivation.", "weaknesses": "- The argument for using attribution strength based on the absolute value of IG may need to be strengthened. The paper motivates this choice by noting that exploratory reasoning can sometimes have negative IG values, and that such reasoning should not be thrown away. However, consider a segment that overall shows strongly negative IG values and only moderate consistency. Is it necessarily the case that such a segment corresponds to “necessary exploratory reasoning” (lines 155–156)? While useful exploratory reasoning may indeed yield negative IG values, it is less clear that the reverse implication always holds. In other words, segments with strong negative attribution are not guaranteed to be beneficial; they may simply reflect confident but unhelpful reasoning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sdOcuyQmKl", "reviewer": "Reviewer_2YhS", "summary": "They introduces a method to improve learning from CoT reasoning sequences in LRMs. It identifies which segments of a reasoning trace contribute to the correct answer using segment-level attributions derived from Integrated Gradients. Segments are evaluated for attribution strength and direction consistency, and only important segments are used for supervised fine tuning.\n\nThe key contribution is demonstrating that models can learn more effectively by focusing on the reasoning steps that truly matter.", "questions": "1. you mention that: \"We estimate the correct answer confidence by multiple temperature-sampled generations\". how do you calculate confidence?\n2. Why did you choose to work in segment level? why not token level?\n3. you mention that: \"Extremely consistent attribution directions (all positive or all negative) may indicate shallow reasoning (e.g., dispensable explanation of final answers) while moderate consistency often reflect reflective and critical reasoning patterns where the model explores different possibilities, refines its understanding, and eliminates incorrect content, which are more valuable for learning effective reasoning strategies.\". How do you claim this?\n4. Could you please provide a quantification of the computational cost of the IG attribution step? Specifically, relative to the total time required for a standard Full CoT SFT epoch, how much overhead does the segment attribution calculation add? \n5. The IG values were calculated using R1-Distill-Qwen2.5-7B. Did you explore how effective the segments identified by this 7B model are when fine-tuning much larger or smaller models, or models from entirely different families (i.e., treating the attribution model as an external tool)? This would help demonstrate the robustness of the identified importance definitions beyond a single model lineage.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "good analytical experiments\nreliance on gradient based method rather LLM as judge\nThe use of Attribution Strength combined with Attribution Direction Consistency provides a mechanism to distinguish critical reasoning from superficial or shallow content\nefficiency gains in terms of reasoning length (while having high/comparable accuracy)\nComprehensive ablation studies", "weaknesses": "- Focus on only one LLM (Qwen)\n- High computational cost for importance calculation (J=50)\n- Sensitivity to hyperparameter selection: the performance of the framework hinges on two crucial, empirically determined thresholds: τ (strength threshold) and β (consistency threshold). The hyperparameter search indicates that model performance is sensitive to these values, as choosing a higher τ introduced more false positives and negatively impacted training performance. The selection process (maximizing the difference in confidence gain between important/unimportant segments) is justified, but sensitivity remains a potential concern when moving to drastically different models or tasks.\n- The segment partitioning relies on a pre-defined set of common transition keywords (e.g., “\\n\\nWait”, “\\n\\nAlternatively”)\n- best performance is not indicated in the tables", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RD2kOt9FwQ", "reviewer": "Reviewer_KP2C", "summary": "This paper introduces a segment-level attribution framework for selectively learning from long chain-of-thought (CoT) traces produced by large reasoning models. The approach leverages integrated gradients to assign an attribution “strength” and “direction consistency” to each CoT segment, identifying those segments with substantial and mixed influence on answer prediction. Selective supervised fine-tuning is then applied on only the important segments while masking loss for uninformative parts. \n\nExperiments on multiple models and mathematics/reasoning datasets show improved accuracy and output efficiency versus standard full-CoT fine-tuning. The paper is supported by careful empirical and ablation studies and provides a detailed analysis of redundancy in CoTs.", "questions": "1. Is the method for selecting keywords reasonable? How are the keywords chosen? Does this method apply to various models of different types?\n\n2. Hyperparameter search in Sec. 3.2: The hyperparameter search method is conducted only for a single model and dataset. If the model or dataset is changed, is it necessary to re-search the hyperparameters? Given that the hyperparameter search requires first using the model to calculate the accuracy gain on the tasks, does this result in significant cost?\n\n3. Applicability to other domains: Is this method still reasonable in other domains, such as code or general domains?\n\n4. Failure case analysis: Have you reviewed any cases where your segmentation choices led to a performance drop in the model? Could such failures encourage further improvements to the masking strategy’s features?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. Evaluates both pruning-based and selective learning variants, tests random and prior-importance alternatives, and shows that selective learning is preferable to simple pruning, with the IG-based method outperforming confidence, entropy, and perplexity-based segment selection (Table 3, Page 8).\n\n2. The use of normalized strength and moderate consistency to select segments is justified theoretically and supported by empirical analysis (Section 3.1, Figures 2 and 3; the ablation study in Table 2, Page 8, directly supports the effectiveness of the joint criterion over alternatives).\n\n3. Figures are used to substantiate both the intuition and the rigor (e.g., Figure 1 summarizes the distribution and accumulation of attribution across segments, while Figure 3 provides boxplots linking log perplexity, entropy, and BLEU similarity between important/unimportant segments).\n\n4. Proposes segment-level attribution metrics based on integrated gradients—specifically, attribution strength and direction consistency—combined in a well-motivated and interpretable manner. The mathematical formulations (Equations on Pages 3–4) are clearly specified and grounded in attribution theory.", "weaknesses": "> Completeness of the chosen keywords for segmentation\n\nAre the keywords used for segmentation comprehensive? How were the keywords selected? The generalizability of this paragraph segmentation method across different models (e.g., DeepSeek-R1 and Qwen2.5-7B) requires further discussion.\n\n> Limited theoretical analysis of attribution methods\n\nAlthough the attribution metrics are intuitive and based on empirical motivations, there is insufficient discussion of their theoretical limitations or the potential for attribution leakage. For instance, integrated gradients may produce artifacts for redundant input or model features, which could falsely elevate or lower paragraph scores. There is no thorough analysis or mitigation (e.g., using SmoothGrad) to reduce variance or noise in the attribution process.\n\n> Threshold sensitivity and generalization\n\nThe critical thresholds for attribution strength and consistency (i.e., $\\tau$ and $\\beta$) are set via greedy search. However, the analysis of their stability or transferability to other datasets or inference types beyond mathematical QA is limited. Figure 5 (on page 15) addresses threshold effects, but only for one model/dataset, and does not provide a detailed analysis of generalization robustness.\n\n> Insufficient evaluation metrics for diversity sampling.\n\nIn the results presented in Table 1, the pass@k value for temperature sampling is not provided, which is insufficient to fully capture the performance of the method across different benchmarks.\n\n> Limited performance improvement\n\nAs shown in Table 1, for certain benchmarks, the accuracy after segment-selective SFT training is either lower or unchanged compared to the full CoT.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RTlVTrmsSK", "reviewer": "Reviewer_sHrL", "summary": "This paper study on the segment selection in long chains of thought (CoT), they propose a series of new metrics: attribution strength and direction consistency, from integrated gradient attribution perspective. Based on these two metrics, this paper proposes a segment-level selective learning framework to identify important segments with high attribution strength but reflective 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 the approach improves accuracy and output efficiency, enabling more effective learning from long reasoning traces.", "questions": "1. What is the definition of confidence in Fig. 2 and Fig. 4.\n2. In equation (2), how exactly is IG be calculated from hidden state to outcomes, i.e. the partial derivative, how to define J in practice. If it includes frequently altering input to acquire output score, computation cost should be discussed.\n3. Most of the experiment is conducted under Qwen-family framework, which leave the possibility that this method benefits from the model's characteristics cannot be completely ruled out. Experiment under different model architecture is suggested.\n4. Is there more application scenarios could be employed by those IG metrics.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Compared to prior studies which either neglects semantic integrity and fails to interpret redundancy or ignore the indirect effect to the answer, this paper integrated gradient attribution to analyse segment's effect in depth.", "weaknesses": "The mechanism in section 2.2 \"INTEGRATED GRADIENTS BASED SEGMENT IMPORTANCE\" should be discussed more, see questions part for details.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)