Learning from Partial Chain-of-Thought via Truncated-Reasoning Self-Distillation

Gianluigi Silvestri; Edoardo Cetin

cs.LGFeb 27, 2026

Learning from Partial Chain-of-Thought via Truncated-Reasoning Self-Distillation

Authors: Gianluigi Silvestri, Edoardo Cetin

Organizations: Radboud University · Sakana AI

Abstract

Reasoning-oriented language models achieve strong performance by generating long chain-of-thought traces at inference time. However, this capability comes with substantial and often excessive computational cost, which can materialize in redundant or inefficient reasoning. We study this setting and introduce Truncated-Reasoning Self-Distillation (TRSD), a lightweight post-training procedure that encourages models to produce correct predictions from partial reasoning traces. In TRSD, a frozen teacher model first generates a full reasoning trace and evaluates the corresponding answer distribution conditioned on the prompt and the complete reasoning to construct a synthetic training target. A student model with the same architecture is then trained to match the teacher's answer distribution while being conditioned only on a truncated prefix of its reasoning trace. Across multiple reasoning benchmarks and token budgets, we demonstrate that TRSD improves robustness to truncated inference, with far reduced accuracy tradeoffs when applied to a diverse set of reasoning models. Moreover, although never explicitly regularized for shorter generation during training, we also find that TRSD-trained models inherently output shorter reasoning traces without truncation, significantly reducing inference-time costs even without artificial interventions.

Figures & tables

Figure 1: Accuracy as a function of the available reasoning budget for a Qwen3-4B model on GSM8K. Truncated-Reasoning Self-Distillation (TRSD) substantially improves performance in low-budget regimes, enabling accurate predictions with limited reasoning.

Figure 2: Truncated-Reasoning Self-Distillation (TRSD). Given an input prompt $x$ , a frozen teacher model first generates a full chain-of-thought reasoning trace $r$ and an answer $y$ , and then evaluates the answer-token distribution $p_{\text{teacher}}(y\mid x,r)$ conditioned on the prompt and the complete reasoning trace. The trainable student model, initialized as a copy of the teacher, is conditioned only on a truncated prefix $\bar{r}$ of the teacher-generated reasoning trace, and evaluates the corresponding answer distribution $p_{\text{student}}(y\mid x,\bar{r})$ . Training minimizes the KL divergence between the teacher and student answer distributions, encouraging the student to recover the same predictions from partial reasoning and to remain accurate when inference-time reasoning is truncated.

Budget / Data	Countdown		Dolci		GSM8K		Math500
Budget / Data	Baseline	TRSD	Baseline	TRSD	Baseline	TRSD	Baseline	TRSD
Teacher Model: Qwen3-4B
32	11.72	27.24 ( - +15.5)	8.85	9.90 ( - +1.1)	19.03	22.67 ( - +3.6)	19.00	23.60 ( - +4.6)
64	13.46	28.06 ( - +14.6)	8.97	10.54 ( - +1.6)	20.85	32.68 ( - +11.8)	18.60	25.40 ( - +6.8)
128	26.04	35.12 ( - +9.1)	9.33	11.82 ( - +2.5)	26.31	63.53 ( - +37.2)	21.20	25.80 ( - +4.6)
256	43.94	49.72 ( - +5.8)	9.94	13.18 ( - +3.2)	54.66	86.28 ( - +31.6)	26.00	35.20 ( - +9.2)

Table 1: Baseline vs TRSD under truncated inference for in-distribution prompts. Entries report Baseline/TRSD accuracy (higher is better), bold indicates the better entry. When Baseline and TRSD accuracies differ by at most 1 percentage point, both are bold.

Teacher Model: Qwen3-4B, TRSD: GSM8K
Budget / Data	Countdown		Dolci		Math500
Budget / Data	Baseline	TRSD	Baseline	TRSD	Baseline	TRSD
32	11.72	12.62 (+0.9)	8.85	12.18 (+3.3)	19.00	23.20 (+4.2)
64	13.46	15.50 (+2.0)	8.97	12.14 (+3.2)	18.60	24.40 (+5.8)
128	26.04	24.88 (-1.2)	9.33	12.10 (+2.8)	21.20	27.40 (+6.2)
256	43.94	44.36 (+0.4)	9.94	13.46 (+3.5)	26.00	39.00 (+13.0)

Table 2: Baseline vs TRSD performance under truncated inference on out-of-distribution benchmarks. Entries report Baseline/TRSD accuracy (higher is better), bold indicates the better entry. When Baseline and TRSD accuracies differ by at most 1 percentage point, both are bold. Left: Baseline Qwen3-4B compared to TRSD trained on GSM8K. Right: Baseline Phi-4-mini-reasoning compared to TRSD trained on Countdown.

Panel of Figure 3 — Figure 3: Per-dataset accuracy as a function of the reasoning budget for Qwen3-4B. The evaluation dataset is specified below the respective plot.

Panel of Figure 4 — Figure 4: Per-dataset accuracy as a function of the reasoning budget for Phi-4-mini-reasoning. The evaluation dataset is specified below the respective plot.

		Countdown		Dolci		GSM8K		Math500
Model	Outcome	Baseline	TRSD	Baseline	TRSD	Baseline	TRSD	Baseline	TRSD
In-distribution evaluation
Qwen3-4B	Correct	1209	660	6915	3328	783	247	2810	1570
Qwen3-4B	Wrong	7169	3725	7913	4458	2546	661	5770	5243
Phi-4-mini-reasoning	Correct	1350	1095	6272	4346	683	513	1859	1831
Phi-4-mini-reasoning	Wrong	4418	3799	7585	6477	1983	1537	3523	3079

Table 3: Average number of reasoning tokens conditioned on answer correctness under unconstrained inference (

r_{\max}=8192

). Top block: in-distribution TRSD (trained and evaluated on the same prompt distribution). Bottom block: out-of-distribution TRSD using a single fixed checkpoint per model (Qwen3-4B trained on GSM8K; Phi-4-mini-reasoning trained on Countdown). Bold indicates the configuration (Baseline or TRSD) with lower average reasoning length.

Panel of Figure 5 — Figure 5: Example where both models answer correctly, but the TRSD-trained model uses a substantially shorter reasoning trace. The example is taken verbatim from the Qwen3-4B GSM8K evaluation set.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Countdown	Dolci	GSM8K	MATH500	OOD checkpoint
Qwen3-0.6B	500	2000	250	250	500 (Countdown)
Qwen3-1.7B	500	1000	250	250	250 (GSM8K)
Qwen3-4B	250	2000	250	250	250 (GSM8K)
Phi-4-mini-reasoning	1000	1000	500	500	500 (Countdown)
OpenThinker3-1.5B	1000	250	2000	1000	2000 (GSM8K)

Appendix

Table 4: Checkpoint selection summary. Columns corresponding to a specific dataset report the checkpoints achieving the best in-distribution performance for each (model, dataset) pair. The final column reports the checkpoint used for out-of-distribution evaluation, selected by best average performance across all four datasets; the dataset used for self-distillation is shown in parentheses. All entries denote the training step of the selected checkpoint.

Budget / Data	Countdown		Dolci		GSM8K		Math500
Budget / Data	Baseline	TRSD	Baseline	TRSD	Baseline	TRSD	Baseline	TRSD
Teacher Model: Qwen3-0.6B
32	14.06	22.14 ( - +8.1)	3.37	4.45 ( - +1.1)	3.03	8.34 ( - +5.3)	7.20	11.20 ( - +4.0)
64	13.84	22.12 ( - +8.3)	3.89	4.69 ( - +0.8)	3.41	10.24 ( - +6.8)	6.40	11.00 ( - +4.6)
128	17.18	29.86 ( - +12.7)	3.37	5.17 ( - +1.8)	4.47	25.85 ( - +21.4)	7.40	11.60 ( - +4.2)
256	25.82	42.46 ( - +16.6)	3.93	5.57 ( - +1.6)	23.96	54.59 ( - +30.6)	12.40	22.60 ( - +10.2)

Appendix

Table 5: Baseline vs TRSD accuracy under truncated inference for smaller and mid-size models (in-distribution prompts). Entries report Baseline/TRSD accuracy (higher is better). Colored values report the absolute change in accuracy (TRSD

-

Baseline, percentage points). Bold indicates the better entry; when Baseline and TRSD differ by at most 1 percentage point, both are bold.

Teacher Model: Qwen3-0.6B \| TRSD: Countdown
Budget / Data	Dolci		GSM8K		Math500
Budget / Data	Baseline	TRSD	Baseline	TRSD	Baseline	TRSD
32	3.37	4.01 ( - +0.6)	3.03	5.23 ( - +2.2)	7.20	11.40 ( - +4.2)
64	3.89	4.69 ( - +0.8)	3.41	5.46 ( - +2.1)	6.40	11.80 ( - +5.4)
128	3.37	4.77 ( - +1.4)	4.47	12.36 ( - +7.9)	7.40	13.80 ( - +6.4)
256	3.93	5.09 ( - +1.2)	23.96	41.17 ( - +17.2)	12.40	21.60 ( - +9.2)

Appendix

Table 6: Baseline vs TRSD accuracy under truncated inference on out-of-distribution benchmarks (smaller/mid-size models). Each block reports a fixed TRSD source dataset (shown in the block header) and evaluates on the remaining benchmarks. Entries report Baseline/TRSD accuracy (higher is better). Colored values report the absolute change in accuracy (TRSD

-

Baseline, percentage points). Bold indicates the better entry; when Baseline and TRSD accuracies differ by at most 1 percentage point, both are bold.

		Countdown		Dolci		GSM8K		Math500
Model	Outcome	Baseline	TRSD	Baseline	TRSD	Baseline	TRSD	Baseline	TRSD
In-distribution evaluation
Qwen3-0.6B	Correct	1169	724	5655	3765	848	575	2460	1575
Qwen3-0.6B	Wrong	4719	3230	7295	5268	2183	1888	5606	4506
Qwen3-1.7B	Correct	1098	747	6633	4590	1005	710	2879	2227
Qwen3-1.7B	Wrong	6379	2482	7777	6071	3364	2501	5829	5633

Appendix

Table 7: Average number of reasoning tokens conditioned on answer correctness for smaller models. Top block: in-distribution TRSD using per-dataset checkpoints selected as in Table A.3 . Bottom block: out-of-distribution-style TRSD using a single fixed checkpoint per model (Countdown step 500 for Qwen3-0.6B, GSM8K step 250 for Qwen3-1.7B, GSM8K step 2000 for OpenThinker3-1.5B). Bold indicates the configuration (Baseline or TRSD) with lower average reasoning length.

Panel of Figure 6 — Figure 6: Per-dataset accuracy as a function of the reasoning budget for Qwen3-0.6B. The evaluation dataset is specified below the respective plot.

Panel of Figure 7 — Figure 7: Per-dataset accuracy as a function of the reasoning budget for Qwen3-1.7B. The evaluation dataset is specified below the respective plot.

Panel of Figure 8 — Figure 8: Per-dataset accuracy as a function of the reasoning budget for OpenThinker3-1.5B. The evaluation dataset is specified below the respective plot.

Panel of Figure 9 — Figure 9: Example where both models answer correctly on a prompt from the Countdown dataset.

Panel of Figure 10 — Figure 10: Example where both baseline and TRSD-trained models fail on the same prompt, illustrating a shared failure mode.

Efficient Language Model Reasoning

Efficient language model reasoning studies how to improve problem-solving accuracy while reducing inference, memory, latency, and training costs.

See all 356 papers

Masked Self-Distillation: Internalizing the Chain-of-Thought in Language Models
Jun 18, 2026Durgesh Kalwar, Vardhan Palod, Jaya Adithya Pavuluri +1CoT Distillation Efficient Language Model Reasoning
SCAI, Arizona State University
Better Starts, Better Ends: Bootstrapped Iterative Self-Reasoning Distillation for Compressed Reasoning
Jul 17, 2026Leichao Dong, Dongxu Zhang, Yiding Sun +4CoT Distillation Efficient Language Model Reasoning
1Xi’an Jiaotong University · 2Peking University
Distilling the Essence: Efficient Reasoning Distillation via Sequence Truncation
Dec 24, 2025Wei-Rui Chen, Vignesh Kothapalli, Ata Fatahibaarzi +5CoT Distillation Selective Knowledge Distillation
1The University of British Columbia · 2LinkedIn · 3Canada Research Chair in NLP and ML

Jun 18, 2026J/K move · Enter open · S save

Masked Self-Distillation: Internalizing the Chain-of-Thought in Language Models

Durgesh Kalwar, Vardhan Palod, Jaya Adithya Pavuluri, Subbarao Kambhampati

SCAI, Arizona State University

CoT Distillation Efficient Language Model Reasoning LLM Inference Acceleration Language Model Distillation Self-Distillation CoT Compression

Open paper AskPDF arXiv

Large Reasoning Models produce long, explicit chains of intermediate steps before generating a final answer at inference time. These intermediate traces dominate latency, memory usage, and serving cost, even though final answer correctness is not causally related to the trace correctness and the trace length is not a reliable indicator of the problem complexity. This raises an obvious question: can the computation expressed in these intermediate tokens be internalized into the parameters of a language model, enabling it to produce answers with much shorter intermediate traces? We propose masked self-distillation, a knowledge-distillation based post-training framework in which copies of the same model are instantiated as teacher and student, and the student model is trained to internalize all or part of the intermediate trace, thus becoming more efficient at inference. We vary the fraction of intermediate trace the student is trained to internalize, interpolating between full internalization and no internalization. We conduct controlled experiments on two reasoning domains: math and graph coloring. We use the masked self-distillation framework to post-train Qwen3-4B & 8B models. Our results demonstrate that this method can be used to improve task performance while increasing inference efficiency across various domains and model sizes. We systematically analyze whether improved efficiency gain in the post-trained models generalize to OOD problems. We find that masked self-distillation models generalize well for in-domain OOD problems, and the masked self-distillation training does not induce catastrophic forgetting in the student model on out-of-domain problems. Furthermore, our ablation study shows that supervised fine-tuning can train models to produce shorter traces, but at the cost of generalization, highlighting the importance of on-policy training in masked self-distillation.

Abstract

Efficient Language Model Reasoning

Explore similar work

Masked Self-Distillation: Internalizing the Chain-of-Thought in Language Models

Better Starts, Better Ends: Bootstrapped Iterative Self-Reasoning Distillation for Compressed Reasoning

Distilling the Essence: Efficient Reasoning Distillation via Sequence Truncation