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

2026-07-14

AI Summary

This work addresses the significant performance degradation of large language models after structured pruning in open-ended generation tasks. To mitigate this issue, the authors propose Short-to-Long Online Policy Distillation (ShortOPD), a novel mechanism that applies dense token-level supervision on the model’s self-generated trajectories. ShortOPD dynamically identifies and truncates repetitive suffixes confirmed by the teacher model, thereby reallocating training resources to informative prefixes. By integrating structured pruning, online policy distillation, and adaptive sequence truncation, the method achieves up to approximately 9× the performance of the pruned model without recovery and outperforms standard recovery approaches by 1.6–4.4×. Remarkably, it matches the effectiveness of fixed-length sequence training while using only one-fourth of the training time and generating 71% fewer tokens.

Abstract

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.

Problem

• structured pruning
• LLM compression
• free-form generation
• on-policy distillation
• model recovery

Innovation

• On-Policy Distillation
• Structured Pruning
• Short-to-Long Training
• Token-Level Supervision
• LLM Compression