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

Root Cause: Structural Pruning Fails in Free-Form Generation

Structured pruning is hardware-friendly for LLM compression, yet its validation has largely been confined to multiple-choice recognition tasks (e.g., MMLU, ARC). In contrast, the same pruned checkpoints often collapse under real-world free-form generation—indicating not total failure, but severe degradation in output quality.

Key Observations and Diagnostics

• Observation 1: Greedy \textsc{pass}@1 nearly vanishes post-pruning, yet \textsc{pass}@k recovers substantially under repeated sampling—demonstrating that useful generations are demoted, not erased;
• Observation 2: The recoverable regime fails predominantly via suffix repetition: models loop in late-generation stages rather than err early.

Method Design: Three Innovations of ShortOPD

• Core Mechanism: Builds on On-Policy Distillation (OPD), reusing the pre-pruning model as a frozen teacher to supervise the pruned student at token level on its own on-policy states;
• Key Optimization: Standard long on-policy rollouts waste early training budget on low-information repetitive suffixes. ShortOPD dynamically detects teacher-confirmed repetitive suffixes—those flagged as redundant by the teacher—and treats the surviving non-repetitive prefix as each rollout’s effective length;
• Scheduling Strategy: Implements a short-to-long rollout schedule: starting with short effective lengths and progressively extending them, enabling earlier gradient updates on high-information-density early-generation tokens.

Empirical Validation

Evaluated on Llama-2-7b and Qwen2-7b, ShortOPD outperforms standard OPD under identical training budgets: \textsc{pass}@1 improves by 12.3–18.7 points on free-form generation benchmarks (GSM8K, HumanEval), while repetition rate drops by 34%.