ReadPaper Blog
ShortOPD: Recovering Pruned LLMs with Short-to-Long On-Policy Distillation
ShortOPD addresses a practical failure mode in structured pruning: compressed large language models can retain multiple-choice recognition performance while collapsing on free-form generation needed in deployment. The paper proposes short-to-long on-policy distillation, using the pre-compression model as a frozen teacher and adapting rollout length to avoid wasting training budget on repetitive suffixes. Experiments on a 25%-pruned Qwen3-4B-Instruct report substantially better generation recovery across math, code, and open-ended tasks with less rollout compute than fixed-horizon OPD and standard distillation recipes.
Source: ShortOPD: Recovering Pruned LLMs with Short-to-Long On-Policy Distillation
Page 1 — The Problem on the Rooftops
The paper starts from the deployment appeal of structured pruning, especially depth pruning, because removing whole transformer blocks preserves dense transformer execution and avoids specialized sparse or low-bit kernels. Its central problem is that this hardware-friendly compression is often judged on recognition benchmarks such as MMLU and HellaSwag, while real deployments require free-form generation across reasoning, coding, and instruction-following tasks. Using Block-Influence depth pruning on Qwen3-4B-Instruct, the authors show that removing 4 of 36 layers can severely damage greedy generation even when the model’s architecture remains easy to serve. This recognition-generation gap motivates the paper’s recovery question: after useful capacity is damaged by pruning, what training signal can restore generation quality rather than merely preserve perplexity or multiple-choice accuracy? The implication is that post-pruning evaluation must include generated trajectories, because local next-token errors can compound into unusable outputs despite apparently acceptable benchmark retention.
Page 2 — The Clue Hidden in Sampling
A key empirical clue in the paper is that pruning demotes useful generations rather than erasing them from the model’s distribution. The authors compare greedy pass@1 with best-of-k sampling on math, code, and open-ended benchmarks, and they find that performance rises steadily as more samples are drawn from the pruned Qwen3-4B-Instruct. On GSM8K, the pruned model with 4 removed layers reaches 91.2% at k=64, exceeding the unpruned teacher’s greedy score of 88.1%, while other tasks such as MATH500, HumanEval, MBPP, Alpaca, QA, summarization, and MT-Bench also improve with repeated attempts. This evidence supports the paper’s claim that correct trajectories remain reachable under the compressed model’s own sampling distribution, but are ranked too low for greedy decoding. That finding makes on-policy recovery attractive, because training should correct the states the pruned student actually visits rather than relying only on fixed teacher outputs or static corpus labels.
Page 3 — Why Ordinary Recovery Fails
The paper argues that ordinary recovery recipes miss the mechanism of generation collapse because the error is token-level and path-dependent. ShortGPT-gen’s observation that routing only generated tokens through the full model can restore quality suggests that failures arise from local next-token mis-rankings along the student’s decoding path, not simply from missing global knowledge. On-Policy Distillation (OPD) directly targets this failure mode by letting the compressed student sample its own rollouts while the frozen pre-compression model supplies dense next-token distribution targets at each response position. However, the authors identify a second obstacle: early on-policy rollouts after pruning are dominated by suffix repetition, with the 25%-parameter testbed showing high rates of repeated tails on a fixed 192-prompt probe. These repetitive suffix tokens carry little marginal distillation signal, with much lower teacher–student generalized Jensen–Shannon divergence and teacher negative log-likelihood than ordinary tokens, so long fixed rollouts spend substantial teacher and token budget on low-information regions.
Page 4 — ShortOPD Steps In
ShortOPD is the paper’s method for preserving OPD’s dense on-policy supervision while avoiding wasted long rollouts during early recovery. The algorithm detects teacher-confirmed repetitive suffixes, treats the remaining prefix as the rollout’s effective length, and uses this measurement to adapt the future response-token budget. Its controller uses repetition and truncation feedback: high repetition opens a shrink gate whose target follows an exponential moving average of effective length, while low repetition combined with high truncation opens a growth gate so the horizon expands as the student becomes capable of longer useful generations. A second exponential moving average moves the actual budget smoothly toward the chosen target, preventing abrupt horizon changes from noisy batch estimates. This short-to-long schedule is designed to match rollout length to the current recoverable regime of the pruned policy, keeping dense teacher supervision on informative on-policy states without adding extra forward cost.
Page 5 — The Victory Lap
The experimental results reported in the paper position ShortOPD as both more effective and more compute-efficient than standard post-pruning recovery. On a 25%-pruned Qwen3-4B-Instruct trained with a 45,447-prompt math, code, and instruction corpus, ShortOPD restores about two-thirds of the unpruned teacher’s average generation score across eight task families. The paper reports that ShortOPD raises the compressed model’s score to about 9× its unrecovered value and 1.6–4.4× the scores achieved by SFT without KD, conventional KD, and SeqKD under matched budgets. It also matches a fixed 8192-token rollout horizon within two points while using a quarter of the training time, 8.5 versus 35.9 hours, and 71% fewer rollout tokens. The authors further emphasize that corpus coverage matters, because leave-one-domain-out ablations show that removing math or code prompts sharply weakens recovery on the corresponding capabilities, making post-compression data design a first-class requirement for deployment-ready pruning.
