ShortOPD: Recovering Pruned LLMs with Short-to-Long On-Policy Distillation
AI 深度解读
背景
结构化剪枝是一种对硬件友好的大语言模型（LLM）压缩方法，但其有效性目前主要在多选识别任务上得到验证。当同样经过压缩的检查点被用于实际部署所需的自由文本生成任务时，模型性能往往会急剧恶化。作者通过两项观察揭示了这一差距的来源：第一，压缩后模型的贪婪解码 pass@1 几乎消失，但 pass@k 在多次采样下显著恢复——说明有用的生成内容并未被完全抹去，而是被降级；第二，可恢复区间内的失败主要由后缀重复导致。因此，恢复训练应利用压缩模型自身的在策略状态（on-policy states），配合密集的 token 级监督。这正是 On-Policy Distillation (OPD) 所提供的——通过将压缩前的模型作为冻结教师（teacher）来复用。然而，长序列在策略 rollout 会将早期恢复预算浪费在信息量低的重复后缀上，从而延迟 loss 下降。为缓解这一浪费，作者提出了 ShortOPD。
核心内容
ShortOPD（Short-to-Long On-Policy Distillation）是一种分阶段的在策略蒸馏调度方案。其核心流程如下：
- 检测教师模型确认的重复后缀：在每次 rollout 中，利用冻结的预压缩教师模型判断当前压缩模型生成序列中哪些部分是重复后缀（即低信息量的无意义重复）。
- 保留有效前缀：将教师模型确认的重复后缀之前的剩余部分视为该次 rollout 的“有效长度”（effective length）。
- 动态分配预算：后续 rollout 的预算仅分配给当前策略（压缩模型）所能有效利用的有效长度，从而避免在重复后缀上浪费计算资源。
通过这种短到长的渐进式调度，ShortOPD 在数学、代码和开放式生成任务上均取得了显著效果：
- 压缩模型的得分提升至其未恢复值的约 9 倍。
- 相较于标准恢复方法（SFT w/o KD、KD 和 SeqKD），得分提升 1.6 至 4.4 倍。
- 在训练时间上，ShortOPD 仅用 8.5 小时即可达到固定 8192 token rollout 长度下 35.9 小时的效果（节省约 76% 的时间），同时减少了 71% 的 rollout token 数量，最终性能差距在 2 个点以内。
关键要点
- 问题核心：结构化剪枝后的 LLM 在自由生成任务中性能崩溃，主要原因是生成序列中出现重复后缀，而非有用内容完全丢失。
- 现有方案局限：传统 On-Policy Distillation (OPD) 虽然能提供 token 级监督，但长 rollout 在早期阶段大部分预算被重复后缀消耗，导致损失下降缓慢。
- ShortOPD 创新：
- 利用教师模型在线检测重复后缀，动态截断 rollout 的有效长度。
- 采用短到长的渐进式训练策略，让压缩模型先学会短序列的稳定生成，再逐步增加长度。
- 实验效果：
- 在数学（如 GSM8K）、代码（如 HumanEval）和开放式生成（如 AlpacaEval）上均大幅超越现有恢复配方。
- 训练效率极高：仅需约一半的训练时间（8.5 vs 35.9 小时）即可达到接近固定长 rollout 的效果，token 消耗减少 71%。
- 适用范围：适用于任何采用结构化剪枝并希望将压缩模型部署到实际生成场景的 LLM。
意义与影响
ShortOPD 的意义在于将结构化剪枝从“仅提升困惑度和多选基准的边际收益”推进到“真正面向部署的生成质量”。它提供了一种高效、可落地的恢复方案，使得压缩后的 LLM 不仅能在标准评测中表现良好，更能在自由文本生成任务中稳定输出。该工作的核心洞察——重复后缀是压缩模型失败的主因，且可通过教师引导的短到长训练来克服——为后续剪枝后恢复研究提供了新的方向。同时，ShortOPD 在训练时间和资源消耗上的巨大优势，使其更容易被工业界采纳，有望推动硬件友好的 LLM 压缩技术在实际产品中广泛落地。