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

The Problem: Pruned LLMs Break on Real Generation Tasks

Structured pruning compresses large language models into hardware-friendly versions, but it works deceptively well on multiple-choice benchmarks while catastrophically failing at free-form text generation—the task that actually matters in deployment. The authors discover why: after compression, models can generate correct text when sampled repeatedly (pass@k recovers substantially), but the greedy first output (pass@1) nearly vanishes. This isn't because knowledge is erased—it's demoted. The real failure mode is suffix repetition: the model gets stuck repeating the same tokens endlessly rather than producing varied, coherent continuations.

How ShortOPD Works

The authors propose On-Policy Distillation (OPD), which trains the compressed model using its own generations as training data, with dense token-level supervision from the original uncompressed model acting as a frozen teacher. This forces the compressed model to learn from its actual failure modes rather than just mimicking the original on static data.

The key insight: long rollouts waste training budget. Early in recovery, the model generates low-information, repetitive suffixes—predictable garbage the teacher easily confirms. These repetitive tails dominate the training signal and delay meaningful learning. ShortOPD detects when the compressed model starts repeating and truncates rollouts at that point, treating only the non-repetitive prefix as the "effective length" worth training on.

The algorithm then allocates future training budgets (tokens and rollouts) proportionally to the effective lengths the policy can currently sustain. As the model improves and generates longer non-repetitive sequences, the rollout length adaptively increases. This focuses computation on high-signal prefixes and avoids wasting tokens on predictable repetitions.

Results and Impact

Across three diverse domains—math problems, code generation, and open-ended text—ShortOPD recovers compressed models to roughly 9× their unrecovered baseline score. It outperforms standard recovery recipes (fine-tuning without knowledge distillation, standard knowledge distillation, and sequence-level distillation) by 1.6–4.4×.

Crucially, ShortOPD matches the quality of fixed 8192-token rollouts using only a quarter of the training time (8.5 vs. 35.9 hours) and 71% fewer rollout tokens. This efficiency matters: it makes recovery computationally practical rather than a luxury.

The work's contribution is methodological rather than architectural. It identifies a concrete bottleneck in distillation recovery (suffix repetition waste) and designs a targeted, adaptive schedule to fix it. By moving structured pruning from a marginal perplexity-improvement trick into a technique that preserves deployment-ready generation quality, the authors make pruned models genuinely usable. The result suggests that careful training-time optimization—detecting what the model is actually struggling with—can unlock orders of magnitude more value from compressed checkpoints than naive approaches.