Student-in-the-Loop Chain-of-Thought Distillation via Generation-Time Selection

Chaoqun He; Yingfa Chen; Chaojun Xiao; Xu Han; Lijie Wen

cs.CLApr 3, 2026

Student-in-the-Loop Chain-of-Thought Distillation via Generation-Time Selection

Authors: Chaoqun He, Yingfa Chen, Chaojun Xiao, Xu Han, Lijie Wen

Organizations: Tsinghua University

Abstract

Large reasoning models achieve strong performance on complex tasks through long chain-of-thought (CoT) trajectories, but directly transferring such reasoning processes to smaller models remains challenging. A key difficulty is that not all teacher-generated reasoning trajectories are suitable for student learning. Existing approaches typically rely on post-hoc filtering, selecting trajectories after full generation based on heuristic criteria. However, such methods cannot control the generation process itself and may still produce reasoning paths that lie outside the student's learning capacity. To address this limitation, we propose Gen-SSD (Generation-time Self-Selection Distillation), a student-in-the-loop framework that performs generation-time selection. Instead of passively consuming complete trajectories, the student evaluates candidate continuations during the teacher's sampling process, guiding the expansion of only learnable reasoning paths and enabling early pruning of unhelpful branches. Experiments on mathematical reasoning benchmarks demonstrate that Gen-SSD consistently outperforms standard knowledge distillation and recent baselines, with improvements of around 5.9 points over Standard KD and up to 4.7 points over other baselines. Further analysis shows that Gen-SSD produces more stable and learnable reasoning trajectories, highlighting the importance of incorporating supervision during generation for effective distillation.

Figures & tables

Figure 1: Comparison between standard KD and our proposed Gen-SSD.

Figure 2: Overview of Gen-SSD. The student actively participates in the teacher’s multi-sample generation process. At each chunk, the student evaluates candidate continuations with PPL and selects the fragments best aligned with its capability, thereby influencing the teacher’s sampling trajectory. For unsuitable candidates, generation is terminated early, which reduces inference cost and improves sampling efficiency.

Method	AIME25	AIME24	AMC23	OlympiadBench	GSM8K	Avg $\uparrow$
Teacher	49.79	67.50	91.09	66.91	94.20	73.90
Student	0.83	3.33	19.06	15.73	38.74	15.54
Cold Start	3.75	0.42	19.53	16.17	51.63	18.30
Standard KD	7.08	7.08	35.00	24.78	71.11	29.01
Self-Distillation	2.71	2.71	29.53	22.55	63.76	24.25
MCC-KD	5.21	5.63	35.47	27.15	77.79	30.25

Table 1: Performance comparison across various benchmarks. We adopt QwQ-32B as the teacher model and Qwen2.5-Math-1.5B as the student model. Best results are bolded. The chunk size for the Gen-SSD is 4K. All models are evaluated in a zero-shot setting, except for the student model, which is evaluated with two-shot prompting.

Method	PPL $\downarrow$	Token length $\downarrow$
MoRSD	2.46	6139.10
Gen-SSD	1.57	5049.60

Table 2: Comparison between Gen-SSD and MoRSD in terms of student PPL and token length.

Method	AIME25	AIME24	AMC23	OlympiadBench	GSM8K	Avg $\uparrow$
Cold Start	3.75	0.42	19.53	16.17	51.63	18.30
w/o cold-start	6.25	6.88	34.84	28.64	68.84	29.09
w/ cold-start	10.00	10.21	41.88	30.71	81.88	34.94

Table 3: Comparison under settings with and without cold-start initialization.

Figure 3: Average performance of Gen-SSD across benchmarks under different chunk sizes. Detailed results are provided in Table 9 in the appendix.

Setting	AIME25	AIME24	AMC23	OlympiadBench	GSM8K	Avg $\uparrow$
High PPL	7.08	4.37	34.06	25.82	68.39	27.94
Random PPL	9.58	7.92	35.78	28.64	74.75	31.33
Low PPL	10.00	10.21	41.88	30.71	81.88	34.94

Table 4: Comparison of three PPL-based data selection strategies, demonstrating the superiority of low-PPL data. The experimental setting follows Table 1 .

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Value
Learning Rate	$2\times 10^{-5}$
Number of Epochs	5
Number of Devices	4
Batch Size	32
Optimizer	AdamW
Learning Rate Scheduler	cosine

Appendix

Table 5: The hyperparameters used for fine-tuning.

Setting	Time $\downarrow$	Accuracy $\uparrow$
1K
16 candidates	35.46	31.22
Decreasing sampling	28.20	31.73
4K
16 candidates	34.22	33.66
Decreasing sampling	31.59	34.94

Appendix

Table 6: Comparison results. Accuracy is computed as the average performance over the five benchmarks in Table 1 . The best scores are bolded.

Method	AIME25	AIME24	AMC23	OlympiadBench	GSM8K	Avg $\uparrow$
Standard KD*	7.92	6.46	38.12	26.85	78.24	31.52
Gen-SSD*	10.00	6.88	44.37	31.01	82.26	34.90
Cold Start	6.04	4.79	29.06	24.93	65.96	26.16
Standard KD	9.58	7.50	37.97	30.56	77.41	32.60
Gen-SSD	9.79	8.96	45.16	32.49	84.84	36.25

Appendix

Table 7: Performance comparison across various benchmarks with closed-source teacher models. The highest scores are bolded. Here, * indicates direct distillation on Qwen2.5-Math-1.5B without the cold-start stage.

Method	MBPP	ARC-C	StrategyQA	CommonsenseQA	Average $\uparrow$
Standard KD	20.03	38.65	63.23	41.66	40.89
MCC-KD	17.76	38.14	64.37	40.67	40.24
MoRSD	22.07	38.91	64.93	40.43	41.59
Gen-SSD	21.36	39.68	65.68	42.07	42.20

Appendix

Table 8: Evaluation results on MBPP, ARC-C, StrategyQA, and CommonsenseQA, demonstrating generalization beyond mathematical reasoning. The best scores are bolded.

Method	AIME25	AIME24	AMC23	OlympiadBench	GSM8K	Avg $\uparrow$
Standard KD*	4.37	3.75	33.28	23.29	69.52	26.84
Gen-SSD
512	8.33	6.04	33.59	26.71	74.53	29.84
1K	8.54	6.25	37.5	28.19	78.17	31.73
2K	7.50	6.67	40.78	27.60	74.37	31.38
4K	10.00	10.21	41.88	30.71	81.88	34.94

Appendix

Table 9: Benchmarks performance of Gen-SSD with varying chunk sizes, compared against Standard KD as the baseline. Best results are bolded. Here, * indicates direct distillation on Qwen2.5-Math-1.5B without the cold-start stage.

Method	AIME25	AIME24	AMC23	OlympiadBench	GSM8K	Avg $\uparrow$
R1-7B
Standard KD*	7.92	8.33	38.91	26.26	76.27	31.54
Gen-SSD	10.83	9.79	41.88	31.16	82.56	35.24
R1-14B
Standard KD*	10.21	6.88	37.81	26.41	77.10	31.68
Gen-SSD	12.08	9.17	45.31	30.71	80.52	35.56

Appendix

Table 10: Comparison of Standard KD and Gen-SSD using R1-7B/14B/32B as teacher models. The highest scores are bolded. Here, * indicates direct distillation on Qwen2.5-Math-1.5B without the cold-start stage.

Figure 6: Comparison of PPL trends between Gen-SSD and MoRSD.

Language Model Distillation

Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.

See all 512 papers

MI-Distillation: Selecting from Model-Interpolated Instruct-Reasoning Data Spectrum for Chain-of-Thought Distillation
Aug 30, 2026Yangsong Lan, Renkai Hu, HongKai Zheng +4CoT Distillation Language Model Distillation
College of Artificial Intelligence, Nanjing University of Aeronautics and Astronautics, Nanjing, China · 2The Key Laboratory of Brain-Machine Intelligence Technology, Ministry of Education, Nanjing, China
Explain in Your Own Words: Improving Reasoning via Token-Selective Dual Knowledge Distillation
Feb 25, 2026Minsang Kim, Seung Jun BaekLanguage Model Distillation Selective Knowledge Distillation
Korea University · SK Telecom
From Solver Trajectories to Teaching Trajectories: Cognition-Aligned Reasoning Distillation
Oct 27, 2026Yashuo Luo, Qianren Mao, Ziqi Qin +6Reasoning Distillation Language Model Distillation
National College for Excellent Engineers, Beihang University · Zhongguancun Laboratory · Beihang University +3

Aug 30, 2026J/K move · Enter open · S save

MI-Distillation: Selecting from Model-Interpolated Instruct-Reasoning Data Spectrum for Chain-of-Thought Distillation

Yangsong Lan, Renkai Hu, HongKai Zheng, Bo Zhang+3

College of Artificial Intelligence, Nanjing University of Aeronautics and Astronautics, Nanjing, China · 2The Key Laboratory of Brain-Machine Intelligence Technology, Ministry of Education, Nanjing, China

CoT Distillation Language Model Distillation Knowledge Distillation Large Reasoning Models

Open paper AskPDF arXiv

Recent advances in large reasoning models (LRMs) have shown strong performance on complex problems through long chain-of-thought (Long CoT) reasoning. However, distilling such trajectories into smaller student models remains challenging: direct Long CoT supervision often provides limited gains and can be less effective than concise Short CoT rationales. In this work, we investigate this phenomenon from a gradient-centric perspective. Our analysis shows that Long CoT induces larger gradient magnitudes and more concentrated update directions than Short CoT, with this effect becoming more pronounced as student model capacity increases. These findings suggest that effective Long CoT distillation requires balancing the reasoning information density of reasoning trajectories with their distributional alignment to the student model. Motivated by this insight, we propose Model Interporlation Distillation (MI-Distillation), a framework that constructs a continuous Instruct-Reasoning data spectrum through model interpolation. To select suitable trajectories from this spectrum, we further introduce Sequential Learnable Surprisal Score (SeqLSS), which favors reasoning paths that are both informative and learnable for the student. Extensive experiments on reasoning benchmarks show that MI-Distillation consistently improves small model CoT distillation over strong Long CoT baselines.

Abstract

Language Model Distillation

Explore similar work

MI-Distillation: Selecting from Model-Interpolated Instruct-Reasoning Data Spectrum for Chain-of-Thought Distillation

Explain in Your Own Words: Improving Reasoning via Token-Selective Dual Knowledge Distillation

From Solver Trajectories to Teaching Trajectories: Cognition-Aligned Reasoning Distillation