DRP: Distilled Reasoning Pruning with Skill-aware Step Decomposition for Efficient Large Reasoning Models

In plain words

Reasoning models often think out loud at great length. Much of that text repeats itself, backtracks, or checks things twice, and every extra word costs time and money to generate. Shorter answers from smaller models tend to be less accurate, though, so simply training a small model to copy a short answer does not always work.

DRP takes a middle path. The small model writes its own long reasoning first. A larger, more concise model then goes through that reasoning step by step, labels what each step is doing, and removes or shortens the unnecessary parts while keeping the logic. The small model is then trained on these cleaned-up traces. Because the traces still look like the small model's own style of thinking, it is easier for it to learn them.

The result is a smaller model that answers with fewer words on the maths benchmarks the authors tested, without losing accuracy in their runs. The numbers are the authors' own and come from specific settings, so check the results section before relying on them.

TL;DR

• Student-written long chain-of-thought is split into skill-labelled steps by a teacher LLM, then pruned by the same teacher §3 §3.1 §3.2.
• The pruned trace and answer become the SFT target for LoRA fine-tuning of R1-Distill-Qwen models §3.3 app. B.
• Author-reported, 7B student: GSM8K 917 to 328 tokens, accuracy 91.7% to 94.1% tab. 1.
• Accuracy on AIME24 is unchanged at 15/30 while tokens drop from 8674 to 4966 tab. 1.
• Code is public, but the repository has no LICENSE file, so reuse terms are unclear.

Problem

Large reasoning models produce long reasoning traces that include redundant and looping steps. This raises inference cost and can push the model toward wrong conclusions §1. Two families of fixes exist. Inference-time pruning stops generation early, which risks halting reasoning before it is complete. Distilling from short teacher outputs can lose accuracy when the student's long reasoning style does not match the concise teacher, a gap the paper calls a learnability gap §1.

The paper's setting is small student models (R1-Distill-Qwen 1.5B and 7B) trained on maths data, evaluated on in-domain GSM8K and out-of-domain MATH500, AIME24 and AMC §4.1 §4.2.

Method

DRP keeps the student's own long reasoning as the starting point and only edits it under teacher guidance. The student (e.g. R1-Distill-Qwen-7B) writes verbose, reflective traces, and the teacher (e.g. GPT-4o) writes concise ones §3. The paper describes three stages: decompose the trace into skill-based steps, prune and rewrite the steps with the teacher, and fine-tune the student on the revised traces §3 fig. 2.

1. The student produces a response R = (T, A): a reasoning trace T inside think tags and a final answer summary A §3.1.
2. The teacher splits T into non-overlapping segments, each with a skill label such as arithmetic, comparison or logical inference §3.1 fig. 2. The formal mapping D is written without an equation number §3.1.
3. For each step, the teacher picks one of four actions: KEEP, DELETE, SINGLE-STEP COMPRESS or MULTI-STEP COMPRESS §3.2 app. C.
4. The teacher rewrites the surviving steps into one fluent trace T_hat in the student's tone, and may revise the answer to A_hat so it agrees with T_hat §3.2.
5. The target is the concatenation R_hat = (T_hat, A_hat) §3.2.
6. The student is fine-tuned on pairs (question, R_hat) by minimising the negative log-likelihood of the target tokens, teacher-forced §3.3. The loss is also written without an equation number §3.3.

# Distilled Reasoning Pruning (DRP) data step  [§3] [fig. 2]
# Input:  question; student model; teacher model (e.g. GPT-4o)
# Output: (question, pruned trace, pruned answer)

def build_training_pair(question, student, teacher):
    # The student writes its own long response R = (T, A)  [§3]
    response = student.generate(question)
    trace, answer = split_think_and_answer(response)

    # Teacher cuts the trace T into steps, each with a skill label  [§3.1]
    steps = teacher.segment_with_skills(trace)

    # Teacher picks KEEP, DELETE, SINGLE-STEP COMPRESS or
    # MULTI-STEP COMPRESS for each step and revises it  [§3.2] [app. C]
    revised_steps = teacher.apply_actions(steps)

    # Teacher merges the revised steps into one fluent trace T_hat  [§3.2]
    pruned_trace = teacher.rewrite_as_trace(revised_steps)

    # Teacher may revise the answer so it agrees with T_hat  [§3.2]
    pruned_answer = teacher.revise_answer(answer, pruned_trace)

    # Training pair (x, R_hat) with R_hat = (T_hat, A_hat)  [§3.2] [§3.3]
    return question, pruned_trace, pruned_answer

# Supervised fine-tuning loss on one pair  [§3.3]
# Input:  model; question; pruned_trace; pruned_answer
# Output: scalar loss

def sft_loss(model, question, pruned_trace, pruned_answer):
    # Target tokens y_1 .. y_n of the concatenated response  [§3.2] [§3.3]
    target_tokens = tokenize(pruned_trace + pruned_answer)
    loss = 0
    for position, token in enumerate(target_tokens):
        # Each token is scored given the question and the true earlier
        # target tokens (teacher forcing)  [§3.3]
        log_prob = model.log_prob(token, question, target_tokens[:position])
        loss = loss - log_prob
    # Summed over tokens as written; no normalisation is stated  [§3.3]
    return loss

The four actions and the teacher prompts are given in app. C. The paper also reports that the segmentation prompt asks the teacher to label each step with a skill and not to solve or change the content app. C.

Training uses LoRA on the R1-Distill-Qwen models with LLaMA-Factory §4.4 app. B. Ablations replace the skill-based segmentation with a default sentence split or skip segmentation altogether §6.1 tab. 2.

Hyperparameters

Values not stated in the paper are marked "not stated". Code values are included where the repository sets them, and every row is cited.

Results

Author-reported. Setup: R1-Distill-Qwen-1.5B and 7B students, GPT-4o teacher, zero-shot evaluation with vLLM, Pass@1 averaged over five runs §4 §4.6. In-domain is GSM8K; out-of-domain is MATH500, AIME24 and AMC §4.1. Baselines are TALE (prompt-based), CoT-Valve (SFT) and ThinkPrune (RL) §4.3. Tokens are averages per completion. The base-model token counts in tab. 1 equal the 12k-cutoff row of tab. 5, which suggests the cutoff is applied there too §4.6 tab. 5.

Source: tab. 1. The paper's own percentage reductions for DRP are 50%, 35%, 42% and 44% for the 1.5B model, and 64%, 28%, 43% and 33% for the 7B model tab. 1. The paper calls the 1.5B GSM8K gain "12.7%"; from the table it is 12.7 percentage points (70.7% to 83.4%) §5.1 tab. 1.

Decomposition ablation (7B) tab. 2:

Direct distillation from GPT-4o concise traces versus DRP, 7B tab. 4:

Teacher choice, 7B student tab. 3:

12k cutoff effect on the 7B base tab. 5:

Token-length distributions are shown in fig. 3 and fig. 4. The paper reports that DRP shifts mass toward shorter completions and away from the very long tail on AIME and AMC §5.3 fig. 3.

Paper vs code

Official repository YuxuanJiang1/DRP at code_commit a196237. The owner matches the paper's first author and the repository is linked from the abstract. The repository has no LICENSE file, so code is cited by path and line only, with no quotes.

• Generation settings. Paper: vLLM with maximum generation length 131,072 §4.5. Code: 40000 for GSM8K sft_gsm8k.py and 130000 for the OOD tasks sft_ood.py.
• Runs. Paper: Pass@1 averaged over five independent runs §4.6. Code: one pass per task with temperature 0.0 sft_gsm8k.py sft_ood.py. Under greedy decoding repeated runs are essentially identical, so the averaging cannot be reproduced from these scripts.
• Token cutoff. Paper: 12k cutoff for token counts §4.6 app. A. Code: the token summary averages over all responses with no cap sft_gsm8k.py sft_ood.py. The cutoff used for the paper's tables cannot be found in the repository.
• Learning-rate warmup. Paper: not mentioned app. B. Code: warmup ratio 0.1 qwen_lora_sft.yaml.
• Hardware. Paper: two A100 80GB app. B. Code: the run script exposes only GPU 0 run_all.sh.
• Training data source. Paper: GPT-4o is the primary teacher §4.2. Code: the training config points to a dataset folder and name that refer to a Gemini export qwen_lora_sft.yaml. The script that builds the SFT set is not in the repository, so the teacher behind the released training data cannot be confirmed from the code.
• Segmentation prompt. Paper appendix prompt app. C differs in wording from the code prompt, which additionally asks the teacher to keep the final answer line as a skill labelled "clarify answer" split_verify.py.
• Compression prompt. Paper appendix prompt app. C takes a list of labelled steps. Code also passes the question split_verify.py and asks the teacher to start its output with a "##" heading split_verify.py. The script keeps the text after that heading as the compressed response split_verify.py.
• Segmentation run. The segmentation batch request is commented out in the main workflow. The script only downloads a previously created batch output split_verify.py.
• Tokenizer. Paper: Qwen tokenizer for counting §4.6. Code: tokenizer of the merged student checkpoint sft_gsm8k.py.
• Data availability. The repository readme says data is "coming soon" readme.md. The GSM8K CSVs in datasets/ are the original question and answer files, not the pruned targets.
• Version. The arXiv metadata is version v4, updated 2026-04-24. The repository readme says the paper was accepted at ACL 2026 Findings readme.md. This card follows the LaTeX source fetched by wikillm fetch.

Pitfalls

• Student traces are not in the repository. The generation script, sampling settings and the file of initial traces are absent. The split script reads a precomputed CSV split_verify.py. Reproducing the data needs the student's sampling configuration, which the paper does not give §3.
• No answer check. The paper describes no filter that drops pairs with a wrong final answer. The script named split_verify segments and compresses, and does not compare answers split_verify.py. Decide whether to filter on correctness.
• Silent truncation risk. Each teacher call has a 2048-token output cap split_verify.py. A long trace can be cut short, and if no "##" heading is found the script stores the placeholder text "[no ## result found]" split_verify.py. Filter that placeholder out before SFT.
• Resume logic is positional. The script resumes by counting existing rows and slicing the input from that offset split_verify.py. Reordering or dropping rows will misalign results.
• Unknown loss normalisation. The loss is written as a sum over target tokens §3.3. Whether the trainer averages it is not stated, and the choice changes the effective learning rate.
• LoRA details missing. Only "all" linear layers is given qwen_lora_sft.yaml. Rank, alpha, dropout and the optimizer are not stated, so the adapter cannot be reproduced exactly from the paper.
• Effective batch size is ambiguous. Batch 2 with accumulation 4 gives 8 per device app. B. With two GPUs the effective batch could be 16. The paper does not say how the GPUs were used.
• Chat template. The code uses the qwen template qwen_lora_sft.yaml. The exact wrapping of the target text is therefore set by LLaMA-Factory, not by the paper.
• The 12k cutoff is undefined. The paper does not say whether over-cutoff generations are truncated, counted as 12k tokens, or dropped from the average §4.6 app. A. tab. 5 keeps accuracy nearly unchanged while tokens drop, which suggests it is applied when measuring. The rule for unfinished answers is not given.
• The "up to 5x" figure understates the AIME row. The appendix says the average token count can rise by up to 5x without the cutoff app. A. The AIME24 row of tab. 5 goes from 8674 to 54339 tokens, which is about 6.3x.
• Teacher claim is broader than the table. The paper says every teacher improves both accuracy and token count over the base model §6.3. In tab. 3, ChatGPT is below the base on GSM8K (91.2% vs 91.7%), and Gemini, ChatGPT and DeepSeek-V3 are below the base on MATH500 (91.8%, 90.2%, 90.4% vs 92.4%). Gemini uses fewer tokens than GPT-4o on AIME24 and AMC. The text says GPT-4o is strongest "especially in token compression", which holds on GSM8K and MATH500 only.
• Reported token averages differ. Section 6.2 gives DRP's training traces an average of 330 tokens and the direct-distillation traces 186 §6.2. The main table reports 328 tokens for DRP on GSM8K at evaluation tab. 1. The paper does not say whether these measure the same thing.
• Teacher API terms. The paper uses OpenAI and Gemini API outputs as training targets §9. It does not discuss whether those terms allow training other models on the outputs. Check this before commercial use of any model trained on the data.
• Cost is hidden. The pipeline makes two teacher calls per example, one to segment and one to compress split_verify.py split_verify.py. The segment call is commented out in the main script, so only the compress call runs by default. Teacher API cost and student generation cost are not reported.
• Small samples. Tables report integer counts such as 15/30 and 33/40, and the five-run averaging is not visible in those counts tab. 1. A difference of one or two problems on AIME24 or AMC is a small sample, and the paper gives no variance.

Lineage

• Compared against: TALE, a prompt-based token budget method [arXiv:2412.18547]; CoT-Valve, SFT on length-varied chains [arXiv:2502.09601]; ThinkPrune, RL with a length target [arXiv:2504.01296] §4.3 tab. 1. None of these have a card here, so nothing from their content is used.
• Builds on: the DeepSeek-R1-Distill-Qwen student models §4.2. The paper does not name a prior pruning method it extends.
• Superseded by: not known from content in this wiki.