ShortOPD: Recovering Pruned LLMs with Short-to-Long On-Policy Distillation

Qingyu Zhang; Qianhao Yuan; Hongyu Lin; Yaojie Lu; Xianpei Han; Le Sun; Ming Xu; Jiarui Li

Source-linked AI summary

ShortOPD: Recovering Pruned LLMs with Short-to-Long On-Policy Distillation

Qingyu Zhang, Qianhao Yuan, Hongyu Lin, Yaojie Lu, Xianpei Han, Le Sun, Ming Xu, Jiarui Li

arXiv:2607.13124v2Jul 14, 2026cs.LGcs.AIcs.CL

TL;DR

Structured pruning is mainly validated on recognition benchmarks, despite compressed models often failing at free-form generation. ShortOPD uses repetition-gated, truncation-aware on-policy distillation to adapt rollout horizons, matching fixed short and long horizons within two points while using fewer rollout tokens.

Problem
Structured pruning lacks sufficient evaluation and recovery evidence for the free-form generation quality required in deployment.
Method
ShortOPD uses repetition-gated, truncation-aware on-policy distillation to adapt rollout horizons from short to long as usable prefixes improve.
Results
ShortOPD matches fixed short and long horizons within two points while using fewer rollout tokens and outperforms standard recovery baselines across tested pruning settings.
Takeaways & Limitations
Recovery quality depends on training over the compressed model’s visited states and matching rollout budgets to the generation lengths it can currently use.
Takeaways & Limitations
Generality beyond Qwen3 models, the studied pruning settings, compression ratios, and scales remains unestablished.

Abstract

from arXiv · show

Structured pruning is a hardware-friendly way to compress LLMs, but it is mostly validated on multiple-choice recognition tasks, while the same compressed checkpoints can collapse on the free-form generation that deployment actually requires. Two observations trace this gap. First, greedy \textsc{pass}@$1$ nearly vanishes after compression, yet \textsc{pass}@$k$ recovers substantially under repeated sampling: useful generations are demoted, not erased. Second, the recoverable regime fails mainly through suffix repetition. Recovery should therefore train on the compressed model's own on-policy states with dense token-level supervision, which On-Policy Distillation (OPD) provides by reusing the pre-compression model as a frozen teacher. However, long on-policy rollouts spend early recovery budget on low-information repetitive suffixes, delaying loss descent. To mitigate this waste, we propose \textbf{\shortopd}, a short-to-long OPD schedule that detects teacher-confirmed repetitive suffixes, treats the surviving prefix as each rollout's effective length, and allocates future rollout budgets to the effective lengths the policy can currently use. Across math, code, and open-ended generation, \shortopd\ raises the compressed model's score to about $9\times$ its unrecovered value and $1.6$--$4.4\times$ standard recovery recipes (SFT w/o KD, KD, and SeqKD), and it matches a fixed $8192$-token rollout horizon within two points using a quarter of the training time ($8.5$ vs.\ $35.9$ hours) and $71\%$ fewer rollout tokens. We hope this recipe helps move structured pruning beyond marginal gains on perplexity and multiple-choice benchmarks, a step closer to deployment-ready generation quality.

1 Introduction

ShortOPD recovers structurally pruned LLMs by distilling the pre-compression model on the student’s own states while adapting rollout horizons to avoid repetitive suffixes. It restores broad generation quality efficiently, but recovery depends on corpus coverage across capabilities.

Motivation: Pass@k rises with additional sampling on every benchmark, reaching 91% at k=64 on GSM8K versus 88% greedy for the unpruned teacher, showing correct trajectories remain sampleable.The compressed model’s missing ability is demoted rather than erased, motivating recovery from its own sampling distribution.
Motivation: OPD trains on student-sampled rollouts using dense next-token supervision from the frozen pre-compression self, without labels, verifiers, or an external teacher.This directly combines on-policy states with token-level teacher distributions.
Motivation: During early fixed-H=2048 OPD, 55–75% of rollouts end in suffix loops while distillation loss remains low until about step 80, wasting rollout and teacher compute.These repetitive suffixes provide little additional teacher–student correction under a long fixed horizon.
Method: ShortOPD uses repetition-gated, truncation-aware horizon control, with smoothed effective-length targets and gradual budget updates that shrink under repetition and grow when truncation is high.The controller uses exponential moving averages to prevent abrupt horizon changes.
Experiments: ShortOPD restores about two-thirds of the unpruned teacher’s generation score across eight task families, outperforming SFT w/o KD, KD, SeqKD, and sparse-reward RLVR under matched budgets.It matches fixed rollout horizons up to 8192 tokens within two points while generating up to 71% fewer rollout tokens.
Experiments: Removing math or code prompts sharply damages recovery on the corresponding capabilities, making recovery-corpus coverage a first-class design choice.On-policy distillation repairs only states visited by the student’s own search.

2 Related Work

Prior work studies structured pruning through depth- and width-based compression, while distillation research increasingly addresses the mismatch between fixed teacher outputs and student inference states. These lines of work provide the context for recovery after pruning.

Structured pruning and the recognition-generation gap: Structured pruning removes whole transformer blocks or coupled within-layer structures, using criteria such as influence scoring, layer merging, or learned masks.Depth pruning includes ShortGPT [5] and LaCo [33], while width pruning includes Sheared LLaMA [6] and Minitron [7] [37].
Distillation and on-policy recovery signals: Distillation methods have progressed from fixed teacher-generated responses to student-state supervision, including imitation-learning KD [39], MiniLLM [40], divergence generalizations [41] [42], and generalized on-policy KD [23].Sequence-level KD [26] is off-policy because it trains on fixed teacher-generated responses, whereas on-policy approaches distill on student states.

3 Method

ShortOPD recovers structurally compressed LLMs by combining dense on-policy distillation from the frozen pre-compression teacher with a repetition-gated short-to-long rollout schedule. It uses the compressed model’s own states while adapting horizons to usable prefixes, without requiring ground-truth responses or reward verification.

Structured pruning: The pruning experiment removes roughly 25% of parameters using low Block Influence layers, but ShortOPD assumes only a structurally compressed student and its pre-compression self.The original model remains a frozen teacher.
On-policy self-distillation: OPD runs the frozen pre-compression teacher on student-generated trajectories and matches next-token distributions at every response position.The loss is on-policy, dense, and reward-free, using top-100 logits plus aggregated tail mass with clipped importance weighting.
Short-to-long schedule: ShortOPD adapts each rollout budget because long horizons can enter repetitive suffixes before recovery, while permanently short horizons can truncate legitimate long generations.The global response_length remains a padding and context ceiling, while per-step budgets H_t vary within H_max.
Generation-side feedback: The controller detects severe repetition in terminal tokens, refines its onset with OPD divergence and teacher NLL, and treats the surviving prefix as the effective length.These statistics reuse existing rollout tokens and teacher scores; the current batch still receives the unmodified dense OPD loss over all generated tokens.
Repetition-gated horizon control: High repetition shrinks the next horizon toward the observed usable length, whereas clean truncation permits growth only when the current horizon is binding.A hysteresis band and EMA smoothing prevent oscillation, and high repetition takes priority when repetition and truncation coincide.
Design space: ShortOPD combines on-policy student rollouts, dense teacher distributions, and repetition-gated horizon control, unlike recovery methods that omit one of these dimensions.It requires neither ground-truth responses nor a reward verifier.

4 Experiments

Across eight free-form generation task families, ShortOPD substantially restores structured-pruned models and consistently outperforms offline recovery baselines across scales and pruning structures. Its short-to-long controller achieves comparable quality to long fixed rollouts with substantially lower training cost, while additional on-policy exposure and domain-matched data further improve recovery.

Main results: 48.46 Avg restores nearly 9× the pruned model’s 5.71 and exceeds the strongest off-policy baseline by 18 points, while ShortOPD also improves code and all four open-ended tasks.On 8B, ShortOPD reaches 42.84 and 53.71 Avg for no-thinking and Thinking models, respectively, leading every generation task in the Thinking evaluation.
Closed-loop recovery trajectory: ShortOPD reduces early repetitive warm-up and reaches the post-warm-up regime in about 40–50 steps versus roughly 80 for Vanilla fixed at H=2048.The controller lowers its budget from 2048 to 1024 during high repetition, then returns to 2048 as clean truncations take over.
Efficiency: 8.5 hours and 250M rollout tokens achieve Avg comparable to fixed 8192-token rollouts, whose quality ranges only from 48.5 to 50.2 across schedules.Relative to fixed 2048, ShortOPD reduces mean generation time by 29%; relative to fixed 8192, it uses 76% less time.
Scaling with exposure: ShortOPD scales with exposure, raising Avg from 48.46 after one epoch to 55.41 after three epochs, or 73.7% of the teacher, with broad gains across math, code, and open-ended tasks.The improvements span GSM8K, MATH-500, HumanEval, MBPP, Alpaca, QA, and MT-Bench.
Recovery signal: ShortOPD’s on-policy states matter: SeqKD loses 19.9 Avg points and KD trails ShortOPD by 17.9 points despite using the same frozen teacher.All methods start from the same compressed student and use the same step budget and prompt distribution, while OPD and ShortOPD match the teacher on student-generated states.
Ablations and generality: Domain-matched recovery data and pruning structure both matter: removing math nearly eliminates mathematical recovery, removing code collapses execution scores, and ShortOPD beats KD by 6.44–9.01 Avg points across width-only and hybrid configurations.The experiments cover math, code, and open-ended instruction data, with evaluation restricted to free-form generation across eight task families.

5 Conclusion

Structured pruning can preserve latent capability while failing to promote it during generation, making recovery a distributional repair problem. Across pruning settings and model modes, on-policy self-distillation substantially outperforms standard recovery baselines, despite low-quality early rollouts.

5 Conclusion: On-policy self-distillation substantially outperforms SFT w/o KD, SeqKD, and KD across Qwen3 models, modes, and depth, width, or hybrid pruning.In matched Math+Code controls, it also outperforms sparse-reward RLVR.
5 Conclusion: Structured pruning can leave latent capability in the model’s search space while failing to promote it during generation, framing recovery as distributional repair.
5 Conclusion: The main cost of on-policy self-distillation is low-quality early rollout.

6 Limitations · Appendix

The study’s scope is limited to BI-depth, width-only, and hybrid pruning on Qwen3 models at 4B and 8B, leaving generality across broader settings unresolved. Further evaluation is needed to determine when compression becomes too aggressive for lightweight recovery.

6 Limitations: The experiments cover BI-depth, width-only, and hybrid pruning on Qwen3 models at 4B and 8B.
6 Limitations: Broader model families, compression ratios, and model scales are needed to establish generality and identify when compression becomes too aggressive for lightweight recovery.

A Training details · B Terminal periodic-loop detector

The paper recovers structurally pruned students through dense on-policy distillation with fixed training and baseline protocols, while ShortOPD detects severe terminal repetition and refines each rollout’s effective length using teacher-confirmed loss signals.

A Training details: Pruning removes nine of 36 transformer blocks using held-out-corpus BI calibration, with the frozen original as teacher and the compressed model initializing the actor.The same 36 →27 depth reduction is used for Qwen3-4B-Thinking and Qwen3-8B, with backbone-specific calibration.
A Training details: Main recovery uses rollout group 8, batch 64, 4096-token prompt and response limits, temperature 0.8, learning rate 2 × 10−6, and 710 training steps.The objective is top-100+tail generalized JSD with α = 0.5, importance weights clipped at 2.0, and no policy-gradient loss.
A Training details: Offline SFT w/o KD, SeqKD, and KD baselines use the same prompt distribution, 8192-token sequences, batch 64, learning rate 2 × 10−6, and 710 steps.KD uses forward KL at temperature 1.0 without a hard-label cross-entropy term.
A Training details: Efficiency runs reuse the main corpus, objective, batch size, rollout count, and learning rate while comparing ShortOPD with fixed response ceilings of 2048 and 8192 tokens.The controller checkpoints its current budget and all EMA state across restarts.
A Training details: Dense OPD always trains on generated tokens, and only the independent repetition-rate threshold can open the shrink gate; low loss alone cannot trigger shortening.The probe does not mask the loss.
B Terminal periodic-loop detector: The detector examines the last W=512 valid response tokens for periods 1–10 and uses shifted token comparisons, allowing detection when truncation cuts through a cycle.High agreement at the raw response end prevents repeated material followed by a clean ending from being classified as terminal repetition.
B Terminal periodic-loop detector: A severe loop requires terminal agreement η=0.9, at least three cycles, a 64-token minimum tail, and either a 128-token tail or 30% of the response.The detector selects the longest explained terminal tail, then higher agreement and shorter period, returning its onset as the structural effective-length candidate.
B Terminal periodic-loop detector: ShortOPD locally refines the detector’s onset by finding the first 32-token window meeting absolute or relative thresholds for OPD loss and teacher NLL.This refinement changes the boundary but does not decide whether repetition exists.

C Conditional OPD gradient at teacher--student agreement

Conditioned on a fixed sampled prefix, the generalized-JSD OPD gradient vanishes at teacher–student agreement and decays linearly nearby. Top-100-plus-tail aggregation and clipped importance weighting preserve this stationary point, with binned agreement sufficient for zero signal.

C Conditional OPD gradient at teacher--student agreement: The derivation conditions on one sampled prefix held fixed during backpropagation, matching the implemented on-policy update, and uses generalized JSD rather than forward KL.The two KL terms’ non-logarithmic contributions cancel during differentiation.
C Conditional OPD gradient at teacher--student agreement: At exact teacher–student agreement, conditional logit and parameter gradients are zero, vanishing linearly nearby under retained-support positivity and a bounded local Jacobian.At α=0.5, the local slope is one quarter of the forward-KL gradient.
C Conditional OPD gradient at teacher--student agreement: Top-100 teacher bins plus an aggregated tail extend the analysis to binned distributions, so coarse-grained agreement already removes the update signal.Differences within the tail bin are invisible to the loss, making this condition weaker than full-vocabulary agreement.
C Conditional OPD gradient at teacher--student agreement: Clipped importance weighting preserves the stationary point because its bounded multiplier scales a divergence and gradient that are both zero at exact agreement.This is the relevant condition for the implementation’s loop states.

D Evaluation details

The evaluation uses task-specific decoding, response caps, and answer normalization across math, code, open-ended, and multiple-choice benchmarks. Main and matched comparisons are rerun under unified no-thinking protocols, with remaining matched Math+Code results reported in Table 11.

Main evaluation protocol: GSM8K uses all 1,319 test examples with exact/normalized matching, while MATH-500 uses boxed-answer extraction and normalization under a sampled 16,384-token protocol.The 4B Instruct main results use no-thinking decoding with a 512-token GSM8K cap; MATH-500 uses temperature 0.7, top-p 0.8, top-k 20, minimum-p 0, and repetition stopping only truncates detected loops.
Comparison protocols: The width/hybrid and matched Math+Code comparisons use unified no-thinking reruns, with greedy decoding for GSM8K and seven non-MATH tasks and sampled decoding for MATH-500.Response caps are 512 tokens for GSM8K and 768 for code and open-ended generation; width/hybrid multiple-choice evaluation disables thinking, uses zero-shot MMLU, a 4096-token maximum length, and mean log-probability.
Matched Math+Code results: Table 11 reports remaining full-evaluation results for the matched Math+Code comparison, including judge scores on a 1–10 scale and accuracy percentages for multiple-choice columns and MC Avg.These results use the matched comparison protocol described above.

E Repetition versus pruning depth: the full sweep

Across the full 36-layer pruning sweep, repetition is the dominant degeneration mode only while the model remains coherent (k≤12); deeper pruning instead produces incoherent text. The transition is also visible qualitatively, from recognizable prefixes with repetitive suffixes at 25% compression to mixed-token collapse under deeper compression.

E Repetition versus pruning depth: the full sweep: While the model remains coherent, it retains enough structure to sustain n-gram loops, making looping its dominant degeneration mode.This explains why repetition tracks increasing removed depth only within the coherent regime.
E Repetition versus pruning depth: the full sweep: Repetition rises and distinct-2 diversity falls monotonically through k≤12, whereas pruning beyond k≈13 collapses outputs into incoherent text where loops no longer form.The full sweep covers all 36 Qwen3-4B-Instruct layers on a 192-prompt probe.
E Repetition versus pruning depth: the full sweep: At 25% compression, generations retain recognizable prefixes but develop repetitive suffixes; deeper compression instead enters an incoherent mixed-token regime.Table 12 excerpts illustrate the transition from normal generation to repetition and then incoherent collapse.

F Multiple-choice recognition results · F.1 Width and hybrid pruning

Multiple-choice recognition is a complementary diagnostic to the paper’s primary free-form generation evaluation, covering held-out ARC-Challenge, HellaSwag, MMLU, and WinoGrande. Width and hybrid pruning preserve recognition better than generation, but ShortOPD’s smaller recognition gains show that fixed-option accuracy is not a sufficient proxy.

F Multiple-choice recognition results: Held-out multiple-choice recovery on ARC-Challenge, HellaSwag, MMLU, and WinoGrande serves as a complementary diagnostic rather than the paper’s primary generation target.These benchmarks are not part of the recovery corpus, and the appendix collects the multiple-choice tables separately.
F.1 Width and hybrid pruning: Evaluation disables thinking and scores candidate letters after a fixed assistant ANSWER: prefix using mean token log-probability, with zero-shot MMLU capped at 4096 tokens.This is the evaluation protocol for the recognition results.
F.1 Width and hybrid pruning: 51.14 and 69.41 MC Avg remain for pruned 4B-width and 8B-hybrid models, despite generation Avg scores of only 17.23 and 16.17.Recognition degrades less than generation under width and hybrid pruning.
F.1 Width and hybrid pruning: ShortOPD exceeds KD by 0.87–2.08 MC points, versus 6.44–9.01 generation gains, showing fixed-option recognition is not a sufficient proxy for generation.The recognition comparison is reported for the pruned width and hybrid models.

F.2 Depth-pruning sanity check

The depth-pruning sanity check shows that ShortOPD substantially recovers multiple-choice recognition, especially when initialized from SFT and trained for three epochs. Combining SFT and ShortOPD reaches the teacher’s recognition score.

F.2 Depth-pruning sanity check: 76.8 MC Avg from SFT-init ShortOPD matches the Teacher at 76.7 on Qwen3-4B-Instruct.SFT-init ShortOPD starts from the one-epoch SFT w/o KD checkpoint rather than the pruned student.
F.2 Depth-pruning sanity check: Three-epoch one-epoch ShortOPD rises to 71.6 MC Avg, from 67.6 after one epoch and 22 points above Pruned.Table 14 uses candidate log-likelihood scoring on a held-out multiple-choice recognition task.
F.2 Depth-pruning sanity check: The recognition task favors SFT and KD at 74.9/74.4 MC Avg, despite their weaker generation performance in Table 3.Scores are accuracies under candidate log-likelihood scoring with thinking disabled.