Long-Chain Reasoning Distillation via Adaptive Prefix Alignment

TL;DR

The paper addresses the challenge of distilling long, uncertain reasoning trails from teacher models into smaller student models. It introduces P-ALIGN, which adaptively truncates teacher CoTs to a minimal sufficient prefix using student self-judgment and a binary search, then aligns the student to this prefix to supervise SFT. Empirical results across four mathematical benchmarks show that P-ALIGN consistently outperforms zero-shot and existing distillation baselines, with notable gains over Long-CoT-based SFT and UPFT, across different backbones. The approach reduces supervision length while preserving or enhancing reasoning quality, and analyses show that early prefixes carry stable, informative signals while late steps add uncertainty. Overall, P-ALIGN advances practical distillation of structured reasoning for math tasks and informs data design for prefix-based supervision, albeit with computational costs tied to teacher CoT generation and self-judgment reliability.

Abstract

Large Language Models (LLMs) have demonstrated remarkable reasoning capabilities, particularly in solving complex mathematical problems. Recent studies show that distilling long reasoning trajectories can effectively enhance the reasoning performance of small-scale student models. However, teacher-generated reasoning trajectories are often excessively long and structurally complex, making them difficult for student models to learn. This mismatch leads to a gap between the provided supervision signal and the learning capacity of the student model. To address this challenge, we propose Prefix-ALIGNment distillation (P-ALIGN), a framework that fully exploits teacher CoTs for distillation through adaptive prefix alignment. Specifically, P-ALIGN adaptively truncates teacher-generated reasoning trajectories by determining whether the remaining suffix is concise and sufficient to guide the student model. Then, P-ALIGN leverages the teacher-generated prefix to supervise the student model, encouraging effective prefix alignment. Experiments on multiple mathematical reasoning benchmarks demonstrate that P-ALIGN outperforms all baselines by over 3%. Further analysis indicates that the prefixes constructed by P-ALIGN provide more effective supervision signals, while avoiding the negative impact of redundant and uncertain reasoning components. All code is available at https://github.com/NEUIR/P-ALIGN.

Figures (8)

• Figure 1: The Framework of Our P-ALIGN Model. Later steps in traces are often more uncertain and harder for the student to learn. P-ALIGN therefore performs adaptively selecting a sufficient prefix and aligning student supervision to it for more effective distillation.
• Figure 2: Performance of Distilled Student Models under Different Prefix Truncation Strategies. We compare models distilled using fixed prefix truncation ratios and P-ALIGN (Figures \ref{['fig:effective-of-price:aime']} and \ref{['fig:effective-of-price:math500']}), analyze their response lengths (Figure \ref{['fig:effective-of-price:response-length']}), and evaluate CoT quality with GLM-4.5 as the judge (Figure \ref{['fig:effective-of-price:gpt_score']}).
• Figure 3: Effectiveness of Chunks in Long-Form CoTs at Different Positions. Figure \ref{['fig:effective-of-prefix:entropy']} shows the entropy scores of student models across different chunks at varying positions. Figure \ref{['fig:effective-of-prefix:Positional']} compares the performance of student models when using prefix, middle, or suffix segments as supervision.
• Figure 4: The Prompt Templates Used for Direct Question Answering.
• Figure 5: The Prompt Templates Used for Prefix Evaluation.